add GetChunksResponse

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

Author: varun-edachali-dbx

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

@@ -4,7 +4,7 @@
 These models define the structures used in SEA API responses.
 """
 
-from typing import Dict, Any
+from typing import Dict, Any, List
 from dataclasses import dataclass
 
 from databricks.sql.backend.types import CommandState
@@ -154,3 +154,38 @@ class CreateSessionResponse:
     def from_dict(cls, data: Dict[str, Any]) -> "CreateSessionResponse":
         """Create a CreateSessionResponse from a dictionary."""
         return cls(session_id=data.get("session_id", ""))
+
+
+@dataclass
+class GetChunksResponse:
+    """Response from getting chunks for a statement."""
+
+    statement_id: str
+    external_links: List[ExternalLink]
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "GetChunksResponse":
+        """Create a GetChunksResponse from a dictionary."""
+        external_links = []
+        if "external_links" in data:
+            for link_data in 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"),
+                    )
+                )
+
+        return cls(
+            statement_id=data.get("statement_id", ""),
+            external_links=external_links,
+        )

tests/unit/test_sea_backend.py

@@ -6,7 +6,7 @@
 """
 
 import pytest
-from unittest.mock import patch, MagicMock
+from unittest.mock import Mock, patch, MagicMock
 
 from databricks.sql.backend.sea.backend import (
     SeaDatabricksClient,

tests/unit/test_sea_result_set.py

@@ -3,6 +3,7 @@
 """
 
 import pytest
+import unittest 
 from unittest.mock import patch, MagicMock, Mock
 
 from databricks.sql.result_set import SeaResultSet
@@ -39,30 +40,6 @@ def execute_response(self):
         mock_response.is_staging_operation = False
         return mock_response
 
-        # Create a mock CommandId
-        self.mock_command_id = MagicMock()
-        self.mock_command_id.to_sea_statement_id.return_value = "test-statement-id"
-
-        # Create a mock ExecuteResponse for inline data
-        self.mock_execute_response_inline = ExecuteResponse(
-            command_id=self.mock_command_id,
-            status=CommandState.SUCCEEDED,
-            description=self.sample_description,
-            has_been_closed_server_side=False,
-            lz4_compressed=False,
-            is_staging_operation=False,
-        )
-
-        # Create a mock ExecuteResponse for error
-        self.mock_execute_response_error = ExecuteResponse(
-            command_id=self.mock_command_id,
-            status=CommandState.FAILED,
-            description=None,
-            has_been_closed_server_side=False,
-            lz4_compressed=False,
-            is_staging_operation=False,
-        )
-
     def test_init_with_inline_data(self):
         """Test initialization with inline data."""
         # Create mock result data and manifest