fix imports

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

Author: varun-edachali-dbx

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

@@ -19,7 +19,7 @@
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
 
-from databricks.sql.result_set import SeaResultSet
+from databricks.sql.backend.sea.result_set import SeaResultSet
 
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.backend.types import (

src/databricks/sql/backend/thrift_backend.py

@@ -46,7 +46,6 @@
 
 from databricks.sql.utils import (
     ThriftResultSetQueueFactory,
-    ResultSetQueueFactory,
     _bound,
     RequestErrorInfo,
     NoRetryReason,
@@ -605,7 +604,7 @@ def open_session(self, session_configuration, catalog, schema) -> SessionId:
             session_id = SessionId.from_thrift_handle(
                 response.sessionHandle, properties
             )
-            self._session_id_hex = session_id.hex_guid
+            self._session_id_hex = session_id.guid_hex
             return session_id
         except:
             self._transport.close()

src/databricks/sql/session.py

@@ -40,7 +40,7 @@ def __init__(
         self.session_configuration = session_configuration
         self.catalog = catalog
         self.schema = schema
-        self.http_path = http_path 
+        self.http_path = http_path
 
         self.auth_provider = get_python_sql_connector_auth_provider(
             server_hostname, **kwargs
@@ -183,4 +183,4 @@ def close(self) -> None:
         except Exception as e:
             logger.error(f"Attempt to close session raised a local exception: {e}")
 
-        self.is_open = False
+        self.is_open = False

tests/e2e/common/retry_test_mixins.py

@@ -343,7 +343,9 @@ def test_retry_abort_close_operation_on_404(self, caplog):
                         )
 
     @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
-    def test_retry_max_redirects_raises_too_many_redirects_exception(self, mock_send_telemetry):
+    def test_retry_max_redirects_raises_too_many_redirects_exception(
+        self, mock_send_telemetry
+    ):
         """GIVEN the connector is configured with a custom max_redirects
         WHEN the DatabricksRetryPolicy is created
         THEN the connector raises a MaxRedirectsError if that number is exceeded
@@ -368,7 +370,9 @@ def test_retry_max_redirects_raises_too_many_redirects_exception(self, mock_send
             assert mock_obj.return_value.getresponse.call_count == expected_call_count
 
     @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
-    def test_retry_max_redirects_unset_doesnt_redirect_forever(self, mock_send_telemetry):
+    def test_retry_max_redirects_unset_doesnt_redirect_forever(
+        self, mock_send_telemetry
+    ):
         """GIVEN the connector is configured without a custom max_redirects
         WHEN the DatabricksRetryPolicy is used
         THEN the connector raises a MaxRedirectsError if that number is exceeded

tests/unit/test_downloader.py

@@ -79,7 +79,12 @@ def test_run_get_response_not_ok(self, mock_time):
             return_value=create_response(status_code=404, _content=b"1234"),
         ):
             d = downloader.ResultSetDownloadHandler(
-                settings, result_link, ssl_options=SSLOptions()
+                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()
@@ -99,7 +104,12 @@ def test_run_uncompressed_successful(self, mock_time):
             return_value=create_response(status_code=200, _content=file_bytes),
         ):
             d = downloader.ResultSetDownloadHandler(
-                settings, result_link, ssl_options=SSLOptions()
+                settings,
+                result_link,
+                ssl_options=SSLOptions(),
+                chunk_id=0,
+                session_id_hex=Mock(),
+                statement_id=Mock(),
             )
             file = d.run()
 
@@ -120,7 +130,12 @@ def test_run_compressed_successful(self, mock_time):
             return_value=create_response(status_code=200, _content=compressed_bytes),
         ):
             d = downloader.ResultSetDownloadHandler(
-                settings, result_link, ssl_options=SSLOptions()
+                settings,
+                result_link,
+                ssl_options=SSLOptions(),
+                chunk_id=0,
+                session_id_hex=Mock(),
+                statement_id=Mock(),
             )
             file = d.run()
 
@@ -137,7 +152,12 @@ def test_download_connection_error(self, mock_time):
 
         with patch.object(http_client, "execute", side_effect=ConnectionError("foo")):
             d = downloader.ResultSetDownloadHandler(
-                settings, result_link, ssl_options=SSLOptions()
+                settings,
+                result_link,
+                ssl_options=SSLOptions(),
+                chunk_id=0,
+                session_id_hex=Mock(),
+                statement_id=Mock(),
             )
             with self.assertRaises(ConnectionError):
                 d.run()
@@ -152,7 +172,12 @@ def test_download_timeout(self, mock_time):
 
         with patch.object(http_client, "execute", side_effect=TimeoutError("foo")):
             d = downloader.ResultSetDownloadHandler(
-                settings, result_link, ssl_options=SSLOptions()
+                settings,
+                result_link,
+                ssl_options=SSLOptions(),
+                chunk_id=0,
+                session_id_hex=Mock(),
+                statement_id=Mock(),
             )
             with self.assertRaises(TimeoutError):
                 d.run()

tests/unit/test_fetches_bench.py

@@ -33,12 +33,10 @@ def make_dummy_result_set_from_initial_results(arrow_table):
         rs = client.ResultSet(
             connection=None,
             backend=None,
-            backend=None,
             execute_response=ExecuteResponse(
                 status=None,
                 has_been_closed_server_side=True,
                 is_direct_results=False,
-                is_direct_results=False,
                 description=Mock(),
                 command_id=None,
                 arrow_schema_bytes=arrow_table.schema,

tests/unit/test_filters.py

@@ -77,7 +77,7 @@ def test_filter_by_column_values(self):
             "databricks.sql.backend.sea.utils.filters.isinstance", return_value=True
         ):
             with patch(
-                "databricks.sql.result_set.SeaResultSet"
+                "databricks.sql.backend.sea.result_set.SeaResultSet"
             ) as mock_sea_result_set_class:
                 mock_instance = MagicMock()
                 mock_sea_result_set_class.return_value = mock_instance
@@ -104,7 +104,7 @@ def test_filter_by_column_values(self):
             "databricks.sql.backend.sea.utils.filters.isinstance", return_value=True
         ):
             with patch(
-                "databricks.sql.result_set.SeaResultSet"
+                "databricks.sql.backend.sea.result_set.SeaResultSet"
             ) as mock_sea_result_set_class:
                 mock_instance = MagicMock()
                 mock_sea_result_set_class.return_value = mock_instance

tests/unit/test_sea_backend.py

@@ -737,7 +737,7 @@ def test_get_schemas(self, sea_client, sea_session_id, mock_cursor):
     def test_get_tables(self, sea_client, sea_session_id, mock_cursor):
         """Test the get_tables method with various parameter combinations."""
         # Mock the execute_command method
-        from databricks.sql.result_set import SeaResultSet
+        from databricks.sql.backend.sea.result_set import SeaResultSet
 
         mock_result_set = Mock(spec=SeaResultSet)
 

tests/unit/test_telemetry.py

@@ -291,7 +291,9 @@ def test_factory_shutdown_flow(self):
         assert TelemetryClientFactory._initialized is False
         assert TelemetryClientFactory._executor is None
 
-    @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient.export_failure_log")
+    @patch(
+        "databricks.sql.telemetry.telemetry_client.TelemetryClient.export_failure_log"
+    )
     @patch("databricks.sql.client.Session")
     def test_connection_failure_sends_correct_telemetry_payload(
         self, mock_session, mock_export_failure_log
@@ -306,11 +308,12 @@ def test_connection_failure_sends_correct_telemetry_payload(
 
         try:
             from databricks import sql
+
             sql.connect(server_hostname="test-host", http_path="/test-path")
         except Exception as e:
             assert str(e) == error_message
 
         mock_export_failure_log.assert_called_once()
         call_arguments = mock_export_failure_log.call_args
         assert call_arguments[0][0] == "Exception"
-        assert call_arguments[0][1] == error_message
+        assert call_arguments[0][1] == error_message