Merge branch 'sea-migration' into max-rows-sea

Author: varun-edachali-dbx

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

@@ -353,7 +353,7 @@ def _results_message_to_execute_response(
             description=description,
             has_been_closed_server_side=False,
             lz4_compressed=lz4_compressed,
-            is_staging_operation=False,
+            is_staging_operation=response.manifest.is_volume_operation,
             arrow_schema_bytes=None,
             result_format=response.manifest.format,
         )

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

@@ -92,4 +92,4 @@ class ResultManifest:
     truncated: bool = False
     chunks: Optional[List[ChunkInfo]] = None
     result_compression: Optional[str] = None
-    is_volume_operation: Optional[bool] = None
+    is_volume_operation: bool = False

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

@@ -65,7 +65,7 @@ def _parse_manifest(data: Dict[str, Any]) -> ResultManifest:
         truncated=manifest_data.get("truncated", False),
         chunks=chunks,
         result_compression=manifest_data.get("result_compression"),
-        is_volume_operation=manifest_data.get("is_volume_operation"),
+        is_volume_operation=manifest_data.get("is_volume_operation", False),
     )
 
 

tests/unit/test_sea_backend.py

@@ -633,6 +633,31 @@ def test_utility_methods(self, sea_client):
             sea_client._extract_description_from_manifest(no_columns_manifest) is None
         )
 
+    def test_results_message_to_execute_response_is_staging_operation(self, sea_client):
+        """Test that is_staging_operation is correctly set from manifest.is_volume_operation."""
+        # Test when is_volume_operation is True
+        response = MagicMock()
+        response.statement_id = "test-statement-123"
+        response.status.state = CommandState.SUCCEEDED
+        response.manifest.is_volume_operation = True
+        response.manifest.result_compression = "NONE"
+        response.manifest.format = "JSON_ARRAY"
+
+        # Mock the _extract_description_from_manifest method to return None
+        with patch.object(
+            sea_client, "_extract_description_from_manifest", return_value=None
+        ):
+            result = sea_client._results_message_to_execute_response(response)
+            assert result.is_staging_operation is True
+
+        # Test when is_volume_operation is False
+        response.manifest.is_volume_operation = False
+        with patch.object(
+            sea_client, "_extract_description_from_manifest", return_value=None
+        ):
+            result = sea_client._results_message_to_execute_response(response)
+            assert result.is_staging_operation is False
+
     def test_get_catalogs(self, sea_client, sea_session_id, mock_cursor):
         """Test the get_catalogs method."""
         # Mock the execute_command method