refactors

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

Author: varun-edachali-dbx

src/databricks/sql/backend/filters.py

@@ -75,13 +75,19 @@ def _filter_sea_result_set(
             is_staging_operation=False,
         )
 
+        # Create a new ResultData object with filtered data
+        from databricks.sql.backend.sea.models.base import ResultData
+        
+        result_data = ResultData(data=filtered_rows, external_links=None)
+
         # Create a new SeaResultSet with the filtered data
         filtered_result_set = SeaResultSet(
             connection=result_set.connection,
             execute_response=execute_response,
             sea_client=cast(SeaDatabricksClient, result_set.backend),
             buffer_size_bytes=result_set.buffer_size_bytes,
             arraysize=result_set.arraysize,
+            result_data=result_data,
         )
 
         return filtered_result_set

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

@@ -449,14 +449,15 @@ def _get_schema_bytes(self, sea_response) -> Optional[bytes]:
 
     def _results_message_to_execute_response(self, sea_response, command_id):
         """
-        Convert a SEA response to an ExecuteResponse.
+        Convert a SEA response to an ExecuteResponse and extract result data.
 
         Args:
             sea_response: The response from the SEA API
             command_id: The command ID
 
         Returns:
-            ExecuteResponse: The normalized execute response
+            tuple: (ExecuteResponse, ResultData, ResultManifest) - The normalized execute response,
+                  result data object, and manifest object
         """
         # Extract status
         status_data = sea_response.get("status", {})
@@ -498,8 +499,10 @@ def _results_message_to_execute_response(self, sea_response, command_id):
         # Check for compression
         lz4_compressed = manifest_data.get("result_compression") == "LZ4_FRAME"
 
-        # Create results queue
-        results_queue = None
+        # Initialize result_data_obj and manifest_obj
+        result_data_obj = None
+        manifest_obj = None
+        
         result_data = sea_response.get("result", {})
         if result_data:
             # Convert external links
@@ -528,31 +531,19 @@ def _results_message_to_execute_response(self, sea_response, command_id):
                 data=result_data.get("data_array"), external_links=external_links
             )
 
-            # Create the manifest object
-            manifest_obj = ResultManifest(
-                format=manifest_data.get("format", ""),
-                schema=manifest_data.get("schema", {}),
-                total_row_count=manifest_data.get("total_row_count", 0),
-                total_byte_count=manifest_data.get("total_byte_count", 0),
-                total_chunk_count=manifest_data.get("total_chunk_count", 0),
-                truncated=manifest_data.get("truncated", False),
-                chunks=manifest_data.get("chunks"),
-                result_compression=manifest_data.get("result_compression"),
-            )
-
-            results_queue = SeaResultSetQueueFactory.build_queue(
-                result_data_obj,
-                manifest_obj,
-                command_id.to_sea_statement_id(),
-                description=description,
-                schema_bytes=schema_bytes,
-                max_download_threads=self.max_download_threads,
-                ssl_options=self.ssl_options,
-                sea_client=self,
-                lz4_compressed=lz4_compressed,
-            )
+        # Create the manifest object
+        manifest_obj = ResultManifest(
+            format=manifest_data.get("format", ""),
+            schema=manifest_data.get("schema", {}),
+            total_row_count=manifest_data.get("total_row_count", 0),
+            total_byte_count=manifest_data.get("total_byte_count", 0),
+            total_chunk_count=manifest_data.get("total_chunk_count", 0),
+            truncated=manifest_data.get("truncated", False),
+            chunks=manifest_data.get("chunks"),
+            result_compression=manifest_data.get("result_compression"),
+        )
 
