Coverage for src/commands/list_volumes.py: 24%

34 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-06-05 22:56 -0700

1""" 

2Command for listing volumes in a Unity Catalog schema. 

3""" 

4 

5from typing import Optional, Any 

6from src.clients.databricks import DatabricksAPIClient 

7from src.commands.base import CommandResult 

8from src.config import get_active_catalog, get_active_schema 

9from src.command_registry import CommandDefinition 

10import logging 

11 

12 

13def handle_command( 

14 client: Optional[DatabricksAPIClient], **kwargs: Any 

15) -> CommandResult: 

16 """ 

17 List volumes in a Unity Catalog schema. 

18 

19 Args: 

20 client: DatabricksAPIClient instance for API calls 

21 **kwargs: Command parameters 

22 - catalog_name: Name of the catalog containing the schema 

23 - schema_name: Name of the schema to list volumes from 

24 - include_browse: Whether to include volumes with selective metadata access (optional) 

25 

26 Returns: 

27 CommandResult with list of volumes if successful 

28 """ 

29 if not client: 

30 return CommandResult( 

31 False, 

32 message="No Databricks client available. Please set up your workspace first.", 

33 ) 

34 

35 # Extract parameters 

36 catalog_name = kwargs.get("catalog_name") 

37 schema_name = kwargs.get("schema_name") 

38 include_browse = kwargs.get("include_browse", False) 

39 

40 # If catalog_name not provided, try to use active catalog 

41 if not catalog_name: 

42 catalog_name = get_active_catalog() 

43 if not catalog_name: 

44 return CommandResult( 

45 False, 

46 message="No catalog specified and no active catalog selected. Please provide a catalog_name or select a catalog first using /select-catalog.", 

47 ) 

48 

49 # If schema_name not provided, try to use active schema 

50 if not schema_name: 

51 schema_name = get_active_schema() 

52 if not schema_name: 

53 return CommandResult( 

54 False, 

55 message="No schema specified and no active schema selected. Please provide a schema_name or select a schema first using /select-schema.", 

56 ) 

57 

58 try: 

59 # List volumes in the schema 

60 result = client.list_volumes( 

61 catalog_name=catalog_name, 

62 schema_name=schema_name, 

63 include_browse=include_browse, 

64 ) 

65 

66 volumes = result.get("volumes", []) 

67 

68 if not volumes: 

69 return CommandResult( 

70 True, 

71 message=f"No volumes found in schema '{catalog_name}.{schema_name}'.", 

72 ) 

73 

74 # Format volume information for display 

75 formatted_volumes = [] 

76 for volume in volumes: 

77 formatted_volume = { 

78 "name": volume.get("name"), 

79 "full_name": volume.get("full_name"), 

80 "volume_type": volume.get("volume_type"), 

81 "comment": volume.get("comment", ""), 

82 "created_at": volume.get("created_at"), 

83 "created_by": volume.get("created_by", ""), 

84 "owner": volume.get("owner", ""), 

85 } 

86 formatted_volumes.append(formatted_volume) 

87 

88 return CommandResult( 

89 True, 

90 data={ 

91 "volumes": formatted_volumes, 

92 "total_count": len(formatted_volumes), 

93 "catalog_name": catalog_name, 

94 "schema_name": schema_name, 

95 }, 

96 message=f"Found {len(formatted_volumes)} volume(s) in '{catalog_name}.{schema_name}'.", 

97 ) 

98 except Exception as e: 

99 logging.error(f"Error listing volumes: {str(e)}") 

100 return CommandResult( 

101 False, message=f"Failed to list volumes: {str(e)}", error=e 

102 ) 

103 

104 

105DEFINITION = CommandDefinition( 

106 name="list-volumes", 

107 description="List volumes in a Unity Catalog schema.", 

108 handler=handle_command, 

109 parameters={ 

110 "catalog_name": { 

111 "type": "string", 

112 "description": "Name of the catalog containing the schema.", 

113 }, 

114 "schema_name": { 

115 "type": "string", 

116 "description": "Name of the schema to list volumes from.", 

117 }, 

118 "include_browse": { 

119 "type": "boolean", 

120 "description": "Whether to include volumes with selective metadata access.", 

121 "default": False, 

122 }, 

123 }, 

124 required_params=[], # Not required anymore as we'll try to get from active config 

125 tui_aliases=["/list-volumes", "/volumes"], 

126 needs_api_client=True, 

127 visible_to_user=True, 

128 visible_to_agent=True, 

129 agent_display="full", # Show full volume list in tables 

130 usage_hint="Usage: /list-volumes [--catalog_name <catalog>] [--schema_name <schema>] [--include_browse true|false]\n(Uses active catalog/schema if not specified)", 

131)