mandatory Manifest, clean up statement_id typing

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

Author: varun-edachali-dbx

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

@@ -73,11 +73,15 @@ def _filter_sea_result_set(
         from databricks.sql.result_set import SeaResultSet
 
         # Create a new SeaResultSet with the filtered data
+        manifest = result_set.manifest
+        manifest.total_row_count = len(filtered_rows)
+
         filtered_result_set = SeaResultSet(
             connection=result_set.connection,
             execute_response=execute_response,
             sea_client=cast(SeaDatabricksClient, result_set.backend),
             result_data=result_data,
+            manifest=manifest,
             buffer_size_bytes=result_set.buffer_size_bytes,
             arraysize=result_set.arraysize,
         )

src/databricks/sql/backend/types.py

@@ -369,7 +369,7 @@ def from_sea_statement_id(cls, statement_id: str):
 
         return cls(BackendType.SEA, statement_id)
 
-    def to_thrift_handle(self):
+    def to_thrift_handle(self) -> Optional[ttypes.TOperationHandle]:
         """
         Convert this CommandId to a Thrift TOperationHandle.
 
@@ -390,7 +390,7 @@ def to_thrift_handle(self):
             modifiedRowCount=self.modified_row_count,
         )
 
-    def to_sea_statement_id(self):
+    def to_sea_statement_id(self) -> Optional[str]:
         """
         Get the SEA statement ID string.
 
@@ -401,7 +401,7 @@ def to_sea_statement_id(self):
         if self.backend_type != BackendType.SEA:
             return None
 
-        return self.guid
+        return str(self.guid)
 
     def to_hex_guid(self) -> str:
         """

src/databricks/sql/result_set.py

@@ -473,10 +473,16 @@ def __init__(
             manifest: Manifest from SEA response
         """
 
+        self.manifest = manifest
+
+        statement_id = execute_response.command_id.to_sea_statement_id()
+        if statement_id is None:
+            raise ValueError("Command ID is not a SEA statement ID")
+
         results_queue = SeaResultSetQueueFactory.build_queue(
             result_data,
-            manifest,
-            str(execute_response.command_id.to_sea_statement_id()),
+            self.manifest,
+            statement_id,
             description=execute_response.description,
             max_download_threads=sea_client.max_download_threads,
             sea_client=sea_client,