retry tests fix

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

src/databricks/sql/client.py

@@ -252,17 +252,28 @@ def read(self) -> Optional[OAuthToken]:
             self.client_telemetry_enabled and self.server_telemetry_enabled
         )
 
-        self.session = Session(
-            server_hostname,
-            http_path,
-            http_headers,
-            session_configuration,
-            catalog,
-            schema,
-            _use_arrow_native_complex_types,
-            **kwargs,
-        )
-        self.session.open()
+        try:
+            self.session = Session(
+                server_hostname,
+                http_path,
+                http_headers,
+                session_configuration,
+                catalog,
+                schema,
+                _use_arrow_native_complex_types,
+                **kwargs,
+            )
+            self.session.open()
+        except Exception as e:
+            TelemetryClientFactory.connection_failure_log(
+                error_name="Exception",
+                error_message=str(e),
+                host_url=server_hostname,
+                http_path=http_path,
+                port=self.session.port,
+                user_agent=self.session.useragent_header if self.session else None,
+            )
+            raise e
 
         self.use_inline_params = self._set_use_inline_params_with_warning(
             kwargs.get("use_inline_params", False)

src/databricks/sql/session.py

@@ -10,7 +10,6 @@
 from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.backend.types import SessionId
-from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 
 logger = logging.getLogger(__name__)
 
@@ -42,18 +41,9 @@ def __init__(
         self.schema = schema
         self.http_path = http_path
 
-        try:
-            self.auth_provider = get_python_sql_connector_auth_provider(
-                server_hostname, **kwargs
-            )
-        except Exception as e:
-            TelemetryClientFactory.connection_failure_log(
-                error_name="Exception",
-                error_message=str(e),
-                host_url=server_hostname,
-                http_path=http_path,
-                port=self.port,
-            )
+        self.auth_provider = get_python_sql_connector_auth_provider(
+            server_hostname, **kwargs
+        )
 
         user_agent_entry = kwargs.get("user_agent_entry")
         if user_agent_entry is None:
@@ -85,46 +75,25 @@ def __init__(
             tls_client_cert_key_password=kwargs.get("_tls_client_cert_key_password"),
         )
 
-        try:
-            self.backend: DatabricksClient = ThriftDatabricksClient(
-                self.host,
-                self.port,
-                http_path,
-                (http_headers or []) + base_headers,
-                self.auth_provider,
-                ssl_options=self._ssl_options,
-                _use_arrow_native_complex_types=_use_arrow_native_complex_types,
-                **kwargs,
-            )
-        except Exception as e:
-            TelemetryClientFactory.connection_failure_log(
-                error_name="Exception",
-                error_message=str(e),
-                host_url=server_hostname,
-                http_path=http_path,
-                port=self.port,
-                user_agent=self.useragent_header,
-            )
+        self.backend: DatabricksClient = ThriftDatabricksClient(
+            self.host,
+            self.port,
+            http_path,
+            (http_headers or []) + base_headers,
+            self.auth_provider,
+            ssl_options=self._ssl_options,
+            _use_arrow_native_complex_types=_use_arrow_native_complex_types,
+            **kwargs,
+        )
 
         self.protocol_version = None
 
     def open(self):
-        try:
-            self._session_id = self.backend.open_session(
-                session_configuration=self.session_configuration,
-                catalog=self.catalog,
-                schema=self.schema,
-            )
-        except Exception as e:
-            TelemetryClientFactory.connection_failure_log(
-                error_name="Exception",
-                error_message=str(e),
-                host_url=self.host,
-                http_path=self.http_path,
-                port=self.port,
-                user_agent=self.useragent_header,
-            )
-            raise e
+        self._session_id = self.backend.open_session(
+            session_configuration=self.session_configuration,
+            catalog=self.catalog,
+            schema=self.schema,
+        )
 
         self.protocol_version = self.get_protocol_version(self._session_id)
         self.is_open = True

tests/e2e/common/retry_test_mixins.py

@@ -127,7 +127,8 @@ class PySQLRetryTestsMixin:
         "_retry_delay_default": 0.5,
     }
 
-    def test_retry_urllib3_settings_are_honored(self):
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
+    def test_retry_urllib3_settings_are_honored(self, mock_send_telemetry):
         """Databricks overrides some of urllib3's configuration. This tests confirms that what configuration
         we DON'T override is preserved in urllib3's internals
         """
@@ -147,7 +148,8 @@ def test_retry_urllib3_settings_are_honored(self):
         assert rp.read == 11
         assert rp.redirect == 12
 
-    def test_oserror_retries(self):
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
+    def test_oserror_retries(self, mock_send_telemetry):
         """If a network error occurs during make_request, the request is retried according to policy"""
         with patch(
             "urllib3.connectionpool.HTTPSConnectionPool._validate_conn",
@@ -159,7 +161,8 @@ def test_oserror_retries(self):
 
             assert mock_validate_conn.call_count == 6
 
-    def test_retry_max_count_not_exceeded(self):
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
+    def test_retry_max_count_not_exceeded(self, mock_send_telemetry):
         """GIVEN the max_attempts_count is 5
         WHEN the server sends nothing but 429 responses
         THEN the connector issues six request (original plus five retries)
@@ -171,7 +174,8 @@ def test_retry_max_count_not_exceeded(self):
                     pass
             assert mock_obj.return_value.getresponse.call_count == 6
 
-    def test_retry_exponential_backoff(self):
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
+    def test_retry_exponential_backoff(self, mock_send_telemetry):
         """GIVEN the retry policy is configured for reasonable exponential backoff
         WHEN the server sends nothing but 429 responses with retry-afters
         THEN the connector will use those retry-afters values as floor
@@ -338,7 +342,8 @@ def test_retry_abort_close_operation_on_404(self, caplog):
                             "Operation was canceled by a prior request" in caplog.text
                         )
 
-    def test_retry_max_redirects_raises_too_many_redirects_exception(self):
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._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
@@ -362,7 +367,8 @@ def test_retry_max_redirects_raises_too_many_redirects_exception(self):
             # Total call count should be 2 (original + 1 retry)
             assert mock_obj.return_value.getresponse.call_count == expected_call_count
 
-    def test_retry_max_redirects_unset_doesnt_redirect_forever(self):
+    @patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._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