-        return ExecuteResponse(
+        execute_response = ExecuteResponse(
             command_id=command_id,
             status=state,
             description=description,
@@ -563,6 +554,8 @@ def _results_message_to_execute_response(self, sea_response, command_id):
             arrow_schema_bytes=schema_bytes,
             result_format=manifest_data.get("format"),
         )
+        
+        return execute_response, result_data_obj, manifest_obj
 
     def execute_command(
         self,
@@ -782,8 +775,8 @@ def get_execution_result(
         # Create and return a SeaResultSet
         from databricks.sql.result_set import SeaResultSet
 
-        # Convert the response to an ExecuteResponse
-        execute_response = self._results_message_to_execute_response(
+        # Convert the response to an ExecuteResponse and extract result data
+        execute_response, result_data, manifest = self._results_message_to_execute_response(
             response_data, command_id
         )
 
@@ -793,6 +786,8 @@ def get_execution_result(
             sea_client=self,
             buffer_size_bytes=cursor.buffer_size_bytes,
             arraysize=cursor.arraysize,
+            result_data=result_data,
+            manifest=manifest,
         )
 
     # == Metadata Operations ==

src/databricks/sql/result_set.py

@@ -450,6 +450,8 @@ def __init__(
         sea_client: "SeaDatabricksClient",
         buffer_size_bytes: int = 104857600,
         arraysize: int = 10000,
+        result_data: Optional["ResultData"] = None,
+        manifest: Optional["ResultManifest"] = None,
     ):
         """
         Initialize a SeaResultSet with the response from a SEA query execution.
@@ -460,6 +462,8 @@ def __init__(
             sea_client: The SeaDatabricksClient instance for direct access
             buffer_size_bytes: Buffer size for fetching results
             arraysize: Default number of rows to fetch
+            result_data: Result data from SEA response (optional)
+            manifest: Manifest from SEA response (optional)
         """
         # Extract and store SEA-specific properties
         self.statement_id = (
@@ -468,58 +472,10 @@ def __init__(
             else None
         )
 
-        # Get the response data from the SEA backend
-        response_data = sea_client.http_client._make_request(
-            method="GET",
-            path=sea_client.STATEMENT_PATH_WITH_ID.format(self.statement_id),
-            data={"statement_id": self.statement_id},
-        )
-        
         # Build the results queue
         results_queue = None
 
-        # Extract data from the response
-        result_data = response_data.get("result", {})
-        manifest_data = response_data.get("manifest", {})
-        
         if result_data:
-            # Convert external links
-            external_links = None
-            if "external_links" in result_data:
-                external_links = []
-                for link_data in result_data["external_links"]:
-                    external_links.append(
-                        ExternalLink(
-                            external_link=link_data.get("external_link", ""),
-                            expiration=link_data.get("expiration", ""),
-                            chunk_index=link_data.get("chunk_index", 0),
-                            byte_count=link_data.get("byte_count", 0),
-                            row_count=link_data.get("row_count", 0),
-                            row_offset=link_data.get("row_offset", 0),
-                            next_chunk_index=link_data.get("next_chunk_index"),
-                            next_chunk_internal_link=link_data.get("next_chunk_internal_link"),
-                            http_headers=link_data.get("http_headers", {}),
-                        )
-                    )
-
-            # Create the result data object
-            result_data_obj = ResultData(
-                data=result_data.get("data_array"), external_links=external_links
-            )
-
-            # Create the manifest object
-            manifest_obj = ResultManifest(
-                format=manifest_data.get("format", ""),
-                schema=manifest_data.get("schema", {}),
-                total_row_count=manifest_data.get("total_row_count", 0),
-                total_byte_count=manifest_data.get("total_byte_count", 0),
-                total_chunk_count=manifest_data.get("total_chunk_count", 0),
-                truncated=manifest_data.get("truncated", False),
-                chunks=manifest_data.get("chunks"),
-                result_compression=manifest_data.get("result_compression"),
-            )
-
-            # Build the queue based on the response data
             from typing import cast, List
 
             # Convert description to the expected format
@@ -528,8 +484,8 @@ def __init__(
                 desc = cast(List[Tuple[Any, ...]], execute_response.description)
 
             results_queue = SeaResultSetQueueFactory.build_queue(
-                result_data_obj,
-                manifest_obj,
+                result_data,
+                manifest,
                 str(self.statement_id),
                 description=desc,
                 schema_bytes=execute_response.arrow_schema_bytes if execute_response.arrow_schema_bytes else None,

src/databricks/sql/utils.py

@@ -129,7 +129,7 @@ class SeaResultSetQueueFactory(ABC):
     @staticmethod
     def build_queue(
         sea_result_data: ResultData,
-        manifest: ResultManifest,
+        manifest: Optional[ResultManifest],
         statement_id: str,
         description: Optional[List[Tuple[Any, ...]]] = None,
         schema_bytes: Optional[bytes] = None,
@@ -176,6 +176,10 @@ def build_queue(
                 raise ValueError(
                     "SEA client is required for EXTERNAL_LINKS disposition"
                 )
+            if not manifest:
+                raise ValueError(
+                    "Manifest is required for EXTERNAL_LINKS disposition"
+                )
 
             return SeaCloudFetchQueue(
                 initial_links=sea_result_data.external_links,