From 004c3dccdbb10212567868de3b594b57b8addc34 Mon Sep 17 00:00:00 2001 From: varun-edachali-dbx Date: Fri, 20 Jun 2025 03:19:28 +0000 Subject: [PATCH 1/2] assign manifest.is_volume_operation to is_staging_operation Signed-off-by: varun-edachali-dbx --- src/databricks/sql/backend/sea/backend.py | 2 +- src/databricks/sql/backend/sea/models/base.py | 2 +- .../sql/backend/sea/models/responses.py | 2 +- tests/unit/test_sea_backend.py | 25 +++++++++++++++++++ 4 files changed, 28 insertions(+), 3 deletions(-) diff --git a/src/databricks/sql/backend/sea/backend.py b/src/databricks/sql/backend/sea/backend.py index 76903ccd2..7529e61e8 100644 --- a/src/databricks/sql/backend/sea/backend.py +++ b/src/databricks/sql/backend/sea/backend.py @@ -350,7 +350,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 index b12c26eb0..f99e85055 100644 --- 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 index 0baf27ab2..302b32d0c 100644 --- 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 index f30c92ed0..1f42c5a00 100644 --- a/tests/unit/test_sea_backend.py +++ b/tests/unit/test_sea_backend.py @@ -632,6 +632,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_unimplemented_metadata_methods( self, sea_client, sea_session_id, mock_cursor ): From 3cf771bcff36a9af590c696de4c22a3866647571 Mon Sep 17 00:00:00 2001 From: varun-edachali-dbx Date: Thu, 26 Jun 2025 10:05:57 +0000 Subject: [PATCH 2/2] introduce unit test to ensure correct assignment of is_staging_op Signed-off-by: varun-edachali-dbx --- tests/unit/test_sea_backend.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/unit/test_sea_backend.py b/tests/unit/test_sea_backend.py index 6847cded0..bc6768d2b 100644 --- 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