snowpark-connect 0.21.0__py3-none-any.whl → 0.23.0__py3-none-any.whl

snowflake/snowpark_connect/utils/telemetry.py

@@ -1,19 +1,22 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-
+import functools
 import json
 import os
 import queue
 import threading
+import uuid
 from abc import ABC, abstractmethod
 from collections import defaultdict
+from collections.abc import Iterable
 from contextvars import ContextVar
 from enum import Enum, unique
 from typing import Dict
 
 import google.protobuf.message
 
+from snowflake.connector.cursor import SnowflakeCursor
 from snowflake.connector.telemetry import (
     TelemetryClient as PCTelemetryClient,
     TelemetryData as PCTelemetryData,
@@ -21,7 +24,6 @@ from snowflake.connector.telemetry import (
 )
 from snowflake.connector.time_util import get_time_millis
 from snowflake.snowpark import Session
-from snowflake.snowpark._internal.telemetry import safe_telemetry
 from snowflake.snowpark._internal.utils import get_os_name, get_python_version
 from snowflake.snowpark.version import VERSION as snowpark_version
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
@@ -43,6 +45,7 @@ class TelemetryField(Enum):
     KEY_OS = "operating_system"
     KEY_DATA = "data"
     KEY_START_TIME = "start_time"
+    KEY_EVENT_ID = "event_id"
 
 
 class TelemetryType(Enum):
@@ -107,7 +110,34 @@ REDACTED_PLAN_SUFFIXES = [
 ]
 
 
+def _basic_telemetry_data() -> Dict:
+    return {
+        **STATIC_TELEMETRY_DATA,
+        TelemetryField.KEY_EVENT_ID.value: str(uuid.uuid4()),
+    }
+
+
+def safe(func):
+    """
+    Decorator to safely execute telemetry functions, catching and logging exceptions
+    without affecting the main application flow.
+    """
+
+    @functools.wraps(func)
+    def wrap(*args, **kwargs):
+        try:
+            func(*args, **kwargs)
+        except Exception:
+            # We don't really care if telemetry fails, just want to be safe for the user
+            logger.warning(f"Telemetry operation failed: {func}", exc_info=True)
+
+    return wrap
+
+
 class TelemetrySink(ABC):
+    MAX_BUFFER_ELEMENTS = 20
+    MAX_WAIT_MS = 10000  # 10 seconds
+
     @abstractmethod
     def add_telemetry_data(self, message: dict, timestamp: int) -> None:
         pass
@@ -128,23 +158,44 @@ class NoOpTelemetrySink(TelemetrySink):
 class ClientTelemetrySink(TelemetrySink):
     def __init__(self, telemetry_client: PCTelemetryClient) -> None:
         self._telemetry_client = telemetry_client
+        self._lock = threading.Lock()
+        self._reset()
 
     def add_telemetry_data(self, message: dict, timestamp: int) -> None:
         telemetry_data = PCTelemetryData(message=message, timestamp=timestamp)
         self._telemetry_client.try_add_log_to_batch(telemetry_data)
+        with self._lock:
+            self._events_since_last_flush += 1
+        # flush more often than the underlying telemetry client
+        if self._should_flush():
+            self.flush()
 
     def flush(self) -> None:
+        with self._lock:
+            self._reset()
         self._telemetry_client.send_batch()
 
+    def _should_flush(self) -> bool:
+        current_time = get_time_millis()
+
+        return (
+            self._events_since_last_flush >= TelemetrySink.MAX_BUFFER_ELEMENTS
+            or (current_time - self._last_flush_time) >= TelemetrySink.MAX_WAIT_MS
+        )
+
+    def _reset(self):
+        self._events_since_last_flush = 0
+        self._last_flush_time = get_time_millis()
+
 
 class QueryTelemetrySink(TelemetrySink):
 
-    MAX_BUFFER_SIZE = 100 * 1024  # 100KB
-    MAX_WAIT_MS = 10000  # 10 seconds
+    MAX_BUFFER_SIZE = 20 * 1024  # 20KB
     TELEMETRY_JOB_ID = "43e72d9b-56d0-4cdb-a615-6b5b5059d6df"
 
     def __init__(self, session: Session) -> None:
         self._session = session
+        self._lock = threading.Lock()
         self._reset()
 
     def add_telemetry_data(self, message: dict, timestamp: int) -> None:
@@ -152,31 +203,37 @@ class QueryTelemetrySink(TelemetrySink):
 
         # stringify entry, and escape single quotes
         entry_str = json.dumps(telemetry_entry).replace("'", "''")
-        self._buffer.append(entry_str)
-        self._buffer_size += len(entry_str)
 
-        current_time = get_time_millis()
-        if (
-            self._buffer_size > QueryTelemetrySink.MAX_BUFFER_SIZE
-            or (current_time - self._last_export_time) > QueryTelemetrySink.MAX_WAIT_MS
-        ):
+        with self._lock:
+            self._buffer.append(entry_str)
+            self._buffer_size += len(entry_str)
+
+        if self._should_flush():
             self.flush()
 
     def flush(self) -> None:
-        if not self._buffer:
-            return
+        with self._lock:
+            if not self._buffer:
+                return
+            # prefix query with a unique identifier for easier tracking
+            query = f"select '{self.TELEMETRY_JOB_ID}' as scos_telemetry_export, '[{','.join(self._buffer)}]'"
+            self._reset()
 
-        # prefix query with a unique identifier for easier tracking
-        query = f"select '{self.TELEMETRY_JOB_ID}' as scos_telemetry_export, '[{','.join(self._buffer)}]'"
         self._session.sql(query).collect_nowait()
 
-        self._reset()
-
     def _reset(self) -> None:
         self._buffer = []
         self._buffer_size = 0
         self._last_export_time = get_time_millis()
 
+    def _should_flush(self):
+        current_time = get_time_millis()
+        return (
+            self._buffer_size >= QueryTelemetrySink.MAX_BUFFER_SIZE
+            or len(self._buffer) >= TelemetrySink.MAX_BUFFER_ELEMENTS
+            or (current_time - self._last_export_time) >= TelemetrySink.MAX_WAIT_MS
+        )
+
 
 class Telemetry:
     def __init__(self, is_enabled=True) -> None:
@@ -185,6 +242,8 @@ class Telemetry:
             "request_summary", default={}
         )
         self._is_enabled = is_enabled
