Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commitfb70b70

Browse files
committed
Merge branch 'main' into telemetry-server-flag
Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>
2 parents08c0bd8 +c6f4a27 commitfb70b70

File tree

6 files changed

+103
-20
lines changed

6 files changed

+103
-20
lines changed

‎src/databricks/sql/client.py‎

Lines changed: 25 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
NotSupportedError,
2323
ProgrammingError,
2424
)
25+
2526
fromdatabricks.sql.thrift_api.TCLIServiceimportttypes
2627
fromdatabricks.sql.backend.thrift_backendimportThriftDatabricksClient
2728
fromdatabricks.sql.backend.databricks_clientimportDatabricksClient
@@ -245,17 +246,30 @@ def read(self) -> Optional[OAuthToken]:
245246
self.use_cloud_fetch=kwargs.get("use_cloud_fetch",True)
246247
self._cursors= []# type: List[Cursor]
247248

248-
self.session=Session(
249-
server_hostname,
250-
http_path,
251-
http_headers,
252-
session_configuration,
253-
catalog,
254-
schema,
255-
_use_arrow_native_complex_types,
256-
**kwargs,
257-
)
258-
self.session.open()
249+
try:
250+
self.session=Session(
251+
server_hostname,
252+
http_path,
253+
http_headers,
254+
session_configuration,
255+
catalog,
256+
schema,
257+
_use_arrow_native_complex_types,
258+
**kwargs,
259+
)
260+
self.session.open()
261+
exceptExceptionase:
262+
TelemetryClientFactory.connection_failure_log(
263+
error_name="Exception",
264+
error_message=str(e),
265+
host_url=server_hostname,
266+
http_path=http_path,
267+
port=kwargs.get("_port",443),
268+
user_agent=self.session.useragent_header
269+
ifhasattr(self,"session")
270+
elseNone,
271+
)
272+
raisee
259273

