typing, change DESCRIBE TABLE to SHOW COLUMNS

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

Author: varun-edachali-dbx

src/databricks/sql/backend/models/base.py

@@ -7,6 +7,8 @@
 from typing import Dict, List, Any, Optional, Union
 from dataclasses import dataclass, field
 
+from databricks.sql.backend.types import CommandState
+
 
 @dataclass
 class ServiceError:
@@ -20,7 +22,7 @@ class ServiceError:
 class StatementStatus:
     """Status information for a statement execution."""
 
-    state: str
+    state: CommandState
     error: Optional[ServiceError] = None
     sql_state: Optional[str] = None
 

src/databricks/sql/backend/models/requests.py

@@ -36,7 +36,7 @@ class ExecuteStatementRequest:
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert the request to a dictionary for JSON serialization."""
-        result = {
+        result: Dict[str, Any] = {
             "warehouse_id": self.warehouse_id,
             "session_id": self.session_id,
             "statement": self.statement,
@@ -118,7 +118,7 @@ class CreateSessionRequest:
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert the request to a dictionary for JSON serialization."""
-        result = {"warehouse_id": self.warehouse_id}
+        result: Dict[str, Any] = {"warehouse_id": self.warehouse_id}
 
         if self.session_confs:
             result["session_confs"] = self.session_confs

src/databricks/sql/backend/models/responses.py

@@ -7,10 +7,12 @@
 from typing import Dict, List, Any, Optional, Union
 from dataclasses import dataclass, field
 
+from databricks.sql.backend.types import CommandState
 from databricks.sql.backend.models.base import (
     StatementStatus,
     ResultManifest,
     ResultData,
+    ServiceError,
 )
 
 
@@ -27,14 +29,17 @@ class ExecuteStatementResponse:
     def from_dict(cls, data: Dict[str, Any]) -> "ExecuteStatementResponse":
         """Create an ExecuteStatementResponse from a dictionary."""
         status_data = data.get("status", {})
+        error = None
+        if "error" in status_data:
+            error_data = status_data["error"]
+            error = ServiceError(
+                message=error_data.get("message", ""),
+                error_code=error_data.get("error_code"),
+            )
+
         status = StatementStatus(
-            state=status_data.get("state", ""),
-            error=None
-            if "error" not in status_data
-            else {
-                "message": status_data["error"].get("message", ""),
-                "error_code": status_data["error"].get("error_code"),
-            },
+            state=CommandState.from_sea_state(status_data.get("state", "")),
+            error=error,
             sql_state=status_data.get("sql_state"),
         )
 
@@ -59,14 +64,17 @@ class GetStatementResponse:
     def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse":
         """Create a GetStatementResponse from a dictionary."""
         status_data = data.get("status", {})
+        error = None
+        if "error" in status_data:
+            error_data = status_data["error"]
+            error = ServiceError(
+                message=error_data.get("message", ""),
+                error_code=error_data.get("error_code"),
+            )
+
         status = StatementStatus(
-            state=status_data.get("state", ""),
-            error=None
-            if "error" not in status_data
-            else {
-                "message": status_data["error"].get("message", ""),
-                "error_code": status_data["error"].get("error_code"),
-            },
+            state=CommandState.from_sea_state(status_data.get("state", "")),
+            error=error,
             sql_state=status_data.get("sql_state"),
         )
 

src/databricks/sql/backend/sea_backend.py

@@ -273,13 +273,17 @@ 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,
             statement=operation,
-            disposition="EXTERNAL_LINKS" if use_cloud_fetch else "INLINE",
-            format="ARROW_STREAM" if use_cloud_fetch else "JSON_ARRAY",
+            disposition=disposition,
+            format=format,
             wait_timeout="0s" if async_op else "30s",
             on_wait_timeout="CONTINUE",
             row_limit=max_rows if max_rows > 0 else None,
