Adding telemetry to stored procs (#3563)

Author: sfc-gh-daviwang

src/snowflake/snowpark/_internal/telemetry.py

@@ -7,6 +7,7 @@
 import json
 import threading
 from enum import Enum, unique
+from logging import getLogger
 import time
 from typing import Any, Dict, List, Optional
 
@@ -35,6 +36,7 @@
     get_version,
     is_in_stored_procedure,
     is_interactive,
+    generate_random_alphanumeric,
 )
 
 try:
@@ -44,6 +46,8 @@
 except ImportError:
     PS_UTIL_AVAILABLE = False
 
+_logger = getLogger(__name__)
+
 
 @unique
 class TelemetryField(Enum):
@@ -474,13 +478,50 @@ def __init__(self, conn: SnowflakeConnection) -> None:
         self.python_version: str = get_python_version()
         self.os: str = get_os_name()
         self.is_interactive = is_interactive()
+        self._enabled = True
+
+        # Initializing telemetry client for stored procedures
+        # In stored procs, we can't import this package at the top level, so we need to do it here
+        internal_metrics_available = None
+        if is_in_stored_procedure():
+            try:
+                from _snowflake import internal_metrics
+
+                internal_metrics_available = True
+            except ImportError:
+                internal_metrics_available = False
+        self.stored_proc_meter = (
+            internal_metrics.get_meter("snowpark-python-client")
+            if internal_metrics_available
+            else None
+        )
+        # We periodically clean out the stored procedure meter of unused gauges
+        self.clean_up_stored_proc_meter_interval = 1000
 
     def send(self, msg: Dict, timestamp: Optional[int] = None):
+        if not self._enabled:
+            _logger.info("Telemetry client is disabled, skipping telemetry")
+            return
+        if not timestamp:
+            timestamp = get_time_millis()
         if self.telemetry:
-            if not timestamp:
-                timestamp = get_time_millis()
             telemetry_data = PCTelemetryData(message=msg, timestamp=timestamp)
             self.telemetry.try_add_log_to_batch(telemetry_data)