260274
self.use_inline_params=self._set_use_inline_params_with_warning(
261275
kwargs.get("use_inline_params",False)

‎src/databricks/sql/session.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ def __init__(
3939
self.session_configuration=session_configuration
4040
self.catalog=catalog
4141
self.schema=schema
42+
self.http_path=http_path
4243

4344
self.auth_provider=get_python_sql_connector_auth_provider(
4445
server_hostname,**kwargs
@@ -93,6 +94,7 @@ def open(self):
9394
catalog=self.catalog,
9495
schema=self.schema,
9596
)
97+
9698
self.protocol_version=self.get_protocol_version(self._session_id)
9799
self.is_open=True
98100
logger.info("Successfully opened session %s",str(self.guid_hex))

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,9 +149,9 @@ class TelemetryEvent(JsonSerializableMixin):
149149
operation_latency_ms (Optional[int]): Operation latency in milliseconds
150150
"""
151151

152-
session_id:str
153152
system_configuration:DriverSystemConfiguration
154153
driver_connection_params:DriverConnectionParameters
154+
session_id:Optional[str]=None
155155
sql_statement_id:Optional[str]=None
156156
auth_type:Optional[str]=None
157157
vol_operation:Optional[DriverVolumeOperation]=None

‎src/databricks/sql/telemetry/telemetry_client.py‎

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,20 @@
88
TelemetryEvent,
99
DriverSystemConfiguration,
1010
DriverErrorInfo,
11+
DriverConnectionParameters,
12+
HostDetails,
1113
)
1214
fromdatabricks.sql.telemetry.models.frontend_logsimport (
1315
TelemetryFrontendLog,
1416
TelemetryClientContext,
1517
FrontendLogContext,
1618
FrontendLogEntry,
1719
)
18-
fromdatabricks.sql.telemetry.models.enumsimportAuthMech,AuthFlow
20+
fromdatabricks.sql.telemetry.models.enumsimport (
21+
AuthMech,
22+
AuthFlow,
23+
DatabricksClientType,
24+
)
1925
fromdatabricks.sql.telemetry.models.endpoint_modelsimport (
2026
TelemetryRequest,
2127
TelemetryResponse,
@@ -451,3 +457,35 @@ def close(session_id_hex):
451457
logger.debug("Failed to shutdown thread pool executor: %s",e)
452458
TelemetryClientFactory._executor=None
453459
TelemetryClientFactory._initialized=False
460+
461+
@staticmethod
462+
defconnection_failure_log(
463+
error_name:str,
464+
error_message:str,
465+
host_url:str,
466+
http_path:str,
467+
port:int,
468+
user_agent:Optional[str]=None,
469+
):
470+
"""Send error telemetry when connection creation fails, without requiring a session"""
471+
472+
UNAUTH_DUMMY_SESSION_ID="unauth_session_id"
473+
474+
TelemetryClientFactory.initialize_telemetry_client(
475+
telemetry_enabled=True,
476+
session_id_hex=UNAUTH_DUMMY_SESSION_ID,
477+
auth_provider=None,
478+
host_url=host_url,
479+
)
480+
481+
telemetry_client=TelemetryClientFactory.get_telemetry_client(
482+
UNAUTH_DUMMY_SESSION_ID
483+
)
484+
telemetry_client._driver_connection_params=DriverConnectionParameters(
485+
http_path=http_path,
486+
mode=DatabricksClientType.THRIFT,# TODO: Add SEA mode
487+
host_info=HostDetails(host_url=host_url,port=port),
488+
)
489+
telemetry_client._user_agent=user_agent
490+
491+
telemetry_client.export_failure_log(error_name,error_message)

‎tests/e2e/common/retry_test_mixins.py‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,8 @@ class PySQLRetryTestsMixin:
127127
"_retry_delay_default":0.5,
128128
}
129129

130-
deftest_retry_urllib3_settings_are_honored(self):
130+
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
131+
deftest_retry_urllib3_settings_are_honored(self,mock_send_telemetry):
131132
"""Databricks overrides some of urllib3's configuration. This tests confirms that what configuration
132133
we DON'T override is preserved in urllib3's internals
133134
"""
@@ -147,7 +148,8 @@ def test_retry_urllib3_settings_are_honored(self):
147148
assertrp.read==11
148149
assertrp.redirect==12
149150

150-
deftest_oserror_retries(self):
151+
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
152+
deftest_oserror_retries(self,mock_send_telemetry):
151153
"""If a network error occurs during make_request, the request is retried according to policy"""
152154
withpatch(
153155
"urllib3.connectionpool.HTTPSConnectionPool._validate_conn",
@@ -159,7 +161,8 @@ def test_oserror_retries(self):
159161

160162
assertmock_validate_conn.call_count==6
161163

162-
deftest_retry_max_count_not_exceeded(self):
164+
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
165+
deftest_retry_max_count_not_exceeded(self,mock_send_telemetry):
163166
"""GIVEN the max_attempts_count is 5
164167
WHEN the server sends nothing but 429 responses
165168
THEN the connector issues six request (original plus five retries)
@@ -171,7 +174,8 @@ def test_retry_max_count_not_exceeded(self):
171174
pass
172175
assertmock_obj.return_value.getresponse.call_count==6
173176

174-
deftest_retry_exponential_backoff(self):
177+
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
178+
deftest_retry_exponential_backoff(self,mock_send_telemetry):
175179
"""GIVEN the retry policy is configured for reasonable exponential backoff
176180
WHEN the server sends nothing but 429 responses with retry-afters
177181
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):
338342
"Operation was canceled by a prior request"incaplog.text
339343
)
340344

341-
deftest_retry_max_redirects_raises_too_many_redirects_exception(self):
345+
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
346+
deftest_retry_max_redirects_raises_too_many_redirects_exception(self,mock_send_telemetry):
342347
"""GIVEN the connector is configured with a custom max_redirects
343348
WHEN the DatabricksRetryPolicy is created
344349
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):
362367
# Total call count should be 2 (original + 1 retry)
363368
assertmock_obj.return_value.getresponse.call_count==expected_call_count
364369

365-
deftest_retry_max_redirects_unset_doesnt_redirect_forever(self):
370+
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient._send_telemetry")
371+
deftest_retry_max_redirects_unset_doesnt_redirect_forever(self,mock_send_telemetry):
366372
"""GIVEN the connector is configured without a custom max_redirects
367373
WHEN the DatabricksRetryPolicy is used
368374
THEN the connector raises a MaxRedirectsError if that number is exceeded

‎tests/unit/test_telemetry.py‎

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
NoopTelemetryClient,
99
TelemetryClientFactory,
1010
TelemetryHelper,
11-
BaseTelemetryClient,
1211
)
1312
fromdatabricks.sql.telemetry.models.enumsimportAuthMech,AuthFlow
1413
fromdatabricks.sql.auth.authenticatorsimport (
@@ -290,3 +289,27 @@ def test_factory_shutdown_flow(self):
290289
TelemetryClientFactory.close(session2)
291290
assertTelemetryClientFactory._initializedisFalse
292291
assertTelemetryClientFactory._executorisNone
292+
293+
@patch("databricks.sql.telemetry.telemetry_client.TelemetryClient.export_failure_log")
294+
@patch("databricks.sql.client.Session")
295+
deftest_connection_failure_sends_correct_telemetry_payload(
296+
self,mock_session,mock_export_failure_log
297+
):
298+
"""
299+
Verify that a connection failure constructs and sends the correct
300+
telemetry payload via _send_telemetry.
301+
"""
302+
303+
error_message="Could not connect to host"
304+
mock_session.side_effect=Exception(error_message)
305+
306+
try:
307+
fromdatabricksimportsql
308+
sql.connect(server_hostname="test-host",http_path="/test-path")
309+
exceptExceptionase:
310+
assertstr(e)==error_message
311+
312+
mock_export_failure_log.assert_called_once()
313+
call_arguments=mock_export_failure_log.call_args
314+
assertcall_arguments[0][0]=="Exception"
315+
assertcall_arguments[0][1]==error_message

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp