review metadata ops

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

src/databricks/sql/backend/sea_backend.py

@@ -101,7 +101,6 @@ def _extract_warehouse_id(self, http_path: str) -> str:
 
         The warehouse ID is expected to be the last segment of the path when the
         second-to-last segment is either 'warehouses' or 'endpoints'.
-        This matches the JDBC implementation which supports both formats.
 
         Args:
             http_path: The HTTP path from which to extract the warehouse ID
@@ -182,7 +181,6 @@ def open_session(
             schema=schema,
         )
 
-        # Send the request
         response = self.http_client._make_request(
             method="POST", path=self.SESSION_PATH, data=request.to_dict()
         )
@@ -219,16 +217,14 @@ def close_session(self, session_id: SessionId) -> None:
             raise ValueError("Not a valid SEA session ID")
         sea_session_id = session_id.to_sea_session_id()
 
-        # Create the request model
         request = DeleteSessionRequest(
             warehouse_id=self.warehouse_id, session_id=sea_session_id
         )
 
-        # Send the request
         self.http_client._make_request(
             method="DELETE",
             path=self.SESSION_PATH_WITH_ID.format(sea_session_id),
-            data=request.to_dict(),
+            params=request.to_dict(),
         )
 
     def execute_command(
@@ -279,11 +275,8 @@ def execute_command(
                     )
                 )
 
-        # Determine format and disposition based on use_cloud_fetch
         format = "ARROW_STREAM" if use_cloud_fetch else "JSON_ARRAY"
         disposition = "EXTERNAL_LINKS" if use_cloud_fetch else "INLINE"
