constrain backend diff

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

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

@@ -1,18 +1,16 @@
 import logging
-import uuid
 import time
 import re
-from typing import Dict, Tuple, List, Optional, Any, Union, TYPE_CHECKING, Set
-
-from databricks.sql.backend.sea.models.base import ExternalLink
+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.utils.constants import (
     ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
-    MetadataCommands,
     ResultFormat,
     ResultDisposition,
     ResultCompression,
     WaitTimeout,
+    MetadataCommands,
 )
 
 if TYPE_CHECKING:
@@ -29,7 +27,6 @@
 )
 from databricks.sql.exc import DatabaseError, ServerOperationError
 from databricks.sql.backend.sea.utils.http_client import SeaHttpClient
-from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import SSLOptions
 
 from databricks.sql.backend.sea.models import (
@@ -43,6 +40,8 @@
     ExecuteStatementResponse,
     GetStatementResponse,
     CreateSessionResponse,
+)
+from databricks.sql.backend.sea.models.responses import (
     GetChunksResponse,
 )
 
@@ -91,6 +90,9 @@ class SeaDatabricksClient(DatabricksClient):
     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
+
     def __init__(
         self,
         server_hostname: str,
@@ -121,7 +123,7 @@ def __init__(
             http_path,
         )
 
-        super().__init__(ssl_options, **kwargs)
+        super().__init__(ssl_options=ssl_options, **kwargs)
 
         # Extract warehouse ID from http_path
         self.warehouse_id = self._extract_warehouse_id(http_path)
@@ -288,28 +290,28 @@ def get_allowed_session_configurations() -> List[str]:
         """
         return list(ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP.keys())
 
-    def _extract_description_from_manifest(self, manifest_obj) -> Optional[List]:
+    def _extract_description_from_manifest(
+        self, manifest: ResultManifest
+    ) -> Optional[List]:
         """
-        Extract column description from a manifest object.
+        Extract column description from a manifest object, in the format defined by
+        the spec: https://peps.python.org/pep-0249/#description
 
         Args:
-            manifest_obj: The ResultManifest object containing schema information
+            manifest: The ResultManifest object containing schema information
 
         Returns:
             Optional[List]: A list of column tuples or None if no columns are found
         """
 
-        schema_data = manifest_obj.schema
+        schema_data = manifest.schema
         columns_data = schema_data.get("columns", [])
 
         if not columns_data:
             return None
 
         columns = []
         for col_data in columns_data:
-            if not isinstance(col_data, dict):
-                continue
-
             # Format: (name, type_code, display_size, internal_size, precision, scale, null_ok)
             columns.append(
                 (
@@ -325,38 +327,9 @@ def _extract_description_from_manifest(self, manifest_obj) -> Optional[List]:
 
         return columns if columns else None
 
-    def get_chunk_link(self, statement_id: str, chunk_index: int) -> 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
-        link = next((l for l in links if l.chunk_index == chunk_index), None)
-        if not link:
-            raise ServerOperationError(
-                f"No link found for chunk index {chunk_index}",
-                {
-                    "operation-id": statement_id,
-                    "diagnostic-info": None,
-                },
-            )
-
-        return link
-
-    def _results_message_to_execute_response(self, response: GetStatementResponse, command_id: CommandId):
+    def _results_message_to_execute_response(
+        self, response: GetStatementResponse
+    ) -> ExecuteResponse:
         """
         Convert a SEA response to an ExecuteResponse and extract result data.
 
@@ -365,18 +338,19 @@ def _results_message_to_execute_response(self, response: GetStatementResponse, c
             command_id: The command ID
 
         Returns:
-            tuple: (ExecuteResponse, ResultData, ResultManifest) - The normalized execute response,
-                  result data object, and manifest object
+            ExecuteResponse: The normalized execute response
         """
 
         # Extract description from manifest schema
         description = self._extract_description_from_manifest(response.manifest)
 
         # Check for compression
-        lz4_compressed = response.manifest.result_compression == ResultCompression.LZ4_FRAME.value
+        lz4_compressed = (
+            response.manifest.result_compression == ResultCompression.LZ4_FRAME.value
+        )
 
         execute_response = ExecuteResponse(
-            command_id=command_id,
+            command_id=CommandId.from_sea_statement_id(response.statement_id),
             status=response.status.state,
             description=description,
             has_been_closed_server_side=False,
@@ -433,7 +407,7 @@ def execute_command(
         lz4_compression: bool,
         cursor: "Cursor",
         use_cloud_fetch: bool,
-        parameters: List,
+        parameters: List[Dict[str, Any]],
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
     ) -> Union["ResultSet", None]:
@@ -467,9 +441,9 @@ def execute_command(
             for param in parameters:
                 sea_parameters.append(
                     StatementParameter(
-                        name=param.name,
-                        value=param.value,
-                        type=param.type if hasattr(param, "type") else None,
+                        name=param["name"],
+                        value=param["value"],
+                        type=param["type"] if "type" in param else None,
                     )
                 )
 
@@ -638,8 +612,7 @@ def get_execution_result(
         # Create and return a SeaResultSet
         from databricks.sql.result_set import SeaResultSet
 
-        # Convert the response to an ExecuteResponse and extract result data
-        execute_response = self._results_message_to_execute_response(response, command_id)
+        execute_response = self._results_message_to_execute_response(response)
 
         return SeaResultSet(
             connection=cursor.connection,
@@ -651,6 +624,35 @@ def get_execution_result(
             manifest=response.manifest,
         )
 
+    def get_chunk_link(self, statement_id: str, chunk_index: int) -> 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
+        link = next((l for l in links if l.chunk_index == chunk_index), None)
+        if not link:
+            raise ServerOperationError(
+                f"No link found for chunk index {chunk_index}",
+                {
+                    "operation-id": statement_id,
+                    "diagnostic-info": None,
+                },
+            )
+
+        return link
+
     # == Metadata Operations ==
 
     def get_catalogs(
@@ -692,7 +694,7 @@ def get_schemas(
         operation = MetadataCommands.SHOW_SCHEMAS.value.format(catalog_name)
 
         if schema_name:
-            operation += f" LIKE '{schema_name}'"
+            operation += MetadataCommands.LIKE_PATTERN.value.format(schema_name)
 
         result = self.execute_command(
             operation=operation,
@@ -724,17 +726,19 @@ def get_tables(
         if not catalog_name:
             raise ValueError("Catalog name is required for get_tables")
 
-        operation = MetadataCommands.SHOW_TABLES.value.format(
+        operation = (
             MetadataCommands.SHOW_TABLES_ALL_CATALOGS.value
             if catalog_name in [None, "*", "%"]
-            else MetadataCommands.CATALOG_SPECIFIC.value.format(catalog_name)
+            else MetadataCommands.SHOW_TABLES.value.format(
+                MetadataCommands.CATALOG_SPECIFIC.value.format(catalog_name)
+            )
         )
 
         if schema_name:
-            operation += f" SCHEMA LIKE '{schema_name}'"
+            operation += MetadataCommands.SCHEMA_LIKE_PATTERN.value.format(schema_name)
 
         if table_name:
-            operation += f" LIKE '{table_name}'"
+            operation += MetadataCommands.LIKE_PATTERN.value.format(table_name)
 
         result = self.execute_command(
             operation=operation,
@@ -750,7 +754,7 @@ def get_tables(
         )
         assert result is not None, "execute_command returned None in synchronous mode"
 
-        # Apply client-side filtering by table_types if specified
+        # Apply client-side filtering by table_types
         from databricks.sql.backend.filters import ResultSetFilter
 
         result = ResultSetFilter.filter_tables_by_type(result, table_types)