Merge branch 'main' into sea-migration

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -96,7 +96,7 @@ def execute_command(
             max_rows: Maximum number of rows to fetch in a single fetch batch
             max_bytes: Maximum number of bytes to fetch in a single fetch batch
             lz4_compression: Whether to use LZ4 compression for result data
-            cursor: The cursor object that will handle the results
+            cursor: The cursor object that will handle the results. The command id is set in this cursor.
             use_cloud_fetch: Whether to use cloud fetch for retrieving large result sets
             parameters: List of parameters to bind to the query
             async_op: Whether to execute the command asynchronously
@@ -282,7 +282,9 @@ def get_tables(
             max_bytes: Maximum number of bytes to fetch in a single batch
             cursor: The cursor object that will handle the results
             catalog_name: Optional catalog name pattern to filter by
+                if catalog_name is None, we fetch across all catalogs
             schema_name: Optional schema name pattern to filter by
+                if schema_name is None, we fetch across all schemas
             table_name: Optional table name pattern to filter by
             table_types: Optional list of table types to filter by (e.g., ['TABLE', 'VIEW'])
 
@@ -321,6 +323,7 @@ def get_columns(
             catalog_name: Optional catalog name pattern to filter by
             schema_name: Optional schema name pattern to filter by
             table_name: Optional table name pattern to filter by
+                if table_name is None, we fetch across all tables
             column_name: Optional column name pattern to filter by
 
         Returns:

src/databricks/sql/backend/sea/backend.py

@@ -5,7 +5,7 @@
 import re
 from typing import Any, Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set
 
-from databricks.sql.backend.sea.models.base import ExternalLink, ResultManifest
+from databricks.sql.backend.sea.models.base import ResultManifest, StatementStatus
 from databricks.sql.backend.sea.utils.constants import (
     ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
     ResultFormat,
@@ -19,7 +19,7 @@
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
 
-from databricks.sql.backend.sea.result_set import SeaResultSet
+from databricks.sql.result_set import SeaResultSet
 
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.backend.types import (
@@ -45,14 +45,34 @@
     GetStatementResponse,
     CreateSessionResponse,
 )
-from databricks.sql.backend.sea.models.responses import GetChunksResponse
 
 logger = logging.getLogger(__name__)
 
 
 def _filter_session_configuration(
     session_configuration: Optional[Dict[str, Any]],
 ) -> Dict[str, str]:
+    """
+    Filter and normalise the provided session configuration parameters.
+
+    The Statement Execution API supports only a subset of SQL session
+    configuration options.  This helper validates the supplied
+    ``session_configuration`` dictionary against the allow-list defined in
+    ``ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP`` and returns a new
+    dictionary that contains **only** the supported parameters.
+
+    Args:
+        session_configuration: Optional mapping of session configuration
+            names to their desired values.  Key comparison is
+            case-insensitive.
+
+    Returns:
+        Dict[str, str]: A dictionary containing only the supported
+        configuration parameters with lower-case keys and string values.  If
+        *session_configuration* is ``None`` or empty, an empty dictionary is
+        returned.
+    """
+
     if not session_configuration:
         return {}
 
@@ -90,7 +110,6 @@ class SeaDatabricksClient(DatabricksClient):
     STATEMENT_PATH = BASE_PATH + "statements"
     STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}"
     CANCEL_STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}/cancel"
-    CHUNK_PATH_WITH_ID_AND_INDEX = STATEMENT_PATH + "/{}/result/chunks/{}"
 
     # SEA constants
     POLL_INTERVAL_SECONDS = 0.2
@@ -143,7 +162,7 @@ def __init__(
             http_path=http_path,
             http_headers=http_headers,
             auth_provider=auth_provider,
-            ssl_options=self._ssl_options,
+            ssl_options=ssl_options,
             **kwargs,
         )
 
@@ -275,32 +294,9 @@ def close_session(self, session_id: SessionId) -> None:
             data=request_data.to_dict(),
         )
 
-    @staticmethod
-    def get_default_session_configuration_value(name: str) -> Optional[str]:
-        """
-        Get the default value for a session configuration parameter.
-
-        Args:
-            name: The name of the session configuration parameter
-
-        Returns:
-            The default value if the parameter is supported, None otherwise
-        """
-        return ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP.get(name.upper())
-
-    @staticmethod
-    def get_allowed_session_configurations() -> List[str]:
-        """
-        Get the list of allowed session configuration parameters.
-
-        Returns:
-            List of allowed session configuration parameter names
-        """
-        return list(ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP.keys())
-
     def _extract_description_from_manifest(
         self, manifest: ResultManifest
-    ) -> List[Tuple]:
+    ) -> Optional[List]:
         """
         Extract column description from a manifest object, in the format defined by
         the spec: https://peps.python.org/pep-0249/#description
@@ -309,28 +305,39 @@ def _extract_description_from_manifest(
             manifest: The ResultManifest object containing schema information
 
         Returns:
-            List[Tuple]: A list of column tuples
+            Optional[List]: A list of column tuples or None if no columns are found
         """
 
         schema_data = manifest.schema
         columns_data = schema_data.get("columns", [])
 
+        if not columns_data:
+            return None
+
         columns = []
         for col_data in columns_data:
             # Format: (name, type_code, display_size, internal_size, precision, scale, null_ok)
+            name = col_data.get("name", "")
+            type_name = col_data.get("type_name", "")
+            type_name = (
+                type_name[:-5] if type_name.endswith("_TYPE") else type_name
+            ).lower()
+            precision = col_data.get("type_precision")
+            scale = col_data.get("type_scale")
+
             columns.append(
                 (
-                    col_data.get("name", ""),  # name
-                    col_data.get("type_name", ""),  # type_code
+                    name,  # name
+                    type_name,  # type_code
                     None,  # display_size (not provided by SEA)
                     None,  # internal_size (not provided by SEA)
-                    col_data.get("precision"),  # precision
-                    col_data.get("scale"),  # scale
-                    col_data.get("nullable", True),  # null_ok
+                    precision,  # precision
+                    scale,  # scale
+                    None,  # null_ok
                 )
             )
 
-        return columns
+        return columns if columns else None
 
     def _results_message_to_execute_response(
         self, response: Union[ExecuteStatementResponse, GetStatementResponse]
@@ -351,7 +358,7 @@ def _results_message_to_execute_response(
 
         # Check for compression
         lz4_compressed = (
-            response.manifest.result_compression == ResultCompression.LZ4_FRAME.value
+            response.manifest.result_compression == ResultCompression.LZ4_FRAME
         )
 
         execute_response = ExecuteResponse(
@@ -389,8 +396,9 @@ def _response_to_result_set(
         )
 
     def _check_command_not_in_failed_or_closed_state(
-        self, state: CommandState, command_id: CommandId
+        self, status: StatementStatus, command_id: CommandId
     ) -> None:
+        state = status.state
         if state == CommandState.CLOSED:
             raise DatabaseError(
                 "Command {} unexpectedly closed server side".format(command_id),
@@ -399,8 +407,11 @@ def _check_command_not_in_failed_or_closed_state(
                 },
             )
         if state == CommandState.FAILED:
+            error = status.error
+            error_code = error.error_code if error else "UNKNOWN_ERROR_CODE"
+            error_message = error.message if error else "UNKNOWN_ERROR_MESSAGE"
             raise ServerOperationError(
-                "Command {} failed".format(command_id),
+                "Command failed: {} - {}".format(error_code, error_message),
                 {
                     "operation-id": command_id,
                 },
@@ -414,16 +425,18 @@ def _wait_until_command_done(
         """
 
         final_response: Union[ExecuteStatementResponse, GetStatementResponse] = response
-
-        state = final_response.status.state
         command_id = CommandId.from_sea_statement_id(final_response.statement_id)
 
-        while state in [CommandState.PENDING, CommandState.RUNNING]:
+        while final_response.status.state in [
+            CommandState.PENDING,
+            CommandState.RUNNING,
+        ]:
             time.sleep(self.POLL_INTERVAL_SECONDS)
             final_response = self._poll_query(command_id)
-            state = final_response.status.state
 
-        self._check_command_not_in_failed_or_closed_state(state, command_id)
+        self._check_command_not_in_failed_or_closed_state(
+            final_response.status, command_id
+        )
 
         return final_response
 
@@ -457,7 +470,7 @@ def execute_command(
             enforce_embedded_schema_correctness: Whether to enforce schema correctness
 
         Returns:
-            SeaResultSet: A SeaResultSet instance for the executed command
+            ResultSet: A SeaResultSet instance for the executed command
         """
 
         if session_id.backend_type != BackendType.SEA:
@@ -513,14 +526,6 @@ def execute_command(
         )
         response = ExecuteStatementResponse.from_dict(response_data)
         statement_id = response.statement_id
-        if not statement_id:
-            raise ServerOperationError(
-                "Failed to execute command: No statement ID returned",
-                {
-                    "operation-id": None,
-                    "diagnostic-info": None,
-                },
-            )
 
         command_id = CommandId.from_sea_statement_id(statement_id)
 
@@ -552,8 +557,6 @@ def cancel_command(self, command_id: CommandId) -> None:
             raise ValueError("Not a valid SEA command ID")
 
         sea_statement_id = command_id.to_sea_statement_id()
-        if sea_statement_id is None:
-            raise ValueError("Not a valid SEA command ID")
 
         request = CancelStatementRequest(statement_id=sea_statement_id)
         self._http_client._make_request(
@@ -577,8 +580,6 @@ def close_command(self, command_id: CommandId) -> None:
             raise ValueError("Not a valid SEA command ID")
 
         sea_statement_id = command_id.to_sea_statement_id()
-        if sea_statement_id is None:
-            raise ValueError("Not a valid SEA command ID")
 
         request = CloseStatementRequest(statement_id=sea_statement_id)
         self._http_client._make_request(
@@ -596,8 +597,6 @@ def _poll_query(self, command_id: CommandId) -> GetStatementResponse:
             raise ValueError("Not a valid SEA command ID")
 
         sea_statement_id = command_id.to_sea_statement_id()
-        if sea_statement_id is None:
-            raise ValueError("Not a valid SEA command ID")
 
         request = GetStatementRequest(statement_id=sea_statement_id)
         response_data = self._http_client._make_request(
@@ -620,7 +619,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
             CommandState: The current state of the command
 
         Raises:
-            ProgrammingError: If the command ID is invalid
+            ValueError: If the command ID is invalid
         """
 
         response = self._poll_query(command_id)
@@ -648,27 +647,6 @@ def get_execution_result(
         response = self._poll_query(command_id)
         return self._response_to_result_set(response, cursor)
 
-    def get_chunk_links(
-        self, statement_id: str, chunk_index: int
-    ) -> List[ExternalLink]:
-        """
-        Get links for chunks starting from the specified index.
-        Args:
-            statement_id: The statement ID
-            chunk_index: The starting chunk index
-        Returns:
-            ExternalLink: External link for the chunk
-        """
-
-        response_data = self._http_client._make_request(
-            method="GET",
-            path=self.CHUNK_PATH_WITH_ID_AND_INDEX.format(statement_id, chunk_index),
-        )
-        response = GetChunksResponse.from_dict(response_data)
-
-        links = response.external_links or []
-        return links
-
     # == Metadata Operations ==
 
     def get_catalogs(

src/databricks/sql/backend/sea/models/__init__.py

@@ -9,7 +9,6 @@
     StatementStatus,
     ExternalLink,
     ResultData,
-    ColumnInfo,
     ResultManifest,
 )
 
@@ -27,7 +26,6 @@
     ExecuteStatementResponse,
     GetStatementResponse,
     CreateSessionResponse,
-    GetChunksResponse,
 )
 
 __all__ = [
@@ -36,7 +34,6 @@
     "StatementStatus",
     "ExternalLink",
     "ResultData",
-    "ColumnInfo",
     "ResultManifest",
     # Request models
     "StatementParameter",
@@ -50,5 +47,4 @@
     "ExecuteStatementResponse",
     "GetStatementResponse",
     "CreateSessionResponse",
-    "GetChunksResponse",
 ]

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

@@ -67,25 +67,12 @@ class ResultData:
     attachment: Optional[bytes] = None
 
 
-@dataclass
-class ColumnInfo:
-    """Information about a column in the result set."""
-
-    name: str
-    type_name: str
-    type_text: str
-    nullable: bool = True
-    precision: Optional[int] = None
-    scale: Optional[int] = None
-    ordinal_position: Optional[int] = None
-
-
 @dataclass
 class ResultManifest:
     """Manifest information for a result set."""
 
     format: str
-    schema: Dict[str, Any]  # Will contain column information
+    schema: Dict[str, Any]
     total_row_count: int
     total_byte_count: int
     total_chunk_count: int

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

@@ -54,8 +54,8 @@ def to_dict(self) -> Dict[str, Any]:
             result["parameters"] = [
                 {
                     "name": param.name,
-                    **({"value": param.value} if param.value is not None else {}),
-                    **({"type": param.type} if param.type is not None else {}),
+                    "value": param.value,
+                    "type": param.type,
                 }
                 for param in self.parameters
             ]