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
« 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"""
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
13def handle_command(
14 client: Optional[DatabricksAPIClient], **kwargs: Any
15) -> CommandResult:
16 """
17 List tables in a Unity Catalog schema.
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)
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 )
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)
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 )
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 )
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 )
72 tables = result.get("tables", [])
74 if not tables:
75 return CommandResult(
76 True,
77 message=f"No tables found in schema '{catalog_name}.{schema_name}'.",
78 )
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 }
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
118 formatted_tables.append(table_info)
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)
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)