description, partial results (small fixes)

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

Author: varun-edachali-dbx

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

@@ -364,14 +364,14 @@ def __init__(
         # Initialize table and position
         self.table = self._create_next_table()
 
-    def _create_next_table(self) -> Union["pyarrow.Table", None]:
+    def _create_next_table(self) -> "pyarrow.Table":
         """Create next table by retrieving the logical next downloaded file."""
         if self.link_fetcher is None:
-            return None
+            return self._create_empty_table()
 
         chunk_link = self.link_fetcher.get_chunk_link(self._current_chunk_index)
         if chunk_link is None:
-            return None
+            return self._create_empty_table()
 
         row_offset = chunk_link.row_offset
         # NOTE: link has already been submitted to download manager at this point

src/databricks/sql/utils.py

@@ -275,6 +275,7 @@ def next_n_rows(self, num_rows: int) -> "pyarrow.Table":
 
         logger.debug("CloudFetchQueue: trying to get {} next rows".format(num_rows))
         results = self.table.slice(0, 0)
+        partial_result_chunks = [results]
         while num_rows > 0 and self.table.num_rows > 0:
             # Replace current table with the next table if we are at the end of the current table
             if self.table_row_index == self.table.num_rows:
@@ -300,6 +301,7 @@ def remaining_rows(self) -> "pyarrow.Table":
         """
 
         results = self.table.slice(0, 0)
+        partial_result_chunks = [results]
         while self.table.num_rows > 0:
             table_slice = self.table.slice(
                 self.table_row_index, self.table.num_rows - self.table_row_index
@@ -386,6 +388,8 @@ def __init__(
             chunk_id=chunk_id,
         )
 
+        self.num_links_downloaded = 0
+
         self.start_row_index = start_row_offset
         self.result_links = result_links or []
         self.session_id_hex = session_id_hex

tests/unit/test_client.py

@@ -46,7 +46,7 @@ def new(cls):
             is_staging_operation=False,
             command_id=None,
             has_been_closed_server_side=True,
-            is_direct_results=True,
+            has_more_rows=True,
             lz4_compressed=True,
             arrow_schema_bytes=b"schema",
         )
@@ -266,9 +266,7 @@ def test_negative_fetch_throws_exception(self):
         mock_backend = Mock()
         mock_backend.fetch_results.return_value = (Mock(), False, 0)
 
-        result_set = ThriftResultSet(
-            Mock(), Mock(), mock_backend
-        )
+        result_set = ThriftResultSet(Mock(), Mock(), mock_backend)
 
         with self.assertRaises(ValueError) as e:
             result_set.fetchmany(-1)

tests/unit/test_cloud_fetch_queue.py

@@ -69,6 +69,7 @@ def test_initializer_adds_links(self, mock_create_next_table):
             session_id_hex=Mock(),
             statement_id=Mock(),
             chunk_id=0,
+            description=description,
         )
 
         assert (

tests/unit/test_downloader.py

@@ -73,17 +73,26 @@ def test_run_get_response_not_ok(self, mock_time):
         settings.use_proxy = False
         result_link = Mock(expiryTime=1001)
 
-        d = downloader.ResultSetDownloadHandler(
-            settings,
-            result_link,
-            ssl_options=SSLOptions(),
-            chunk_id=0,
-            session_id_hex=Mock(),
-            statement_id=Mock(),
+        # Create a mock response with 404 status
+        mock_response = create_response(status_code=404, _content=b"Not Found")
+        mock_response.raise_for_status = Mock(
+            side_effect=requests.exceptions.HTTPError("404")
         )
-        with self.assertRaises(requests.exceptions.HTTPError) as context:
-            d.run()
-        self.assertTrue("404" in str(context.exception))
+
+        with patch.object(http_client, "execute") as mock_execute:
+            mock_execute.return_value.__enter__.return_value = mock_response
+
+            d = downloader.ResultSetDownloadHandler(
+                settings,
+                result_link,
+                ssl_options=SSLOptions(),
+                chunk_id=0,
+                session_id_hex=Mock(),
+                statement_id=Mock(),
+            )
+            with self.assertRaises(requests.exceptions.HTTPError) as context:
+                d.run()
+            self.assertTrue("404" in str(context.exception))
 
     @patch("time.time", return_value=1000)
     def test_run_uncompressed_successful(self, mock_time):

tests/unit/test_fetches_bench.py

@@ -36,7 +36,7 @@ def make_dummy_result_set_from_initial_results(arrow_table):
             execute_response=ExecuteResponse(
                 status=None,
                 has_been_closed_server_side=True,
-                is_direct_results=False,
+                has_more_rows=False,
                 description=Mock(),
                 command_id=None,
                 arrow_schema_bytes=arrow_table.schema,

tests/unit/test_sea_queue.py

@@ -27,6 +27,11 @@
 import threading
 import time
 
+try:
+    import pyarrow as pa
+except ImportError:
+    pa = None
+
 
 class TestJsonQueue:
     """Test suite for the JsonQueue class."""

tests/unit/test_telemetry_retry.py

@@ -6,7 +6,8 @@
 from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 from databricks.sql.auth.retry import DatabricksRetryPolicy
 
-PATCH_TARGET = 'urllib3.connectionpool.HTTPSConnectionPool._get_conn'
+PATCH_TARGET = "urllib3.connectionpool.HTTPSConnectionPool._get_conn"
+
 
 def create_mock_conn(responses):
     """Creates a mock connection object whose getresponse() method yields a series of responses."""
@@ -16,15 +17,18 @@ def create_mock_conn(responses):
         mock_http_response = MagicMock()
         mock_http_response.status = resp.get("status")
         mock_http_response.headers = resp.get("headers", {})
-        body = resp.get("body", b'{}')
+        body = resp.get("body", b"{}")
         mock_http_response.fp = io.BytesIO(body)
+
         def release():
             mock_http_response.fp.close()
+
         mock_http_response.release_conn = release
         mock_http_responses.append(mock_http_response)
     mock_conn.getresponse.side_effect = mock_http_responses
     return mock_conn
 
+
 class TestTelemetryClientRetries:
     @pytest.fixture(autouse=True)
     def setup_and_teardown(self):
@@ -49,28 +53,28 @@ def get_client(self, session_id, num_retries=3):
             host_url="test.databricks.com",
         )
         client = TelemetryClientFactory.get_telemetry_client(session_id)
-        
+
         retry_policy = DatabricksRetryPolicy(
             delay_min=0.01,
             delay_max=0.02,
             stop_after_attempts_duration=2.0,
-            stop_after_attempts_count=num_retries, 
+            stop_after_attempts_count=num_retries,
             delay_default=0.1,
             force_dangerous_codes=[],
-            urllib3_kwargs={'total': num_retries}
+            urllib3_kwargs={"total": num_retries},
         )
         adapter = client._http_client.session.adapters.get("https://")
         adapter.max_retries = retry_policy
         return client
 
     @pytest.mark.parametrize(
-    "status_code, description",
-    [
-        (401, "Unauthorized"),
-        (403, "Forbidden"),
-        (501, "Not Implemented"),
-        (200, "Success"),
-    ],
+        "status_code, description",
+        [
+            (401, "Unauthorized"),
+            (403, "Forbidden"),
+            (501, "Not Implemented"),
+            (200, "Success"),
+        ],
     )
     def test_non_retryable_status_codes_are_not_retried(self, status_code, description):
         """
@@ -80,7 +84,9 @@ def test_non_retryable_status_codes_are_not_retried(self, status_code, descripti
         client = self.get_client(f"session-{status_code}")
         mock_responses = [{"status": status_code}]
 
-        with patch(PATCH_TARGET, return_value=create_mock_conn(mock_responses)) as mock_get_conn:
+        with patch(
+            PATCH_TARGET, return_value=create_mock_conn(mock_responses)
+        ) as mock_get_conn:
             client.export_failure_log("TestError", "Test message")
             TelemetryClientFactory.close(client._session_id_hex)
 
@@ -92,16 +98,26 @@ def test_exceeds_retry_count_limit(self):
         Verifies that the client respects the Retry-After header and retries on 429, 502, 503.
         """
         num_retries = 3
-        expected_total_calls = num_retries + 1 
+        expected_total_calls = num_retries + 1
         retry_after = 1
         client = self.get_client("session-exceed-limit", num_retries=num_retries)
-        mock_responses = [{"status": 503, "headers": {"Retry-After": str(retry_after)}}, {"status": 429}, {"status": 502}, {"status": 503}]
-        
-        with patch(PATCH_TARGET, return_value=create_mock_conn(mock_responses)) as mock_get_conn:
+        mock_responses = [
+            {"status": 503, "headers": {"Retry-After": str(retry_after)}},
+            {"status": 429},
+            {"status": 502},
+            {"status": 503},
+        ]
+
+        with patch(
+            PATCH_TARGET, return_value=create_mock_conn(mock_responses)
+        ) as mock_get_conn:
             start_time = time.time()
             client.export_failure_log("TestError", "Test message")
             TelemetryClientFactory.close(client._session_id_hex)
             end_time = time.time()
-            
-            assert mock_get_conn.return_value.getresponse.call_count == expected_total_calls
-            assert end_time - start_time > retry_after
+
+            assert (
+                mock_get_conn.return_value.getresponse.call_count
+                == expected_total_calls
+            )
+            assert end_time - start_time > retry_after

tests/unit/test_thrift_backend.py

@@ -1016,10 +1016,10 @@ def test_fall_back_to_hive_schema_if_no_arrow_schema(self, tcli_service_class):
     def test_handle_execute_response_reads_has_more_rows_in_direct_results(
         self, tcli_service_class, build_queue
     ):
-        for is_direct_results, resp_type in itertools.product(
+        for has_more_rows, resp_type in itertools.product(
             [True, False], self.execute_response_types
         ):
-            with self.subTest(is_direct_results=is_direct_results, resp_type=resp_type):
+            with self.subTest(has_more_rows=has_more_rows, resp_type=resp_type):
                 tcli_service_instance = tcli_service_class.return_value
                 results_mock = Mock()
                 results_mock.startRowOffset = 0
@@ -1031,7 +1031,7 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
                     resultSetMetadata=self.metadata_resp,
                     resultSet=ttypes.TFetchResultsResp(
                         status=self.okay_status,
-                        hasMoreRows=is_direct_results,
+                        hasMoreRows=has_more_rows,
                         results=results_mock,
                     ),
                     closeOperation=Mock(),
@@ -1062,10 +1062,10 @@ def test_handle_execute_response_reads_has_more_rows_in_direct_results(
     def test_handle_execute_response_reads_has_more_rows_in_result_response(
         self, tcli_service_class, build_queue
     ):
-        for is_direct_results, resp_type in itertools.product(
+        for has_more_rows, resp_type in itertools.product(
             [True, False], self.execute_response_types
         ):
-            with self.subTest(is_direct_results=is_direct_results, resp_type=resp_type):
+            with self.subTest(has_more_rows=has_more_rows, resp_type=resp_type):
                 tcli_service_instance = tcli_service_class.return_value
                 results_mock = MagicMock()
                 results_mock.startRowOffset = 0
@@ -1078,7 +1078,7 @@ def test_handle_execute_response_reads_has_more_rows_in_result_response(
 
                 fetch_results_resp = ttypes.TFetchResultsResp(
                     status=self.okay_status,
-                    hasMoreRows=is_direct_results,
+                    hasMoreRows=has_more_rows,
                     results=results_mock,
                     resultSetMetadata=ttypes.TGetResultSetMetadataResp(
                         resultFormat=ttypes.TSparkRowSetType.ARROW_BASED_SET
@@ -1112,7 +1112,7 @@ def test_handle_execute_response_reads_has_more_rows_in_result_response(
                     chunk_id=0,
                 )
 
-                self.assertEqual(is_direct_results, has_more_rows_resp)
+                self.assertEqual(has_more_rows, has_more_rows_resp)
 
     @patch("databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True)
     def test_arrow_batches_row_count_are_respected(self, tcli_service_class):