refactor fetch interface

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

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -17,8 +17,7 @@
     from databricks.sql.client import Cursor
 
 from databricks.sql.thrift_api.TCLIService import ttypes
-from databricks.sql.backend.types import SessionId, CommandId, CommandState
-from databricks.sql.utils import ExecuteResponse
+from databricks.sql.backend.types import SessionId, CommandId, CommandState, ExecuteResponse
 from databricks.sql.types import SSLOptions
 
 # Forward reference for type hints

src/databricks/sql/backend/filters.py

@@ -17,11 +17,12 @@
     TYPE_CHECKING,
 )
 
-# Import SeaResultSet for type checking
-from databricks.sql.backend.sea_result_set import SeaResultSet
+from databricks.sql.utils import JsonQueue, SeaResultSetQueueFactory
+from databricks.sql.backend.types import ExecuteResponse, CommandId
+from databricks.sql.backend.models.base import ResultData
 
 if TYPE_CHECKING:
-    from databricks.sql.result_set import ResultSet
+    from databricks.sql.result_set import ResultSet, SeaResultSet
 
 logger = logging.getLogger(__name__)
 
@@ -36,38 +37,47 @@ class ResultSetFilter:
 
     @staticmethod
     def _filter_sea_result_set(
-        result_set: SeaResultSet, filter_func: Callable[[List[Any]], bool]
-    ) -> SeaResultSet:
+        result_set: "SeaResultSet", filter_func: Callable[[List[Any]], bool]
+    ) -> "SeaResultSet":
         """
         Filter a SEA result set using the provided filter function.
-
+        
         Args:
             result_set: The SEA result set to filter
             filter_func: Function that takes a row and returns True if the row should be included
-
+            
         Returns:
             A filtered SEA result set
         """
-        # Create a filtered version of the result set
-        filtered_response = result_set._response.copy()
-
-        # If there's a result with rows, filter them
-        if (
-            "result" in filtered_response
-            and "data_array" in filtered_response["result"]
-        ):
-            rows = filtered_response["result"]["data_array"]
-            filtered_rows = [row for row in rows if filter_func(row)]
-            filtered_response["result"]["data_array"] = filtered_rows
-
-            # Update row count if present
-            if "row_count" in filtered_response["result"]:
-                filtered_response["result"]["row_count"] = len(filtered_rows)
-
-        # Create a new result set with the filtered data
+        # Get all remaining rows
+        original_index = result_set.results.cur_row_index
+        result_set.results.cur_row_index = 0  # Reset to beginning
+        all_rows = result_set.results.remaining_rows()
+        
+        # Filter rows
+        filtered_rows = [row for row in all_rows if filter_func(row)]
+        
+        # Import SeaResultSet here to avoid circular imports
+        from databricks.sql.result_set import SeaResultSet
+        
+        # Reuse the command_id from the original result set
+        command_id = result_set.command_id
+        
+        # Create an ExecuteResponse with the filtered data
+        execute_response = ExecuteResponse(
+            command_id=command_id,
+            status=result_set.status,
+            description=result_set.description,
+            has_more_rows=result_set._has_more_rows,  
+            results_queue=JsonQueue(filtered_rows),
+            has_been_closed_server_side=result_set.has_been_closed_server_side,
+            lz4_compressed=False,
+            is_staging_operation=False,
+        )
+        
         return SeaResultSet(
             connection=result_set.connection,
-            sea_response=filtered_response,
+            execute_response=execute_response,
             sea_client=result_set.backend,
             buffer_size_bytes=result_set.buffer_size_bytes,
             arraysize=result_set.arraysize,
@@ -97,6 +107,7 @@ def filter_by_column_values(
             allowed_values = [v.upper() for v in allowed_values]
 
         # Determine the type of result set and apply appropriate filtering
+        from databricks.sql.result_set import SeaResultSet
         if isinstance(result_set, SeaResultSet):
             return ResultSetFilter._filter_sea_result_set(
                 result_set,
@@ -145,4 +156,4 @@ def filter_tables_by_type(
         # Table type is typically in the 4th column (index 3)
         return ResultSetFilter.filter_by_column_values(
             result_set, 3, valid_types, case_sensitive=False
-        )
+        )

src/databricks/sql/backend/sea_backend.py

@@ -8,11 +8,13 @@
     from databricks.sql.result_set import ResultSet
 
 from databricks.sql.backend.databricks_client import DatabricksClient
-from databricks.sql.backend.types import SessionId, CommandId, CommandState, BackendType
+from databricks.sql.backend.types import SessionId, CommandId, CommandState, BackendType, ExecuteResponse
 from databricks.sql.exc import Error, NotSupportedError, ServerOperationError
 from databricks.sql.backend.utils.http_client import CustomHttpClient
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import SSLOptions
+from databricks.sql.utils import SeaResultSetQueueFactory
+from databricks.sql.backend.models.base import ResultData
 
 from databricks.sql.backend.models import (
     ExecuteStatementRequest,
@@ -227,6 +229,70 @@ def close_session(self, session_id: SessionId) -> None:
             params=request.to_dict(),
         )
 
+    def _results_message_to_execute_response(self, sea_response, command_id):
+        """
+        Convert a SEA response to an ExecuteResponse.
+        
+        Args:
+            sea_response: The response from the SEA API
+            command_id: The command ID
+            
+        Returns:
+            ExecuteResponse: The normalized execute response
+        """
+        # Extract status 
+        status_data = sea_response.get("status", {})
+        state = CommandState.from_sea_state(status_data.get("state", ""))
+        
+        # Extract description from manifest 
+        description = None
+        manifest_data = sea_response.get("manifest", {})
+        schema_data = manifest_data.get("schema", {})
+        columns_data = schema_data.get("columns", [])
+        
+        if columns_data:
+            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(
+                    (
+                        col_data.get("name", ""),  # name
+                        col_data.get("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
+                    )
+                )
+            description = columns if columns else None
+        
+        # Create results queue
+        results_queue = None
+        result_data = sea_response.get("result", {})
+        if result_data:
+            results_queue = SeaResultSetQueueFactory.build_queue(
+                ResultData(
+                    data=result_data.get("data_array", None),
+                    external_links=result_data.get("external_links", None)
+                ),
+                description=description
+            )
+        
+        return ExecuteResponse(
+            command_id=command_id,
+            status=state,
+            description=description,
+            has_more_rows=False,
+            results_queue=results_queue,
+            has_been_closed_server_side=False,
+            lz4_compressed=False, # TODO: extract from response 
+            is_staging_operation=False,
+        )
+
     def execute_command(
         self,
         operation: str,
@@ -444,11 +510,14 @@ def get_execution_result(
         )
 
         # Create and return a SeaResultSet
-        from databricks.sql.backend.sea_result_set import SeaResultSet
-
+        from databricks.sql.result_set import SeaResultSet
+        
+        # Convert the response to an ExecuteResponse
+        execute_response = self._results_message_to_execute_response(response_data, command_id)
+        
         return SeaResultSet(
             connection=cursor.connection,
-            sea_response=response_data,
+            execute_response=execute_response,
             sea_client=self,
             buffer_size_bytes=cursor.buffer_size_bytes,
             arraysize=cursor.arraysize,
@@ -599,4 +668,4 @@ def get_columns(
             enforce_embedded_schema_correctness=False,
         )
         assert result is not None, "execute_command returned None in synchronous mode"
-        return result
+        return result