merge

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

Author: varun-edachali-dbx

examples/experimental/sea_connector_test.py

@@ -10,7 +10,8 @@
 import subprocess
 from typing import List, Tuple
 
-logging.basicConfig(level=logging.DEBUG)
+# Configure logging
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 TEST_MODULES = [

src/databricks/sql/backend/filters.py

@@ -17,7 +17,6 @@
     TYPE_CHECKING,
 )
 
-from databricks.sql.utils import JsonQueue, SeaResultSetQueueFactory
 from databricks.sql.backend.types import ExecuteResponse, CommandId
 from databricks.sql.backend.sea.models.base import ResultData
 from databricks.sql.backend.sea.backend import SeaDatabricksClient

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

@@ -1,8 +1,7 @@
 import logging
-import uuid
 import time
 import re
-from typing import Dict, Tuple, List, Optional, Any, Union, TYPE_CHECKING, Set
+from typing import Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set
 
 from databricks.sql.backend.sea.utils.constants import (
     ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
@@ -23,9 +22,7 @@
 )
 from databricks.sql.exc import Error, NotSupportedError, 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.utils import SeaResultSetQueueFactory
 from databricks.sql.backend.sea.models.base import (
     ResultData,
     ExternalLink,

src/databricks/sql/backend/thrift_backend.py

@@ -40,6 +40,7 @@
 )
 
 from databricks.sql.utils import (
+    ThriftResultSetQueueFactory,
     _bound,
     RequestErrorInfo,
     NoRetryReason,

src/databricks/sql/result_set.py

@@ -474,39 +474,19 @@ def __init__(
             result_data: Result data from SEA response (optional)
             manifest: Manifest from SEA response (optional)
         """
-        # Extract and store SEA-specific properties
-        self.statement_id = (
-            execute_response.command_id.to_sea_statement_id()
-            if execute_response.command_id
-            else None
-        )
-
-        # Build the results queue
-        results_queue = None
 
         if result_data:
-            from typing import cast, List
-
-            # Convert description to the expected format
-            desc = None
-            if execute_response.description:
-                desc = cast(List[Tuple[Any, ...]], execute_response.description)
-
-            results_queue = SeaResultSetQueueFactory.build_queue(
-                result_data,
-                manifest,
-                str(self.statement_id),
-                description=desc,
-                schema_bytes=execute_response.arrow_schema_bytes
-                if execute_response.arrow_schema_bytes
-                else None,
-                max_download_threads=sea_client.max_download_threads,
-                ssl_options=sea_client.ssl_options,
-                sea_client=sea_client,
-                lz4_compressed=execute_response.lz4_compressed,
+            queue = SeaResultSetQueueFactory.build_queue(
+                sea_result_data=result_data,
+                manifest=manifest,
+                statement_id=execute_response.command_id.to_sea_statement_id(),
+                description=execute_response.description,
+                schema_bytes=execute_response.arrow_schema_bytes,
             )
+        else:
+            logger.warning("No result data provided for SEA result set")
+            queue = JsonQueue([])
 
-        # Call parent constructor with common attributes
         super().__init__(
             connection=connection,
             backend=sea_client,
@@ -515,15 +495,13 @@ def __init__(
             command_id=execute_response.command_id,
             status=execute_response.status,
             has_been_closed_server_side=execute_response.has_been_closed_server_side,
+            results_queue=queue,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
             lz4_compressed=execute_response.lz4_compressed,
-            arrow_schema_bytes=execute_response.arrow_schema_bytes or b"",
+            arrow_schema_bytes=execute_response.arrow_schema_bytes,
         )
 
-        # Initialize queue for result data if not provided
-        self.results = results_queue or JsonQueue([])
-
     def _convert_to_row_objects(self, rows):
         """
         Convert raw data rows to Row objects with named columns based on description.

src/databricks/sql/utils.py

@@ -16,6 +16,9 @@
 
 import lz4.frame
 
+from databricks.sql.backend.sea.backend import SeaDatabricksClient
+from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
+
 try:
     import pyarrow
 except ImportError: