Merge branch 'main' into telemetry-server-flag

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

Author: saishreeeee

src/databricks/sql/client.py

@@ -22,6 +22,7 @@
     NotSupportedError,
     ProgrammingError,
 )
+
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
 from databricks.sql.backend.databricks_client import DatabricksClient
@@ -245,17 +246,30 @@ def read(self) -> Optional[OAuthToken]:
         self.use_cloud_fetch = kwargs.get("use_cloud_fetch", True)
         self._cursors = []  # type: List[Cursor]
 
-        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=kwargs.get("_port", 443),
+                user_agent=self.session.useragent_header
+                if hasattr(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

@@ -39,6 +39,7 @@ def __init__(
         self.session_configuration = session_configuration
         self.catalog = catalog
         self.schema = schema
+        self.http_path = http_path
 
         self.auth_provider = get_python_sql_connector_auth_provider(
             server_hostname, **kwargs
@@ -93,6 +94,7 @@ def open(self):
             catalog=self.catalog,
             schema=self.schema,
         )
+
         self.protocol_version = self.get_protocol_version(self._session_id)
         self.is_open = True
         logger.info("Successfully opened session %s", str(self.guid_hex))

src/databricks/sql/telemetry/models/event.py

@@ -149,9 +149,9 @@ class TelemetryEvent(JsonSerializableMixin):
         operation_latency_ms (Optional[int]): Operation latency in milliseconds
     """
 
-    session_id: str
     system_configuration: DriverSystemConfiguration
     driver_connection_params: DriverConnectionParameters
+    session_id: Optional[str] = None
     sql_statement_id: Optional[str] = None
     auth_type: Optional[str] = None
     vol_operation: Optional[DriverVolumeOperation] = None

src/databricks/sql/telemetry/telemetry_client.py

@@ -8,14 +8,20 @@
     TelemetryEvent,
     DriverSystemConfiguration,
     DriverErrorInfo,
+    DriverConnectionParameters,
+    HostDetails,
 )
 from databricks.sql.telemetry.models.frontend_logs import (
     TelemetryFrontendLog,
     TelemetryClientContext,
     FrontendLogContext,
     FrontendLogEntry,
 )
-from databricks.sql.telemetry.models.enums import AuthMech, AuthFlow
+from databricks.sql.telemetry.models.enums import (
+    AuthMech,
+    AuthFlow,
+    DatabricksClientType,
+)
 from databricks.sql.telemetry.models.endpoint_models import (
     TelemetryRequest,
     TelemetryResponse,
@@ -451,3 +457,35 @@ def close(session_id_hex):
                     logger.debug("Failed to shutdown thread pool executor: %s", e)
                 TelemetryClientFactory._executor = None
                 TelemetryClientFactory._initialized = False
+
+    @staticmethod
+    def connection_failure_log(
+        error_name: str,
+        error_message: str,
+        host_url: str,
+        http_path: str,
+        port: int,
+        user_agent: Optional[str] = None,
+    ):
+        """Send error telemetry when connection creation fails, without requiring a session"""
+
+        UNAUTH_DUMMY_SESSION_ID = "unauth_session_id"
+
+        TelemetryClientFactory.initialize_telemetry_client(
+            telemetry_enabled=True,
+            session_id_hex=UNAUTH_DUMMY_SESSION_ID,
+            auth_provider=None,
+            host_url=host_url,
+        )
+
+        telemetry_client = TelemetryClientFactory.get_telemetry_client(
+            UNAUTH_DUMMY_SESSION_ID
+        )
+        telemetry_client._driver_connection_params = DriverConnectionParameters(
+            http_path=http_path,
+            mode=DatabricksClientType.THRIFT,  # TODO: Add SEA mode
+            host_info=HostDetails(host_url=host_url, port=port),
+        )
+        telemetry_client._user_agent = user_agent
+
+        telemetry_client.export_failure_log(error_name, error_message)

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

tests/unit/test_telemetry.py

@@ -8,7 +8,6 @@
     NoopTelemetryClient,
     TelemetryClientFactory,
     TelemetryHelper,
-    BaseTelemetryClient,
 )
 from databricks.sql.telemetry.models.enums import AuthMech, AuthFlow
 from databricks.sql.auth.authenticators import (
@@ -290,3 +289,27 @@ def test_factory_shutdown_flow(self):
         TelemetryClientFactory.close(session2)
         assert TelemetryClientFactory._initialized is False
         assert TelemetryClientFactory._executor is None
+
+    @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
+    ):
+        """
+        Verify that a connection failure constructs and sends the correct
+        telemetry payload via _send_telemetry.
+        """
+
+        error_message = "Could not connect to host"
+        mock_session.side_effect = Exception(error_message)
+
+        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