+        self._is_initialized = False
+        self._lock = threading.Lock()
 
         # Async processing setup
         self._message_queue = queue.Queue(maxsize=10000)
@@ -202,6 +261,12 @@ class Telemetry:
         if not self._is_enabled:
             return
 
+        with self._lock:
+            if self._is_initialized:
+                logger.warning("Telemetry is already initialized")
+                return
+            self._is_initialized = True
+
         telemetry = getattr(session._conn._conn, "_telemetry", None)
         if telemetry is None:
             # no telemetry client available, so we export with queries
@@ -210,8 +275,9 @@ class Telemetry:
             self._sink = ClientTelemetrySink(telemetry)
 
         self._start_worker_thread()
+        logger.info(f"Telemetry initialized with {type(self._sink)}")
 
-    @safe_telemetry
+    @safe
     def initialize_request_summary(
         self, request: google.protobuf.message.Message
     ) -> None:
@@ -234,8 +300,29 @@ class Telemetry:
                 request.plan, REDACTED_PLAN_SUFFIXES
             )
 
-    @safe_telemetry
+    def _not_in_request(self):
+        # we don't want to add things to the summary if it's not initialized
+        return "created_on" not in self._request_summary.get()
+
+    @safe
+    def report_parsed_sql_plan(self, plan: google.protobuf.message.Message) -> None:
+        if self._not_in_request():
+            return
+
+        summary = self._request_summary.get()
+
+        if "parsed_sql_plans" not in summary:
+            summary["parsed_sql_plans"] = []
+
+        summary["parsed_sql_plans"].append(
+            _protobuf_to_json_with_redaction(plan, REDACTED_PLAN_SUFFIXES)
+        )
+
+    @safe
     def report_function_usage(self, function_name: str) -> None:
+        if self._not_in_request():
+            return
+
         summary = self._request_summary.get()
 
         if "used_functions" not in summary:
@@ -243,8 +330,11 @@ class Telemetry:
 
         summary["used_functions"][function_name] += 1
 
-    @safe_telemetry
+    @safe
     def report_request_failure(self, e: Exception) -> None:
+        if self._not_in_request():
+            return
+
         summary = self._request_summary.get()
 
         summary["was_successful"] = False
@@ -255,37 +345,78 @@ class Telemetry:
         if error_location:
             summary["error_location"] = error_location
 
-    @safe_telemetry
-    def report_config_set(self, key, value):
+    @safe
+    def report_config_set(self, pairs: Iterable) -> None:
+        if self._not_in_request():
+            return
+
         summary = self._request_summary.get()
 
         if "config_set" not in summary:
             summary["config_set"] = []
 
-        summary["config_set"].append(
-            {
-                "key": key,
-                "value": value if key in RECORDED_CONFIG_KEYS else "<redacted>",
-            }
-        )
+        for p in pairs:
+            summary["config_set"].append(
+                {
+                    "key": p.key,
+                    "value": p.value if p.key in RECORDED_CONFIG_KEYS else "<redacted>",
+                }
+            )
+
+    @safe
+    def report_config_unset(self, keys: Iterable[str]) -> None:
+        if self._not_in_request():
+            return
 
-    @safe_telemetry
-    def report_config_unset(self, key):
         summary = self._request_summary.get()
 
         if "config_unset" not in summary:
             summary["config_unset"] = []
 
-        summary["config_unset"].append(key)
+        summary["config_unset"].extend(keys)
+
+    @safe
+    def report_config_get(self, keys: Iterable[str]) -> None:
+        if self._not_in_request():
+            return
 
-    @safe_telemetry
-    def report_config_op_type(self, op_type: str):
         summary = self._request_summary.get()
 
+        if "config_get" not in summary:
+            summary["config_get"] = []
+
+        summary["config_get"].extend(keys)
+
+    @safe
+    def report_config_op_type(self, op_type: str):
+        if self._not_in_request():
+            return
+
+        summary = self._request_summary.get()
         summary["config_op_type"] = op_type
 
-    @safe_telemetry
-    def report_query_id(self, query_id: str):
+    @safe
+    def report_query(
+        self, result: SnowflakeCursor | dict | Exception, **kwargs
+    ) -> None:
+        if result is None or isinstance(result, dict) or self._not_in_request():
+            return
+
+        # SnowflakeCursor and SQL errors will have sfqid
+        # other exceptions will not have it
+        # TODO: handle async queries, but filter out telemetry export queries
+        qid = getattr(result, "sfqid", None)
+
+        if qid is None:
+            logger.warning("Missing query id in result: %s", result)
+
+        is_internal = kwargs.get("_is_internal", False)
+        if is_internal:
+            self._report_internal_query()
+        elif qid:
+            self._report_query_id(qid)
+
+    def _report_query_id(self, query_id: str):
         summary = self._request_summary.get()
 
         if "queries" not in summary:
@@ -293,13 +424,19 @@ class Telemetry:
 
         summary["queries"].append(query_id)
 
-    @safe_telemetry
-    def report_internal_query(self):
+    def _report_internal_query(self):
         summary = self._request_summary.get()
+
+        if "internal_queries" not in summary:
+            summary["internal_queries"] = 0
+
         summary["internal_queries"] += 1
 
-    @safe_telemetry
+    @safe
     def report_udf_usage(self, udf_name: str):
+        if self._not_in_request():
+            return
+
         summary = self._request_summary.get()
 
         if "udf_usage" not in summary:
@@ -307,8 +444,10 @@ class Telemetry:
 
         summary["udf_usage"][udf_name] += 1
 
-    @safe_telemetry
-    def report_io(self, op: str, type: str, options: dict | None):
+    def _report_io(self, op: str, type: str, options: dict | None):
+        if self._not_in_request():
+            return
+
         summary = self._request_summary.get()
 
         if "io" not in summary:
@@ -321,16 +460,18 @@ class Telemetry:
 
         summary["io"].append(io)
 
+    @safe
     def report_io_read(self, type: str, options: dict | None):
-        self.report_io("read", type, options)
+        self._report_io("read", type, options)
 
+    @safe
     def report_io_write(self, type: str, options: dict | None):
-        self.report_io("write", type, options)
+        self._report_io("write", type, options)
 
-    @safe_telemetry
+    @safe
     def send_server_started_telemetry(self):
         message = {
-            **STATIC_TELEMETRY_DATA,
+            **_basic_telemetry_data(),
             TelemetryField.KEY_TYPE.value: TelemetryType.TYPE_EVENT.value,
             TelemetryType.EVENT_TYPE.value: EventType.SERVER_STARTED.value,
             TelemetryField.KEY_DATA.value: {
@@ -339,17 +480,22 @@ class Telemetry:
         }
         self._send(message)
 
-    @safe_telemetry
+    @safe
     def send_request_summary_telemetry(self):
+        if self._not_in_request():
+            logger.warning(
+                "Truing to send request summary telemetry without initializing it"
+            )
+            return
+
         summary = self._request_summary.get()
         message = {
-            **STATIC_TELEMETRY_DATA,
+            **_basic_telemetry_data(),
             TelemetryField.KEY_TYPE.value: TelemetryType.TYPE_REQUEST_SUMMARY.value,
             TelemetryField.KEY_DATA.value: summary,
         }
         self._send(message)
 
-    @safe_telemetry
     def _send(self, msg: Dict) -> None:
         """Queue a telemetry message for asynchronous processing."""
         if not self._is_enabled:
@@ -385,19 +531,6 @@ class Telemetry:
             finally:
                 self._message_queue.task_done()
 
-        # Process any remaining messages
-        while not self._message_queue.empty():
-            try:
-                message, timestamp = self._message_queue.get_nowait()
-                self._sink.add_telemetry_data(message, timestamp)
-                self._message_queue.task_done()
-            except Exception:
-                logger.warning(
-                    "Failed to add remaining telemetry messages to sink during shutdown",
-                    exc_info=True,
-                )
-                break
-
         # Flush the sink
         self._sink.flush()
 
@@ -439,6 +572,18 @@ def _error_location(e: Exception) -> Dict | None:
     }
 
 
+def _is_map_field(field_descriptor) -> bool:
+    """
+    Check if a protobuf field is a map.
+    """
+    return (
+        field_descriptor.label == field_descriptor.LABEL_REPEATED
+        and field_descriptor.message_type is not None
+        and field_descriptor.message_type.has_options
+        and field_descriptor.message_type.GetOptions().map_entry
+    )
+
+
 def _protobuf_to_json_with_redaction(
     message: google.protobuf.message.Message, redacted_suffixes: list[str]
 ) -> dict:
@@ -463,7 +608,9 @@ def _protobuf_to_json_with_redaction(
             return "<redacted>"
 
         # Handle different field types
-        if field_descriptor.type == field_descriptor.TYPE_MESSAGE:
+        if _is_map_field(field_descriptor):
+            return dict(value)
+        elif field_descriptor.type == field_descriptor.TYPE_MESSAGE:
             if field_descriptor.label == field_descriptor.LABEL_REPEATED:
                 # Repeated message field
                 return [_protobuf_to_json_recursive(item, field_path) for item in value]
@@ -481,6 +628,11 @@ def _protobuf_to_json_with_redaction(
         msg: google.protobuf.message.Message, current_path: str = ""
     ) -> dict:
         """Recursively convert protobuf message to dict"""
+
+        if not isinstance(msg, google.protobuf.message.Message):
+            logger.warning("Expected a protobuf message, got: %s", type(msg))
+            return {}
+
         result = {}
 
         # Use ListFields() to get all set fields

snowflake/snowpark_connect/utils/udxf_import_utils.py

@@ -0,0 +1,14 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+from snowflake import snowpark
+from snowflake.snowpark_connect.config import global_config
+
+
+def get_python_udxf_import_files(session: snowpark.Session) -> str:
+    config_imports = global_config.get("snowpark.connect.udf.imports", "")
+    config_imports = config_imports.strip("[] ").split(",") if config_imports else []
+    imports = {*session._python_files, *session._import_files, *config_imports}
+
+    return ",".join([file for file in imports if file])

snowflake/snowpark_connect/version.py

@@ -3,4 +3,4 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
-VERSION = (0, 21, 0)
+VERSION = (0, 23, 0)

snowflake/snowpark_decoder/_internal/proto/generated/DataframeProcessorMsg_pb2.py

@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: DataframeProcessorMsg.proto
+"""Generated protocol buffer code."""
+from google.protobuf.internal import builder as _builder
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from google.protobuf import descriptor_pb2 as google_dot_protobuf_dot_descriptor__pb2
+from google.protobuf import any_pb2 as google_dot_protobuf_dot_any__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1b\x44\x61taframeProcessorMsg.proto\x12\x13\x64\x61taframe_processor\x1a google/protobuf/descriptor.proto\x1a\x19google/protobuf/any.proto\"\x80\x01\n\x07Request\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12:\n\x0e\x64\x61taframe_type\x18\x02 \x01(\x0e\x32\".dataframe_processor.DataframeType\x12%\n\x07payload\x18\x03 \x01(\x0b\x32\x14.google.protobuf.Any\"\xd4\x02\n\x08Response\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12:\n\x0e\x64\x61taframe_type\x18\x02 \x01(\x0e\x32\".dataframe_processor.DataframeType\x12\x36\n\x04\x63ode\x18\x03 \x01(\x0e\x32(.dataframe_processor.Response.StatusCode\x12\x19\n\x0fresult_job_uuid\x18\x04 \x01(\tH\x00\x12\'\n\x07payload\x18\x05 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00\x12\x15\n\rerror_message\x18\x06 \x01(\t\"[\n\nStatusCode\x12\x06\n\x02OK\x10\x00\x12\x0e\n\nINIT_ERROR\x10\x01\x12\x13\n\x0f\x45XECUTION_ERROR\x10\x02\x12 \n\x1cSESSION_RESET_REQUIRED_ERROR\x10\x03\x42\x08\n\x06result\"\x18\n\x16TruncatedSparkRelation*C\n\rDataframeType\x12\r\n\tUNDEFINED\x10\x00\x12\x10\n\x0cSNOWPARK_API\x10\x01\x12\x11\n\rSPARK_CONNECT\x10\x02\x42$\n\x17\x63om.snowflake.dataframeB\x06Protos\x80\x01\x00\x62\x06proto3')
+
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'DataframeProcessorMsg_pb2', globals())
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+  DESCRIPTOR._options = None
+  DESCRIPTOR._serialized_options = b'\n\027com.snowflake.dataframeB\006Protos\200\001\000'
+  _DATAFRAMETYPE._serialized_start=613
+  _DATAFRAMETYPE._serialized_end=680
+  _REQUEST._serialized_start=114
+  _REQUEST._serialized_end=242
+  _RESPONSE._serialized_start=245
+  _RESPONSE._serialized_end=585
+  _RESPONSE_STATUSCODE._serialized_start=484
+  _RESPONSE_STATUSCODE._serialized_end=575
+  _TRUNCATEDSPARKRELATION._serialized_start=587
+  _TRUNCATEDSPARKRELATION._serialized_end=611
+# @@protoc_insertion_point(module_scope)