@@ -315,7 +319,7 @@ def execute_command(
         state = status.state
 
         # Keep polling until we reach a terminal state
-        while state in ["PENDING", "RUNNING"]:
+        while state in [CommandState.PENDING, CommandState.RUNNING]:
             # Add a small delay to avoid excessive API calls
             time.sleep(0.5)
 
@@ -337,12 +341,12 @@ def execute_command(
             state = status.state
 
             # Check for errors
-            if state == "FAILED" and status.error:
-                error_message = status.error["message"]
+            if state == CommandState.FAILED and status.error:
+                error_message = status.error.message
                 raise Error(f"Statement execution failed: {error_message}")
 
             # Check for cancellation
-            if state == "CANCELED":
+            if state == CommandState.CANCELLED:
                 raise Error("Statement execution was canceled")
 
         # Get the final result
@@ -435,11 +439,8 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
         # Parse the response
         response = GetStatementResponse.from_dict(response_data)
 
-        # Extract the status
-        state = response.status.state
-
-        # Map SEA state to CommandState
-        return CommandState.from_sea_state(state)
+        # Return the state directly since it's already a CommandState
+        return response.status.state
 
     def get_execution_result(
         self,
@@ -591,20 +592,20 @@ def get_columns(
         table_name: Optional[str] = None,
         column_name: Optional[str] = None,
     ) -> "ResultSet":
-        """Get columns by executing 'DESCRIBE TABLE [catalog.schema.]table'."""
-        if not table_name:
-            raise ValueError("Table name is required for get_columns")
+        """Get columns by executing 'SHOW COLUMNS IN catalog [SCHEMA LIKE pattern] [TABLE LIKE pattern] [LIKE pattern]'."""
+        if not catalog_name:
+            raise ValueError("Catalog name is required for get_columns")
 
-        operation = "DESCRIBE TABLE "
+        operation = f"SHOW COLUMNS IN `{catalog_name}`"
 
-        if catalog_name and schema_name:
-            operation += f"`{catalog_name}`.`{schema_name}`."
-        elif schema_name:
-            operation += f"`{schema_name}`."
+        if schema_name:
+            operation += f" SCHEMA LIKE '{schema_name}'"
 
-        operation += f"`{table_name}`"
+        if table_name:
+            operation += f" TABLE LIKE '{table_name}'"
 
-        # Column name filtering will be done client-side
+        if column_name:
+            operation += f" LIKE '{column_name}'"
 
         result = self.execute_command(
             operation=operation,

src/databricks/sql/backend/sea_result_set.py

@@ -10,14 +10,15 @@
 
 from databricks.sql.result_set import ResultSet
 from databricks.sql.types import Row
-from databricks.sql.backend.types import CommandId
+from databricks.sql.backend.types import CommandId, CommandState
 from databricks.sql.exc import Error
 
 from databricks.sql.backend.models import (
     StatementStatus,
     ResultManifest,
     ResultData,
     ColumnInfo,
+    ServiceError,
 )
 
 logger = logging.getLogger(__name__)
@@ -42,14 +43,17 @@ def __init__(
 
         # Parse the status
         status_data = sea_response.get("status", {})
+        error = None
+        if "error" in status_data:
+            error_data = status_data["error"]
+            error = ServiceError(
+                message=error_data.get("message", ""),
+                error_code=error_data.get("error_code"),
+            )
+
         self.status = StatementStatus(
-            state=status_data.get("state", ""),
-            error=None
-            if "error" not in status_data
-            else {
-                "message": status_data["error"].get("message", ""),
-                "error_code": status_data["error"].get("error_code"),
-            },
+            state=CommandState.from_sea_state(status_data.get("state", "")),
+            error=error,
             sql_state=status_data.get("sql_state"),
         )
 
@@ -71,7 +75,7 @@ def __init__(
                     )
                 )
 
-            self.manifest = ResultManifest(
+            self.manifest: Optional[ResultManifest] = ResultManifest(
                 schema=columns,
                 total_row_count=manifest_data.get("total_row_count", 0),
                 total_byte_count=manifest_data.get("total_byte_count", 0),
@@ -84,7 +88,7 @@ def __init__(
         # Parse the result data
         result_data = sea_response.get("result")
         if result_data:
-            self.result = ResultData(
+            self.result: Optional[ResultData] = ResultData(
                 data=result_data.get("data"),
                 external_links=result_data.get("external_links"),
             )
@@ -112,7 +116,9 @@ def is_staging_operation(self) -> bool:
 
     def _extract_description_from_manifest(
         self,
-    ) -> List[Tuple[str, str, None, None, None, None, bool]]:
+    ) -> Optional[
+        List[Tuple[str, str, None, None, Optional[int], Optional[int], bool]]
+    ]:
         """
         Extract column descriptions from the SEA manifest.
 
@@ -121,7 +127,7 @@ def _extract_description_from_manifest(
             (name, type_code, display_size, internal_size, precision, scale, null_ok)
         """
         if not self.manifest or not self.manifest.schema:
-            return []
+            return None
 
         description = []
         for col in self.manifest.schema:

src/databricks/sql/result_set.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import List, Optional, Any, Union
+from typing import List, Optional, Any, Union, Tuple
 
 import logging
 import time
@@ -32,7 +32,7 @@ def __init__(self, connection, backend, arraysize: int, buffer_size_bytes: int):
         self.arraysize = arraysize
         self.buffer_size_bytes = buffer_size_bytes
         self._next_row_index = 0
-        self.description = None
+        self.description: Optional[Any] = None
 
     def __iter__(self):
         while True:

tests/unit/test_sea_backend.py

@@ -524,7 +524,7 @@ def test_get_execution_result(
 
         # Verify basic properties of the result
         assert result.statement_id == "test-statement-123"
-        assert result.status.state == "SUCCEEDED"
+        assert result.status.state == CommandState.SUCCEEDED
         assert len(result.description) == 1
         assert result.description[0][0] == "col1"  # column name
 
@@ -657,25 +657,25 @@ def test_get_columns(self, sea_client, mock_cursor, sea_session_id):
             args, kwargs = mock_execute.call_args
             assert (
                 kwargs["operation"]
-                == "DESCRIBE TABLE `test_catalog`.`test_schema`.`test_table`"
+                == "SHOW COLUMNS IN `test_catalog` SCHEMA LIKE 'test_schema' TABLE LIKE 'test_table'"
             )
             assert kwargs["session_id"] == sea_session_id
             assert kwargs["max_rows"] == 100
             assert kwargs["max_bytes"] == 1000
             assert kwargs["cursor"] == mock_cursor
             assert kwargs["async_op"] is False
 
-    def test_get_columns_no_table_name(self, sea_client, mock_cursor, sea_session_id):
-        """Test getting columns without a table name raises an error."""
+    def test_get_columns_no_catalog_name(self, sea_client, mock_cursor, sea_session_id):
+        """Test getting columns without a catalog name raises an error."""
         with pytest.raises(ValueError) as excinfo:
             sea_client.get_columns(
                 session_id=sea_session_id,
                 max_rows=100,
                 max_bytes=1000,
                 cursor=mock_cursor,
-                catalog_name="test_catalog",
+                catalog_name=None,  # No catalog name
                 schema_name="test_schema",
-                table_name=None,  # No table name
+                table_name="test_table",
             )
 
-        assert "Table name is required" in str(excinfo.value)
+        assert "Catalog name is required" in str(excinfo.value)