-
-        # Create the request model
         request = ExecuteStatementRequest(
             warehouse_id=self.warehouse_id,
             session_id=sea_session_id,
@@ -297,15 +290,10 @@ def execute_command(
             parameters=sea_parameters if sea_parameters else None,
         )
 
-        # Execute the statement
         response_data = self.http_client._make_request(
             method="POST", path=self.STATEMENT_PATH, data=request.to_dict()
         )
-
-        # Parse the response
         response = ExecuteStatementResponse.from_dict(response_data)
-
-        # Create a command ID from the statement ID
         statement_id = response.statement_id
         if not statement_id:
             raise ServerOperationError(
@@ -344,7 +332,6 @@ def execute_command(
                 },
             )
 
-        # Get the final result
         return self.get_execution_result(command_id, cursor)
 
     def cancel_command(self, command_id: CommandId) -> None:
@@ -362,10 +349,7 @@ def cancel_command(self, command_id: CommandId) -> None:
 
         sea_statement_id = command_id.to_sea_statement_id()
 
-        # Create the request model
         request = CancelStatementRequest(statement_id=sea_statement_id)
-
-        # Send the cancel request
         self.http_client._make_request(
             method="POST",
             path=self.CANCEL_STATEMENT_PATH_WITH_ID.format(sea_statement_id),
@@ -387,10 +371,7 @@ def close_command(self, command_id: CommandId) -> None:
 
         sea_statement_id = command_id.to_sea_statement_id()
 
-        # Create the request model
         request = CloseStatementRequest(statement_id=sea_statement_id)
-
-        # Send the close request - SEA uses DELETE for closing statements
         self.http_client._make_request(
             method="DELETE",
             path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
@@ -415,10 +396,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
 
         sea_statement_id = command_id.to_sea_statement_id()
 
-        # Create the request model
         request = GetStatementRequest(statement_id=sea_statement_id)
-
-        # Get the statement status
         response_data = self.http_client._make_request(
             method="GET",
             path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
@@ -427,8 +405,6 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
 
         # Parse the response
         response = GetStatementResponse.from_dict(response_data)
-
-        # Return the state directly since it's already a CommandState
         return response.status.state
 
     def get_execution_result(
@@ -509,10 +485,12 @@ def get_schemas(
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
     ) -> "ResultSet":
-        """Get schemas by executing 'SHOW SCHEMAS [IN catalog]'."""
-        operation = "SHOW SCHEMAS"
-        if catalog_name:
-            operation += f" IN `{catalog_name}`"
+        """Get schemas by executing 'SHOW SCHEMAS IN catalog [LIKE pattern]'."""
+        if not catalog_name:
+            raise ValueError("Catalog name is required for get_schemas")
+
+        operation = f"SHOW SCHEMAS IN `{catalog_name}`"
+
         if schema_name:
             operation += f" LIKE '{schema_name}'"
 
@@ -542,13 +520,18 @@ def get_tables(
         table_name: Optional[str] = None,
         table_types: Optional[List[str]] = None,
     ) -> "ResultSet":
-        """Get tables by executing 'SHOW TABLES [IN catalog.schema]'."""
-        operation = "SHOW TABLES"
+        """Get tables by executing 'SHOW TABLES IN catalog [SCHEMA LIKE pattern] [LIKE pattern]'."""
+        if not catalog_name:
+            raise ValueError("Catalog name is required for get_tables")
+
+        operation = "SHOW TABLES IN " + (
+            "ALL CATALOGS"
+            if catalog_name in [None, "*", "%"]
+            else f"CATALOG `{catalog_name}`"
+        )
 
-        if catalog_name and schema_name:
-            operation += f" IN `{catalog_name}`.`{schema_name}`"
-        elif schema_name:
-            operation += f" IN `{schema_name}`"
+        if schema_name:
+            operation += f" SCHEMA LIKE '{schema_name}'"
 
         if table_name:
             operation += f" LIKE '{table_name}'"
@@ -579,11 +562,11 @@ def get_columns(
         table_name: Optional[str] = None,
         column_name: Optional[str] = None,
     ) -> "ResultSet":
-        """Get columns by executing 'SHOW COLUMNS IN catalog [SCHEMA LIKE pattern] [TABLE LIKE pattern] [LIKE pattern]'."""
+        """Get columns by executing 'SHOW COLUMNS IN CATALOG catalog [SCHEMA LIKE pattern] [TABLE LIKE pattern] [LIKE pattern]'."""
         if not catalog_name:
             raise ValueError("Catalog name is required for get_columns")
 
-        operation = f"SHOW COLUMNS IN `{catalog_name}`"
+        operation = f"SHOW COLUMNS IN CATALOG `{catalog_name}`"
 
         if schema_name:
             operation += f" SCHEMA LIKE '{schema_name}'"

src/databricks/sql/backend/utils/http_client.py

@@ -87,7 +87,11 @@ def _get_auth_headers(self) -> Dict[str, str]:
         return headers
 
     def _make_request(
-        self, method: str, path: str, data: Optional[Dict[str, Any]] = None
+        self,
+        method: str,
+        path: str,
+        data: Optional[Dict[str, Any]] = None,
+        params: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
         """
         Make an HTTP request to the SEA endpoint.
@@ -109,12 +113,18 @@ def _make_request(
         logger.debug(f"making {method} request to {url}")
 
         try:
+            args = {
+                "url": url,
+                "headers": headers,
+                "json": data,
+                "params": params,
+            }
             if method.upper() == "GET":
-                response = self.session.get(url, headers=headers, params=data)
+                response = self.session.get(**args)
             elif method.upper() == "POST":
-                response = self.session.post(url, headers=headers, json=data)
+                response = self.session.post(**args)
             elif method.upper() == "DELETE":
-                response = self.session.delete(url, headers=headers, params=data)
+                response = self.session.delete(**args)
             else:
                 raise ValueError(f"Unsupported HTTP method: {method}")
 
@@ -130,7 +140,11 @@ def _make_request(
 
                 # Log response content (but limit it for large responses)
                 content_str = json.dumps(result, indent=4, sort_keys=True)
-                content_str = content_str[:1000] + "..." if len(content_str) > 1000 else content_str
+                content_str = (
+                    content_str[:1000] + "..."
+                    if len(content_str) > 1000
+                    else content_str
+                )
                 logger.debug(f"Response content: {content_str}")
 
                 return result