databricks · varun-edachali-dbx · Jun 26, 2025 · Jun 20, 2025 · Jun 26, 2025 · Jun 26, 2025
diff --git a/src/databricks/sql/backend/sea/backend.py b/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,
        )

diff --git a/src/databricks/sql/backend/sea/models/base.py b/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
diff --git a/src/databricks/sql/backend/sea/models/responses.py b/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),
    )



diff --git a/tests/unit/test_sea_backend.py b/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