fix merge artifacts

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

Author: varun-edachali-dbx

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

@@ -70,19 +70,43 @@ def _create_filtered_manifest(result_set: SeaResultSet, new_row_count: int):
             result_set: Original result set to copy manifest from
             new_row_count: New total row count for filtered data
 
-        from databricks.sql.backend.sea.backend import SeaDatabricksClient
+        Returns:
+            Updated manifest copy
+        """
+        filtered_manifest = deepcopy(result_set.manifest)
+        filtered_manifest.total_row_count = new_row_count
+        return filtered_manifest
+
+    @staticmethod
+    def _create_filtered_result_set(
+        result_set: SeaResultSet,
+        result_data: ResultData,
+        row_count: int,
+    ) -> "SeaResultSet":
+        """
+        Create a new filtered SeaResultSet with the provided data.
+
+        Args:
+            result_set: Original result set to copy parameters from
+            result_data: New result data for the filtered set
+            row_count: Number of rows in the filtered data
+
+        Returns:
+            New filtered SeaResultSet
+        """
         from databricks.sql.backend.sea.result_set import SeaResultSet
 
-        # Create a new SeaResultSet with the filtered data
-        manifest = result_set.manifest
-        manifest.total_row_count = len(filtered_rows)
+        execute_response = ResultSetFilter._create_execute_response(result_set)
+        filtered_manifest = ResultSetFilter._create_filtered_manifest(
+            result_set, row_count
+        )
 
-        filtered_result_set = SeaResultSet(
+        return SeaResultSet(
             connection=result_set.connection,
             execute_response=execute_response,
             sea_client=cast(SeaDatabricksClient, result_set.backend),
             result_data=result_data,
-            manifest=manifest,
+            manifest=filtered_manifest,
             buffer_size_bytes=result_set.buffer_size_bytes,
             arraysize=result_set.arraysize,
         )

src/databricks/sql/backend/thrift_backend.py

@@ -895,7 +895,6 @@ def get_execution_result(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
-            session_id_hex=self._session_id_hex,
             buffer_size_bytes=cursor.buffer_size_bytes,
             arraysize=cursor.arraysize,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,

tests/unit/test_client.py

@@ -127,7 +127,6 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
                     connection=connection,
                     execute_response=mock_execute_response,
                     thrift_client=mock_backend,
-                    session_id_hex=Mock(),
                 )
 
                 # Mock execute_command to return our real result set
@@ -187,7 +186,6 @@ def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
             connection=mock_connection,
             execute_response=Mock(),
             thrift_client=mock_backend,
-            session_id_hex=Mock(),
         )
         result_set.results = mock_results
 
@@ -217,7 +215,6 @@ def test_closing_result_set_hard_closes_commands(self):
             mock_connection,
             mock_results_response,
             mock_thrift_backend,
-            session_id_hex=Mock(),
         )
         result_set.results = mock_results
 
@@ -265,9 +262,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_downloader.py

@@ -26,12 +26,14 @@ class DownloaderTests(unittest.TestCase):
     def _setup_time_mock_for_download(self, mock_time, end_time):
         """Helper to setup time mock that handles logging system calls."""
         call_count = [0]
+
         def time_side_effect():
             call_count[0] += 1
             if call_count[0] <= 2:  # First two calls (validation, start_time)
                 return 1000
             else:  # All subsequent calls (logging, duration calculation)
                 return end_time
+
         mock_time.side_effect = time_side_effect
 
     @patch("time.time", return_value=1000)
@@ -104,7 +106,7 @@ def test_run_get_response_not_ok(self, mock_time):
     @patch("time.time")
     def test_run_uncompressed_successful(self, mock_time):
         self._setup_time_mock_for_download(mock_time, 1000.5)
-        
+
         http_client = DatabricksHttpClient.get_instance()
         file_bytes = b"1234567890" * 10
         settings = Mock(link_expiry_buffer_secs=0, download_timeout=0, use_proxy=False)
@@ -133,7 +135,7 @@ def test_run_uncompressed_successful(self, mock_time):
     @patch("time.time")
     def test_run_compressed_successful(self, mock_time):
         self._setup_time_mock_for_download(mock_time, 1000.2)
-        
+
         http_client = DatabricksHttpClient.get_instance()
         file_bytes = b"1234567890" * 10
         compressed_bytes = b'\x04"M\x18h@d\x00\x00\x00\x00\x00\x00\x00#\x14\x00\x00\x00\xaf1234567890\n\x00BP67890\x00\x00\x00\x00'

tests/unit/test_fetches.py

@@ -63,7 +63,6 @@ def make_dummy_result_set_from_initial_results(initial_results):
             ),
             thrift_client=mock_thrift_backend,
             t_row_set=None,
-            session_id_hex=Mock(),
         )
         return rs
 
@@ -108,7 +107,6 @@ def fetch_results(
                 is_staging_operation=False,
             ),
             thrift_client=mock_thrift_backend,
-            session_id_hex=Mock(),
         )
         return rs
 

tests/unit/test_filters.py

@@ -132,21 +132,25 @@ def test_filter_tables_by_type(self):
             ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, table_types)
             args, kwargs = mock_filter.call_args
             self.assertEqual(args[0], self.mock_sea_result_set)
-            self.assertEqual(args[1], 5)  # Table type column index
-            self.assertEqual(args[2], table_types)
+            self.assertEqual(kwargs.get("column_index"), 5)  # Table type column index
+            self.assertEqual(kwargs.get("allowed_values"), table_types)
             self.assertEqual(kwargs.get("case_sensitive"), True)
 
         # Case 2: Default table types (None or empty list)
         with patch.object(ResultSetFilter, "_filter_json_result_set") as mock_filter:
             # Test with None
             ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, None)
             args, kwargs = mock_filter.call_args
-            self.assertEqual(args[2], ["TABLE", "VIEW", "SYSTEM TABLE"])
+            self.assertEqual(
+                kwargs.get("allowed_values"), ["TABLE", "VIEW", "SYSTEM TABLE"]
+            )
 
             # Test with empty list
             ResultSetFilter.filter_tables_by_type(self.mock_sea_result_set, [])
             args, kwargs = mock_filter.call_args
-            self.assertEqual(args[2], ["TABLE", "VIEW", "SYSTEM TABLE"])
+            self.assertEqual(
+                kwargs.get("allowed_values"), ["TABLE", "VIEW", "SYSTEM TABLE"]
+            )
 
 
 if __name__ == "__main__":

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