+        elif self.stored_proc_meter is not None:
+            gauge_id = generate_random_alphanumeric(10)
+            self.stored_proc_meter.create_gauge(
+                f"snowflake.snowpark.client.gauge{gauge_id}",
+                description=json.dumps(msg, ensure_ascii=False, separators=(",", ":")),
+                unit="data",
+            ).set(
+                200
+            )  # this is a dummy value
+            if (
+                len(self.stored_proc_meter._instrument_id_instrument)
+                >= self.clean_up_stored_proc_meter_interval
+            ):
+                with self.stored_proc_meter._instrument_id_instrument_lock:
+                    self.stored_proc_meter._instrument_id_instrument.clear()
 
     def _create_basic_telemetry_data(self, telemetry_type: str) -> Dict[str, Any]:
         message = {

src/snowflake/snowpark/session.py

@@ -298,6 +298,7 @@
 _PYTHON_SNOWPARK_GENERATE_MULTILINE_QUERIES = (
     "PYTHON_SNOWPARK_GENERATE_MULTILINE_QUERIES"
 )
+_PYTHON_SNOWPARK_INTERNAL_TELEMETRY_ENABLED = "ENABLE_SNOWPARK_FIRST_PARTY_TELEMETRY"
 
 # AST encoding.
 _PYTHON_SNOWPARK_USE_AST = "PYTHON_SNOWPARK_USE_AST"
@@ -675,6 +676,12 @@ def __init__(
         else:
             self._disable_multiline_queries()
 
+        self._internal_telemetry_enabled: bool = (
+            self._conn._get_client_side_session_parameter(
+                _PYTHON_SNOWPARK_INTERNAL_TELEMETRY_ENABLED, False
+            )
+        )
+
         self._large_query_breakdown_enabled: bool = self.is_feature_enabled_for_version(
             _PYTHON_SNOWPARK_USE_LARGE_QUERY_BREAKDOWN_OPTIMIZATION_VERSION
         )
@@ -4083,17 +4090,21 @@ def telemetry_enabled(self) -> bool:
             >>> session.telemetry_enabled
             True
         """
-        return self._conn._conn.telemetry_enabled
+        return self._conn._telemetry_client._enabled
 
     @telemetry_enabled.setter
     def telemetry_enabled(self, value):
         # Set both in-band and out-of-band telemetry to True/False
         if value:
-            self._conn._conn.telemetry_enabled = True
-            self._conn._telemetry_client.telemetry._enabled = True
+            self._conn._telemetry_client._enabled = True
+            if is_in_stored_procedure() and not self._stored_proc_telemetry_enabled:
+                _logger.debug(
+                    "Client side parameter ENABLE_SNOWPARK_FIRST_PARTY_TELEMETRY is set to False, telemetry could not be enabled"
+                )
+                self._conn._telemetry_client._enabled = False
+
         else:
-            self._conn._conn.telemetry_enabled = False
-            self._conn._telemetry_client.telemetry._enabled = False
+            self._conn._telemetry_client._enabled = False
 
     @property
     def file(self) -> FileOperation:

tests/unit/test_stored_proc_telemetry.py

@@ -0,0 +1,154 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import threading
+from unittest.mock import patch, MagicMock
+from typing import Dict
+import json
+
+from snowflake.snowpark._internal.telemetry import TelemetryClient
+from snowflake.snowpark import Session
+
+
+class MockGauge:
+    def __init__(self, name: str, unit: str = "", description: str = "") -> None:
+        self.name = name
+        self.unit = unit
+        self.description = description
+        self._value = None
+
+    def set(self, value):
+        self._value = value
+        return self
+
+    def get_value(self):
+        return self._value
+
+
+class MockMeter:
+    def __init__(self) -> None:
+        self._instrument_id_instrument: Dict[str, MockGauge] = {}
+        self._instrument_id_instrument_lock = threading.Lock()
+
+    def create_gauge(self, name: str, unit: str = "", description: str = ""):
+        instrument = MockGauge(
+            name,
+            unit,
+            description,
+        )
+        instrument_id = f"{name},{unit},{description}"
+
+        with self._instrument_id_instrument_lock:
+            self._instrument_id_instrument[instrument_id] = instrument
+            return instrument
+
+
+@patch("snowflake.snowpark._internal.telemetry.is_in_stored_procedure")
+def test_telemetry_client_with_mock_meter(mock_is_in_stored_proc):
+    mock_is_in_stored_proc.return_value = True
+    mock_meter = MockMeter()
+    mock_conn = MagicMock()
+    client = TelemetryClient(mock_conn)
+    client._enabled = True
+    client.telemetry = None
+    client.stored_proc_meter = mock_meter
+    client.clean_up_stored_proc_meter_interval = 200
+    test_message = {"test": "data", "func_name": "test_function"}
+    client.send(test_message)
+    assert len(mock_meter._instrument_id_instrument) == 1
+    gauge = list(mock_meter._instrument_id_instrument.values())[0]
+    assert "snowflake.snowpark.client.gauge" in gauge.name
+    assert gauge.unit == "data"
+    assert gauge.description == json.dumps(
+        test_message, ensure_ascii=False, separators=(",", ":")
+    )
+    assert gauge.get_value() == 200
+
+
+@patch("snowflake.snowpark._internal.telemetry.is_in_stored_procedure")
+def test_telemetry_client_multiple_sends(mock_is_in_stored_proc):
+    mock_is_in_stored_proc.return_value = True
+    mock_meter = MockMeter()
+    mock_conn = MagicMock()
+    client = TelemetryClient(mock_conn)
+    client._enabled = True
+    client.telemetry = None
+    client.stored_proc_meter = mock_meter
+    client.clean_up_stored_proc_meter_interval = 200
+    for i in range(100):
+        client.send({"test": f"data_{i}"})
+    assert len(mock_meter._instrument_id_instrument) == 100
+    gauges = list(mock_meter._instrument_id_instrument.values())
+    for gauge in gauges:
+        assert "snowflake.snowpark.client.gauge" in gauge.name
+        assert gauge.get_value() == 200
+
+
+@patch("snowflake.snowpark._internal.telemetry.is_in_stored_procedure")
+def test_telemetry_client_cleanup(mock_is_in_stored_proc):
+    """Test TelemetryClient cleanup mechanism"""
+    mock_is_in_stored_proc.return_value = True
+    mock_meter = MockMeter()
+    mock_conn = MagicMock()
+    client = TelemetryClient(mock_conn)
+    client._enabled = True
+    client.telemetry = None
+    client.stored_proc_meter = mock_meter
+    client.clean_up_stored_proc_meter_interval = 100
+    for i in range(client.clean_up_stored_proc_meter_interval):
+        client.send({"test": f"data_{i}"})
+    assert len(mock_meter._instrument_id_instrument) == 0
+    client.send({"test": "data_after_cleanup"})
+    assert len(mock_meter._instrument_id_instrument) == 1
+
+
+@patch("snowflake.snowpark.session.is_in_stored_procedure")
+@patch("snowflake.snowpark.session._logger")
+def test_internal_telemetry_disabled(mock_logger, mock_is_in_stored_proc):
+    """Test that Session logs debug message when telemetry is enabled but stored proc telemetry is disabled"""
+    mock_is_in_stored_proc.return_value = True
+    mock_session = MagicMock(spec=Session)
+    mock_session._stored_proc_telemetry_enabled = False
+    mock_session._conn = MagicMock()
+    mock_session._conn._conn = MagicMock()
+    mock_session._conn._telemetry_client = MagicMock()
+    Session.telemetry_enabled.fset(mock_session, False)
+    assert mock_session._conn._telemetry_client._enabled is False
+    Session.telemetry_enabled.fset(mock_session, True)
+    mock_logger.debug.assert_called_once_with(
+        "Client side parameter ENABLE_SNOWPARK_FIRST_PARTY_TELEMETRY is set to False, telemetry could not be enabled"
+    )
+
+
+@patch("snowflake.snowpark._internal.telemetry.is_in_stored_procedure")
+def test_telemetry_client_internal_metrics_import_fails(mock_is_in_stored_proc):
+    """Test TelemetryClient handles ImportError when _snowflake.internal_metrics import fails"""
+    mock_is_in_stored_proc.return_value = True
+    mock_conn = MagicMock()
+    with patch(
+        "builtins.__import__", side_effect=ImportError("No module named '_snowflake'")
+    ):
+        client = TelemetryClient(mock_conn)
+        assert client is not None
+        assert client._enabled is True
+        assert client.stored_proc_meter is None
+        test_message = {"test": "data", "func_name": "test_function"}
+        client.send(test_message)
+
+
+@patch("snowflake.snowpark._internal.telemetry.is_in_stored_procedure")
+def test_telemetry_client_disabled(mock_is_in_stored_proc):
+    """Test that no message is sent when telemetry client is not enabled"""
+    mock_is_in_stored_proc.return_value = True
+    mock_meter = MockMeter()
+    mock_conn = MagicMock()
+    client = TelemetryClient(mock_conn)
+    client._enabled = False
+    client.telemetry = None
+    client.stored_proc_meter = mock_meter
+    client.clean_up_stored_proc_meter_interval = 200
+    test_message = {"test": "data", "func_name": "test_function"}
+    client.send(test_message)
+    assert len(mock_meter._instrument_id_instrument) == 0