add snowflake-telemetry-python dep optional

Author: sfc-gh-yuwang

setup.py

@@ -230,6 +230,7 @@ def run(self):
             "opentelemetry-api>=1.0.0, <2.0.0",
             "opentelemetry-sdk>=1.0.0, <2.0.0",
             "opentelemetry-exporter-otlp>=1.0.0, <2.0.0",
+            "snowflake-telemetry-python",
         ],
     },
     classifiers=[

src/snowflake/snowpark/_internal/event_table_telemetry.py

@@ -7,7 +7,6 @@
 from logging import getLogger
 from typing import Dict, Optional, Tuple
 from snowflake.connector.options import MissingOptionalDependency, ModuleLikeObject
-
 import snowflake.snowpark
 import requests
 
@@ -23,6 +22,24 @@ class MissingOpenTelemetry(MissingOptionalDependency):
     _dep_name = "opentelemetry"
 
 
+class MissingSnowflakeTelemetry(MissingOptionalDependency):
+    _dep_name = "snowflake-telemetry-python"
+
+
+def _import_or_missing_snowflake_telemetry() -> Tuple[ModuleLikeObject, bool]:
+    try:
+        # "snowflake-telemetry-python" is the *distribution* name. The importable
+        # Python modules live under the `snowflake` namespace package.
+        snowflake_telemetry = importlib.import_module("snowflake")
+        # Import the parent first so `snowflake.telemetry` is attached to the
+        # returned `snowflake` module (mirrors the pattern used for OpenTelemetry).
+        importlib.import_module("snowflake.telemetry")
+        importlib.import_module("snowflake.telemetry.trace")
+        return snowflake_telemetry, True
+    except ImportError:
+        return MissingSnowflakeTelemetry(), False
+
+
 def _import_or_missing_opentelemetry() -> Tuple[ModuleLikeObject, bool]:
     try:
         opentelemetry = importlib.import_module("opentelemetry")
@@ -39,6 +56,7 @@ def _import_or_missing_opentelemetry() -> Tuple[ModuleLikeObject, bool]:
 
 
 opentelemetry, installed_opentelemetry = _import_or_missing_opentelemetry()
+snowflake_telemetry, _ = _import_or_missing_snowflake_telemetry()
 
 BaseLogProvider = opentelemetry._logs.LoggerProvider if installed_opentelemetry else ABC
 BaseTraceProvider = (
@@ -394,7 +412,8 @@ def _init_trace_level(
         opentelemetry.trace.set_tracer_provider(self._proxy_tracer_provider)
 
         self._tracer_provider = opentelemetry.sdk.trace.TracerProvider(
-            resource=resource
+            resource=resource,
+            id_generator=snowflake_telemetry.telemetry.trace.SnowflakeTraceIdGenerator(),
         )
 
         trace_session = requests.Session()

tests/unit/test_external_telemetry.py

@@ -101,3 +101,68 @@ def always_fail(_name):
 
     assert installed is False
     assert isinstance(otel, ett.MissingOpenTelemetry)
+
+
+def test_import_or_missing_snowflake_telemetry(monkeypatch):
+    # Build a fake snowflake module tree
+    snowflake = types.ModuleType("snowflake")
+    snowflake_telemetry = types.ModuleType("snowflake.telemetry")
+    snowflake_trace = types.ModuleType("snowflake.telemetry.trace")
+
+    class SnowflakeTraceIdGenerator:
+        pass
+
+    snowflake_trace.SnowflakeTraceIdGenerator = SnowflakeTraceIdGenerator
+
+    def fake_import_module_factory(registry):
+        def fake_import_module(name):
+            # Simulate ImportError if an expected dependency isn't present
+            if name not in registry:
+                raise ImportError(name)
+
+            mod = registry[name]
+
+            # Emulate Python import behavior: after importing X.Y, X gets attribute Y.
+            if "." in name:
+                parent_name, child = name.rsplit(".", 1)
+                setattr(registry[parent_name], child, mod)
+
+            return mod
+
+        return fake_import_module
+
+    # Case 1: snowflake telemetry trace is missing
+    missing_registry = {
+        "snowflake": snowflake,
+        "snowflake.telemetry": snowflake_telemetry,
+    }
+    monkeypatch.setattr(
+        ett.importlib, "import_module", fake_import_module_factory(missing_registry)
+    )
+
+    sf, installed = ett._import_or_missing_snowflake_telemetry()
+    assert installed is False
+    assert isinstance(sf, ett.MissingSnowflakeTelemetry)
+
+    # Case 2: snowflake telemetry trace is present
+    present_registry = {
+        "snowflake": snowflake,
+        "snowflake.telemetry": snowflake_telemetry,
+        "snowflake.telemetry.trace": snowflake_trace,
+    }
+    monkeypatch.setattr(
+        ett.importlib, "import_module", fake_import_module_factory(present_registry)
+    )
+
+    sf, installed = ett._import_or_missing_snowflake_telemetry()
+    assert installed is True
+    assert sf is snowflake
+
+    # Key regression check: no AttributeError when accessing sf.telemetry.trace
+    assert hasattr(sf, "telemetry")
+    assert hasattr(sf.telemetry, "trace")
+    assert sf.telemetry.trace is snowflake_trace
+    assert sf.telemetry.trace.SnowflakeTraceIdGenerator is SnowflakeTraceIdGenerator
+    assert isinstance(
+        sf.telemetry.trace.SnowflakeTraceIdGenerator(), SnowflakeTraceIdGenerator
+    )