introduce strongly typed ChunkInfo

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

Author: varun-edachali-dbx

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

@@ -42,6 +42,16 @@ class ExternalLink:
     http_headers: Optional[Dict[str, str]] = None
 
 
+@dataclass
+class ChunkInfo:
+    """Information about a chunk in the result set."""
+
+    chunk_index: int
+    byte_count: int
+    row_offset: int
+    row_count: int
+
+
 @dataclass
 class ResultData:
     """Result data from a statement execution."""
@@ -73,5 +83,5 @@ class ResultManifest:
     total_byte_count: int
     total_chunk_count: int
     truncated: bool = False
-    chunks: Optional[List[Dict[str, Any]]] = None
+    chunks: Optional[List[ChunkInfo]] = None
     result_compression: Optional[str] = None

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

@@ -14,6 +14,7 @@
     ResultData,
     ServiceError,
     ExternalLink,
+    ChunkInfo,
 )
 
 
@@ -43,15 +44,28 @@ def _parse_manifest(data: Dict[str, Any]) -> ResultManifest:
     """Parse manifest from response data."""
 
     manifest_data = data.get("manifest", {})
+    chunks = None
+    if "chunks" in manifest_data:
+        chunks = [
+            ChunkInfo(
+                chunk_index=chunk.get("chunk_index", 0),
+                byte_count=chunk.get("byte_count", 0),
+                row_offset=chunk.get("row_offset", 0),
+                row_count=chunk.get("row_count", 0),
+            )
+            for chunk in manifest_data.get("chunks", [])
+        ]
+
     return 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"),
+        chunks=chunks,
         result_compression=manifest_data.get("result_compression"),
+        is_volume_operation=manifest_data.get("is_volume_operation"),
     )