diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index e4166f117..0e0486614 100755 --- a/src/databricks/sql/client.py +++ b/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 @@ -251,17 +252,30 @@ 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=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) diff --git a/src/databricks/sql/session.py b/src/databricks/sql/session.py index 251f502df..9278ff167 100644 --- a/src/databricks/sql/session.py +++ b/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)) diff --git a/src/databricks/sql/telemetry/models/event.py b/src/databricks/sql/telemetry/models/event.py index f5496deec..a155c7597 100644 --- a/src/databricks/sql/telemetry/models/event.py +++ b/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 diff --git a/src/databricks/sql/telemetry/telemetry_client.py b/src/databricks/sql/telemetry/telemetry_client.py index 5eb8c6ed0..2c389513a 100644 --- a/src/databricks/sql/telemetry/telemetry_client.py +++ b/src/databricks/sql/telemetry/telemetry_client.py @@ -8,6 +8,8 @@ TelemetryEvent, DriverSystemConfiguration, DriverErrorInfo, + DriverConnectionParameters, + HostDetails, ) from databricks.sql.telemetry.models.frontend_logs import ( TelemetryFrontendLog, @@ -15,7 +17,11 @@ 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, @@ -431,3 +437,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) diff --git a/tests/e2e/common/retry_test_mixins.py b/tests/e2e/common/retry_test_mixins.py index b5d01a45d..66c15ad1c 100755 --- a/tests/e2e/common/retry_test_mixins.py +++ b/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 diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index dc1c7d630..4e6e928ab 100644 --- a/tests/unit/test_telemetry.py +++ b/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 \ No newline at end of file