Coverage for src/commands/list_tables.py: 91%

44 statements  

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

1""" 

2Command for listing tables in a Unity Catalog schema. 

3""" 

4 

5from typing import Optional, Any 

6from src.clients.databricks import DatabricksAPIClient 

7from src.command_registry import CommandDefinition 

8from src.commands.base import CommandResult 

9from src.config import get_active_catalog, get_active_schema 

10import logging 

11 

12 

13def handle_command( 

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

15) -> CommandResult: 

16 """ 

17 List tables 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 tables from 

24 - include_delta_metadata: Whether delta metadata should be included (optional) 

25 - omit_columns: Whether to omit columns from the response (optional) 

26 - include_browse: Whether to include tables with selective metadata access (optional) 

27 

28 Returns: 

29 CommandResult with list of tables if successful 

30 """ 

31 if not client: 

32 return CommandResult( 

33 False, 

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

35 ) 

36 

37 # Extract parameters 

38 catalog_name = kwargs.get("catalog_name") 

39 schema_name = kwargs.get("schema_name") 

40 include_delta_metadata = kwargs.get("include_delta_metadata", False) 

41 omit_columns = kwargs.get("omit_columns", False) 

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

43 

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

45 if not catalog_name: 

46 catalog_name = get_active_catalog() 

47 if not catalog_name: 

48 return CommandResult( 

49 False, 

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

51 ) 

52 

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

54 if not schema_name: 

55 schema_name = get_active_schema() 

56 if not schema_name: 

57 return CommandResult( 

58 False, 

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

60 ) 

61 

62 try: 

63 # List tables in the schema 

64 result = client.list_tables( 

65 catalog_name=catalog_name, 

66 schema_name=schema_name, 

67 include_delta_metadata=include_delta_metadata, 

68 omit_columns=omit_columns, 

69 include_browse=include_browse, 

70 ) 

71 

72 tables = result.get("tables", []) 

73 

74 if not tables: 

75 return CommandResult( 

76 True, 

77 message=f"No tables found in schema '{catalog_name}.{schema_name}'.", 

78 ) 

79 

80 # Format table information for display 

81 formatted_tables = [] 

82 for table in tables: 

83 # Extract key information 

84 table_info = { 

85 "name": table.get("name"), 

86 "full_name": table.get("full_name"), 

87 "table_type": table.get("table_type", ""), 

88 "data_source_format": table.get("data_source_format", ""), 

89 "comment": table.get("comment", ""), 

90 "created_at": table.get("created_at"), 

91 "updated_at": table.get("updated_at"), 

92 "created_by": table.get("created_by", ""), 

93 "owner": table.get("owner", ""), 

94 "row_count": table.get("properties", {}).get( 

95 "spark.sql.statistics.numRows", "-" 

96 ), 

97 "size_bytes": table.get("properties", {}).get("size_bytes", "Unknown"), 

98 } 

99 

100 # Include columns if available and not omitted 

101 if not omit_columns: 

102 columns = table.get("columns", []) 

103 table_info["column_count"] = len(columns) 

104 if columns: 

105 column_list = [] 

106 for col in columns: 

107 column_list.append( 

108 { 

109 "name": col.get("name"), 

110 "type": col.get( 

111 "type_text", col.get("type", {}).get("name", "") 

112 ), 

113 "nullable": col.get("nullable", True), 

114 } 

115 ) 

116 table_info["columns"] = column_list 

117 

118 formatted_tables.append(table_info) 

119 

120 return CommandResult( 

121 True, 

122 data={ 

123 "tables": formatted_tables, 

124 "total_count": len(formatted_tables), 

125 "catalog_name": catalog_name, 

126 "schema_name": schema_name, 

127 }, 

128 message=f"Found {len(formatted_tables)} table(s) in '{catalog_name}.{schema_name}'.", 

129 ) 

130 except Exception as e: 

131 logging.error(f"Error listing tables: {str(e)}") 

132 return CommandResult(False, message=f"Failed to list tables: {str(e)}", error=e) 

133 

134 

135DEFINITION = CommandDefinition( 

136 name="list-tables", 

137 description="List tables in a Unity Catalog schema.", 

138 handler=handle_command, 

139 parameters={ 

140 "catalog_name": { 

141 "type": "string", 

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

143 }, 

144 "schema_name": { 

145 "type": "string", 

146 "description": "Name of the schema to list tables from.", 

147 }, 

148 "include_delta_metadata": { 

149 "type": "boolean", 

150 "description": "Whether delta metadata should be included.", 

151 "default": False, 

152 }, 

153 "omit_columns": { 

154 "type": "boolean", 

155 "description": "Whether to omit columns from the response.", 

156 "default": False, 

157 }, 

158 "include_browse": { 

159 "type": "boolean", 

160 "description": "Whether to include tables with selective metadata access.", 

161 "default": False, 

162 }, 

163 }, 

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

165 tui_aliases=["/tables"], 

166 needs_api_client=True, 

167 visible_to_user=True, 

168 visible_to_agent=True, 

169 agent_display="full", # Show full table list to agents 

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

171)