Skip to content

Commit cf0ddc0

Browse files
Move telemetry from auth request to in-band
1 parent cb8ec27 commit cf0ddc0

File tree

4 files changed

+51
-3
lines changed

4 files changed

+51
-3
lines changed

src/snowflake/connector/auth/_auth.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
load_pem_private_key,
1818
)
1919

20-
from .._utils import _core_loader, get_application_path
20+
from .._utils import get_application_path
2121
from ..compat import urlencode
2222
from ..constants import (
2323
DAY_IN_SECONDS,
@@ -144,8 +144,6 @@ def base_auth_data(
144144
platform_detection_timeout_seconds=platform_detection_timeout_seconds,
145145
session_manager=session_manager.clone(max_retries=0),
146146
),
147-
"CORE_LOAD_ERROR": _core_loader.get_load_error(),
148-
"CORE_VERSION": _core_loader.get_core_version(),
149147
},
150148
},
151149
}

src/snowflake/connector/connection.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
from ._utils import (
4343
_DEFAULT_VALUE_SERVER_DOP_CAP_FOR_FILE_TRANSFER,
4444
_VARIABLE_NAME_SERVER_DOP_CAP_FOR_FILE_TRANSFER,
45+
_core_loader,
4546
)
4647
from .auth import (
4748
FIRST_PARTY_AUTHENTICATORS,
@@ -677,6 +678,7 @@ def __init__(
677678

678679
# get the imported modules from sys.modules
679680
self._log_telemetry_imported_packages()
681+
self._log_minicore_import()
680682
# check SNOW-1218851 for long term improvement plan to refactor ocsp code
681683
atexit.register(self._close_at_exit)
682684

@@ -2488,6 +2490,22 @@ def async_query_check_helper(
24882490

24892491
return not found_unfinished_query
24902492

2493+
def _log_minicore_import(self):
2494+
ts = get_time_millis()
2495+
self._log_telemetry(
2496+
TelemetryData.from_telemetry_data_dict(
2497+
from_dict={
2498+
TelemetryField.KEY_TYPE.value: TelemetryField.CORE_IMPORT.value,
2499+
TelemetryField.KEY_VALUE: {
2500+
"CORE_LOAD_ERROR": _core_loader.get_load_error(),
2501+
"CORE_VERSION": _core_loader.get_core_version(),
2502+
},
2503+
},
2504+
timestamp=ts,
2505+
connection=self,
2506+
)
2507+
)
2508+
24912509
def _log_telemetry_imported_packages(self) -> None:
24922510
if self._log_imported_packages_in_telemetry:
24932511
# filter out duplicates caused by submodules

src/snowflake/connector/telemetry.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ class TelemetryField(Enum):
3939
PANDAS_WRITE = "client_write_pandas"
4040
# imported packages along with client
4141
IMPORTED_PACKAGES = "client_imported_packages"
42+
# Core import
43+
CORE_IMPORT = "mini_core_import"
4244
# multi-statement usage
4345
MULTI_STATEMENT = "client_multi_statement_query"
4446
# Keys for telemetry data sent through either in-band or out-of-band telemetry

test/integ/test_connection.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1230,6 +1230,36 @@ def check_packages(message: str, expected_packages: list[str]) -> bool:
12301230
assert len(telemetry_test.records) == 0
12311231

12321232

1233+
@pytest.mark.skipolddriver
1234+
def test_minicore_import_telemetry(conn_cnx, capture_sf_telemetry):
1235+
"""Test that minicore import telemetry (CORE_LOAD_ERROR and CORE_VERSION) is logged."""
1236+
with (
1237+
conn_cnx() as conn,
1238+
capture_sf_telemetry.patch_connection(conn, False) as telemetry_test,
1239+
):
1240+
conn._log_minicore_import()
1241+
assert len(telemetry_test.records) > 0
1242+
# Check that the telemetry record contains the proper structure
1243+
found_minicore_telemetry = False
1244+
for t in telemetry_test.records:
1245+
if (
1246+
t.message.get(TelemetryField.KEY_TYPE.value)
1247+
== TelemetryField.CORE_IMPORT.value
1248+
and TelemetryField.KEY_VALUE.value in t.message
1249+
):
1250+
found_minicore_telemetry = True
1251+
# Verify that the value contains CORE_LOAD_ERROR and CORE_VERSION
1252+
value = t.message[TelemetryField.KEY_VALUE.value]
1253+
assert (
1254+
"CORE_LOAD_ERROR" in value
1255+
), "CORE_LOAD_ERROR not in telemetry value"
1256+
assert "CORE_VERSION" in value, "CORE_VERSION not in telemetry value"
1257+
break
1258+
assert (
1259+
found_minicore_telemetry
1260+
), "Minicore telemetry not found in telemetry records"
1261+
1262+
12331263
@pytest.mark.skipolddriver
12341264
def test_disable_query_context_cache(conn_cnx) -> None:
12351265
with conn_cnx(disable_query_context_cache=True) as conn:

0 commit comments

Comments
 (0)