fix: convert sea statement id to CommandId type

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

Author: varun-edachali-dbx

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

@@ -40,11 +40,6 @@
     GetStatementResponse,
     CreateSessionResponse,
 )
-from databricks.sql.backend.sea.models.responses import (
-    parse_status,
-    parse_manifest,
-    parse_result,
-)
 
 logger = logging.getLogger(__name__)
 
@@ -341,10 +336,12 @@ def _results_message_to_execute_response(
         description = self._extract_description_from_manifest(response.manifest)
 
         # Check for compression
-        lz4_compressed = response.manifest.result_compression == "LZ4_FRAME"
+        lz4_compressed = (
+            response.manifest.result_compression == ResultCompression.LZ4_FRAME
+        )
 
         execute_response = ExecuteResponse(
-            command_id=response.statement_id,
+            command_id=CommandId.from_sea_statement_id(response.statement_id),
             status=response.status.state,
             description=description,
             has_been_closed_server_side=False,

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

@@ -18,7 +18,7 @@
 )
 
 
-def parse_status(data: Dict[str, Any]) -> StatementStatus:
+def _parse_status(data: Dict[str, Any]) -> StatementStatus:
     """Parse status from response data."""
     status_data = data.get("status", {})
     error = None
@@ -40,7 +40,7 @@ def parse_status(data: Dict[str, Any]) -> StatementStatus:
     )
 
 
-def parse_manifest(data: Dict[str, Any]) -> ResultManifest:
+def _parse_manifest(data: Dict[str, Any]) -> ResultManifest:
     """Parse manifest from response data."""
 
     manifest_data = data.get("manifest", {})
@@ -69,7 +69,7 @@ def parse_manifest(data: Dict[str, Any]) -> ResultManifest:
     )
 
 
-def parse_result(data: Dict[str, Any]) -> ResultData:
+def _parse_result(data: Dict[str, Any]) -> ResultData:
     """Parse result data from response data."""
     result_data = data.get("result", {})
     external_links = None
@@ -118,9 +118,9 @@ def from_dict(cls, data: Dict[str, Any]) -> "ExecuteStatementResponse":
         """Create an ExecuteStatementResponse from a dictionary."""
         return cls(
             statement_id=data.get("statement_id", ""),
-            status=parse_status(data),
-            manifest=parse_manifest(data),
-            result=parse_result(data),
+            status=_parse_status(data),
+            manifest=_parse_manifest(data),
+            result=_parse_result(data),
         )
 
 
@@ -138,9 +138,9 @@ def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse":
         """Create a GetStatementResponse from a dictionary."""
         return cls(
             statement_id=data.get("statement_id", ""),
-            status=parse_status(data),
-            manifest=parse_manifest(data),
-            result=parse_result(data),
+            status=_parse_status(data),
+            manifest=_parse_manifest(data),
+            result=_parse_result(data),
         )