move get chunk links into backend

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

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -17,7 +17,12 @@
     from databricks.sql.client import Cursor
 
 from databricks.sql.thrift_api.TCLIService import ttypes
-from databricks.sql.backend.types import SessionId, CommandId, CommandState, 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

@@ -41,40 +41,40 @@ def _filter_sea_result_set(
     ) -> "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
         """
         # 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,  
+            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,
             execute_response=execute_response,
@@ -108,6 +108,7 @@ def filter_by_column_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,
@@ -156,4 +157,4 @@ def filter_tables_by_type(
         # Table type is the 6th column (index 5)
         return ResultSetFilter.filter_by_column_values(
             result_set, 5, valid_types, case_sensitive=False
-        )
+        )

src/databricks/sql/backend/models/responses.py

@@ -44,7 +44,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "ExecuteStatementResponse":
             error=error,
             sql_state=status_data.get("sql_state"),
         )
-        
+
         # Parse manifest
         manifest = None
         if "manifest" in data:
@@ -59,13 +59,13 @@ def from_dict(cls, data: Dict[str, Any]) -> "ExecuteStatementResponse":
                 chunks=manifest_data.get("chunks"),
                 result_compression=manifest_data.get("result_compression"),
             )
-        
+
         # Parse result data
         result = None
         if "result" in data:
             result_data = data["result"]
             external_links = None
-            
+
             if "external_links" in result_data:
                 external_links = []
                 for link_data in result_data["external_links"]:
@@ -78,11 +78,13 @@ def from_dict(cls, data: Dict[str, Any]) -> "ExecuteStatementResponse":
                             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"),
+                            next_chunk_internal_link=link_data.get(
+                                "next_chunk_internal_link"
+                            ),
                             http_headers=link_data.get("http_headers"),
                         )
                     )
-            
+
             result = ResultData(
                 data=result_data.get("data_array"),
                 external_links=external_links,
@@ -122,7 +124,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse":
             error=error,
             sql_state=status_data.get("sql_state"),
         )
-        
+
         # Parse manifest
         manifest = None
         if "manifest" in data:
@@ -137,13 +139,13 @@ def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse":
                 chunks=manifest_data.get("chunks"),
                 result_compression=manifest_data.get("result_compression"),
             )
-        
+
         # Parse result data
         result = None
         if "result" in data:
             result_data = data["result"]
             external_links = None
-            
+
             if "external_links" in result_data:
                 external_links = []
                 for link_data in result_data["external_links"]:
@@ -156,11 +158,13 @@ def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse":
                             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"),
+                            next_chunk_internal_link=link_data.get(
+                                "next_chunk_internal_link"
+                            ),
                             http_headers=link_data.get("http_headers"),
                         )
                     )
-            
+
             result = ResultData(
                 data=result_data.get("data_array"),
                 external_links=external_links,
@@ -208,11 +212,13 @@ def from_dict(cls, data: Dict[str, Any]) -> "GetChunksResponse":
                         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"),
+                        next_chunk_internal_link=link_data.get(
+                            "next_chunk_internal_link"
+                        ),
                         http_headers=link_data.get("http_headers"),
                     )
                 )
-                
+
         return cls(
             statement_id=data.get("statement_id", ""),
             external_links=external_links,