snowpark-connect 0.27.0__py3-none-any.whl → 1.7.0__py3-none-any.whl

snowflake/snowpark_connect/utils/telemetry.py

@@ -11,10 +11,11 @@ from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import Iterable
 from contextvars import ContextVar
+from dataclasses import dataclass
 from enum import Enum, unique
-from typing import Dict
 
 import google.protobuf.message
+import pyspark.sql.connect.proto.base_pb2 as proto_base
 
 from snowflake.connector.cursor import SnowflakeCursor
 from snowflake.connector.telemetry import (
@@ -26,6 +27,7 @@ from snowflake.connector.time_util import get_time_millis
 from snowflake.snowpark import Session
 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.error.error_codes import ErrorCodes
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.version import VERSION as sas_version
 
@@ -56,22 +58,16 @@ class TelemetryType(Enum):
 
 class EventType(Enum):
     SERVER_STARTED = "scos_server_started"
+    WARNING = "scos_warning"
 
 
 # global labels
-SOURCE = "SparkConnectForSnowpark"
+DEFAULT_SOURCE = "SparkConnectForSnowpark"
 SCOS_VERSION = ".".join([str(d) for d in sas_version if d is not None])
 SNOWPARK_VERSION = ".".join([str(d) for d in snowpark_version if d is not None])
 PYTHON_VERSION = get_python_version()
 OS = get_os_name()
 
-STATIC_TELEMETRY_DATA = {
-    TelemetryField.KEY_SOURCE.value: SOURCE,
-    TelemetryField.KEY_VERSION.value: SCOS_VERSION,
-    TelemetryField.KEY_SNOWPARK_VERSION.value: SNOWPARK_VERSION,
-    TelemetryField.KEY_PYTHON_VERSION.value: PYTHON_VERSION,
-    TelemetryField.KEY_OS.value: OS,
-}
 
 # list of config keys for which we record values, other config values are not recorded
 RECORDED_CONFIG_KEYS = {
@@ -88,6 +84,7 @@ RECORDED_CONFIG_KEYS = {
     "spark.sql.session.localRelationCacheThreshold",
     "spark.sql.mapKeyDedupPolicy",
     "snowpark.connect.sql.passthrough",
+    "snowpark.connect.cte.optimization_enabled",
     "snowpark.connect.iceberg.external_volume",
     "snowpark.connect.sql.identifiers.auto-uppercase",
     "snowpark.connect.udtf.compatibility_mode",
@@ -104,11 +101,13 @@ REDACTED_PLAN_SUFFIXES = [
 ]
 
 
-def _basic_telemetry_data() -> Dict:
-    return {
-        **STATIC_TELEMETRY_DATA,
-        TelemetryField.KEY_EVENT_ID.value: str(uuid.uuid4()),
-    }
+@dataclass
+class TelemetryMessage:
+    """Container for telemetry messages in the processing queue."""
+
+    message: dict
+    timestamp: int
+    is_warning: bool
 
 
 def safe(func):
@@ -121,9 +120,11 @@ def safe(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)
+        except Exception as e:
+            # report failed operation to telemetry
+            telemetry.send_warning_msg(
+                f"Telemetry operation {func} failed due to exception", e
+            )
 
     return wrap
 
@@ -238,6 +239,7 @@ class Telemetry:
         self._is_enabled = is_enabled
         self._is_initialized = False
         self._lock = threading.Lock()
+        self._source = DEFAULT_SOURCE
 
         # Async processing setup
         self._message_queue = queue.Queue(maxsize=10000)
@@ -246,11 +248,32 @@ class Telemetry:
     def __del__(self):
         self.shutdown()
 
-    def initialize(self, session: Session):
+    def _get_static_telemetry_data(self) -> dict:
+        """Get static telemetry data with current configuration."""
+        return {
+            TelemetryField.KEY_SOURCE.value: self._source,
+            TelemetryField.KEY_VERSION.value: SCOS_VERSION,
+            TelemetryField.KEY_SNOWPARK_VERSION.value: SNOWPARK_VERSION,
+            TelemetryField.KEY_PYTHON_VERSION.value: PYTHON_VERSION,
+            TelemetryField.KEY_OS.value: OS,
+        }
+
+    def _basic_telemetry_data(self) -> dict:
+        return {
+            **self._get_static_telemetry_data(),
+            TelemetryField.KEY_EVENT_ID.value: str(uuid.uuid4()),
+        }
+
+    def initialize(self, session: Session, source: str = None):
         """
         Must be called after the session is created to initialize telemetry.
         Gets the telemetry client from the session's connection and uses it
         to report telemetry data.
+
+        Args:
+            session: Snowpark Session to use for telemetry
+            source: Optional source identifier for telemetry (e.g., "SparkConnectThinClient").
+                    Defaults to "SparkConnectForSnowpark".
         """
         if not self._is_enabled:
             return
@@ -261,12 +284,15 @@ class Telemetry:
                 return
             self._is_initialized = True
 
-        telemetry = getattr(session._conn._conn, "_telemetry", None)
-        if telemetry is None:
+            if source is not None:
+                self._source = source
+
+        telemetry_client = getattr(session._conn._conn, "_telemetry", None)
+        if telemetry_client is None:
             # no telemetry client available, so we export with queries
             self._sink = QueryTelemetrySink(session)
         else:
-            self._sink = ClientTelemetrySink(telemetry)
+            self._sink = ClientTelemetrySink(telemetry_client)
 
         self._start_worker_thread()
         logger.info(f"Telemetry initialized with {type(self._sink)}")
@@ -289,10 +315,7 @@ class Telemetry:
 
         self._request_summary.set(summary)
 
-        if hasattr(request, "plan"):
-            summary["query_plan"] = _protobuf_to_json_with_redaction(
-                request.plan, REDACTED_PLAN_SUFFIXES
-            )
+        _set_query_plan(request, summary)
 
     def _not_in_request(self):
         # we don't want to add things to the summary if it's not initialized
@@ -335,6 +358,11 @@ class Telemetry:
         summary["error_message"] = str(e)
         summary["error_type"] = type(e).__name__
 
+        if not hasattr(e, "custom_error_code") or (e.custom_error_code is None):
+            summary["error_code"] = ErrorCodes.INTERNAL_ERROR
+        else:
+            summary["error_code"] = e.custom_error_code
+
         error_location = _error_location(e)
         if error_location:
             summary["error_location"] = error_location
@@ -426,6 +454,63 @@ class Telemetry:
 
         summary["internal_queries"] += 1
 
+    @safe
+    def report_describe_query_cache_lookup(self):
+        """Report a describe query cache lookup."""
+        if self._not_in_request():
+            return
+
+        summary = self._request_summary.get()
+
+        if "describe_cache_lookups" not in summary:
+            summary["describe_cache_lookups"] = 0
+
+        summary["describe_cache_lookups"] += 1
+
+    @safe
+    def report_describe_query_cache_hit(self):
+        """Report a describe query cache hit."""
+        if self._not_in_request():
+            return
+
+        summary = self._request_summary.get()
+
+        if "describe_cache_hits" not in summary:
+            summary["describe_cache_hits"] = 0
+
+        summary["describe_cache_hits"] += 1
+
+    @safe
+    def report_describe_query_cache_expired(self, expired_by: float):
+        """Report a describe query cache hit."""
+        if self._not_in_request():
+            return
+
+        summary = self._request_summary.get()
+
+        if "describe_cache_expired" not in summary:
+            summary["describe_cache_expired"] = 0
+
+        summary["describe_cache_expired"] += 1
+
+        if "describe_cache_expired_by" not in summary:
+            summary["describe_cache_expired_by"] = []
+
+        summary["describe_cache_expired_by"].append(expired_by)
+
+    @safe
+    def report_describe_query_cache_clear(self):
+        """Report a describe query cache clear."""
+        if self._not_in_request():
+            return
+
+        summary = self._request_summary.get()
+
+        if "describe_cache_cleared" not in summary:
+            summary["describe_cache_cleared"] = 0
+
+        summary["describe_cache_cleared"] += 1
+
     @safe
     def report_udf_usage(self, udf_name: str):
         if self._not_in_request():
@@ -460,7 +545,7 @@ class Telemetry:
     @safe
     def send_server_started_telemetry(self):
         message = {
-            **_basic_telemetry_data(),
+            **self._basic_telemetry_data(),
             TelemetryField.KEY_TYPE.value: TelemetryType.TYPE_EVENT.value,
             TelemetryType.EVENT_TYPE.value: EventType.SERVER_STARTED.value,
             TelemetryField.KEY_DATA.value: {
@@ -472,27 +557,69 @@ class Telemetry:
     @safe
     def send_request_summary_telemetry(self):
         if self._not_in_request():
-            logger.warning(
-                "Truing to send request summary telemetry without initializing it"
+            self.send_warning_msg(
+                "Trying to send request summary telemetry without initializing it"
             )
             return
 
         summary = self._request_summary.get()
         message = {
-            **_basic_telemetry_data(),
+            **self._basic_telemetry_data(),
             TelemetryField.KEY_TYPE.value: TelemetryType.TYPE_REQUEST_SUMMARY.value,
             TelemetryField.KEY_DATA.value: summary,
         }
         self._send(message)
 
-    def _send(self, msg: Dict) -> None:
+    def send_warning_msg(self, msg: str, e: Exception = None) -> None:
+        # using this within @safe decorator may result in recursive loop
+        try:
+            message = self._build_warning_message(msg, e)
+            if not message:
+                return
+
+            self._send(message, is_warning=True)
+        except Exception:
+            # if there's an exception here, there's nothing we can really do about it
+            pass
+
+    def _build_warning_message(self, warning_msg: str, e: Exception = None) -> dict:
+        try:
+            data = {"warning_message": warning_msg}
+            if e is not None:
+                data["exception"] = repr(e)
+
+            # add session and operation id if available
+            spark_session_id = self._request_summary.get().get("spark_session_id", None)
+            if spark_session_id is not None:
+                data["spark_session_id"] = spark_session_id
+
+            spark_operation_id = self._request_summary.get().get(
+                "spark_operation_id", None
+            )
+            if spark_operation_id is not None:
+                data["spark_operation_id"] = spark_operation_id
+
+            message = {
+                **self._basic_telemetry_data(),
+                TelemetryField.KEY_TYPE.value: TelemetryType.TYPE_EVENT.value,
+                TelemetryType.EVENT_TYPE.value: EventType.WARNING.value,
+                TelemetryField.KEY_DATA.value: data,
+            }
+            return message
+        except Exception:
+            return {}
+
+    def _send(self, msg: dict, is_warning: bool = False) -> None:
         """Queue a telemetry message for asynchronous processing."""
         if not self._is_enabled:
             return
 
         timestamp = get_time_millis()
         try:
-            self._message_queue.put_nowait((msg, timestamp))
+            telemetry_msg = TelemetryMessage(
+                message=msg, timestamp=timestamp, is_warning=is_warning
+            )
+            self._message_queue.put_nowait(telemetry_msg)
         except queue.Full:
             # If queue is full, drop the message to avoid blocking
             logger.warning("Telemetry queue is full, dropping message")
@@ -510,13 +637,16 @@ class Telemetry:
         while True:
             try:
                 # block to allow the GIL to switch threads
-                message, timestamp = self._message_queue.get()
-                if timestamp is None and message is None:
-                    # shutdown
+                telemetry_msg = self._message_queue.get()
+                if telemetry_msg is None:
+                    # shutdown signal
                     break
-                self._sink.add_telemetry_data(message, timestamp)
-            except Exception:
-                logger.warning("Failed to add telemetry message to sink", exc_info=True)
+                self._sink.add_telemetry_data(
+                    telemetry_msg.message, telemetry_msg.timestamp
+                )
+            except Exception as e:
+                if not telemetry_msg.is_warning:
+                    self.send_warning_msg("Failed to add telemetry message to sink", e)
             finally:
                 self._message_queue.task_done()
 
@@ -529,7 +659,7 @@ class Telemetry:
             return
 
         try:
-            self._message_queue.put_nowait((None, None))
+            self._message_queue.put_nowait(None)
             # Wait for worker thread to finish
             self._worker_thread.join(timeout=3.0)
         except Exception:
@@ -538,7 +668,7 @@ class Telemetry:
             )
 
 
-def _error_location(e: Exception) -> Dict | None:
+def _error_location(e: Exception) -> dict | None:
     """
     Inspect the exception traceback and extract the file name, line number, and function name
     from the last frame (the one that raised the exception).
@@ -619,7 +749,7 @@ def _protobuf_to_json_with_redaction(
         """Recursively convert protobuf message to dict"""
 
         if not isinstance(msg, google.protobuf.message.Message):
-            logger.warning("Expected a protobuf message, got: %s", type(msg))
+            telemetry.send_warning_msg(f"Expected a protobuf message, got: {type(msg)}")
             return {}
 
         result = {}
@@ -644,6 +774,28 @@ def _protobuf_to_json_with_redaction(
     )
 
 
+def _set_query_plan(request: google.protobuf.message.Message, summary: dict) -> None:
+    if isinstance(request, proto_base.ExecutePlanRequest):
+        # ExecutePlanRequest has plan at top level
+        if hasattr(request, "plan"):
+            summary["query_plan"] = (
+                _protobuf_to_json_with_redaction(request.plan, REDACTED_PLAN_SUFFIXES),
+            )
+
+    elif isinstance(request, proto_base.AnalyzePlanRequest):
+        # AnalyzePlanRequest has plan under oneof analyze
+        analyze_type = request.WhichOneof("analyze")
+        if not analyze_type:
+            return
+
+        summary["analyze_type"] = analyze_type
+        analyze_field = getattr(request, analyze_type)
+        if hasattr(analyze_field, "plan"):
+            summary["query_plan"] = _protobuf_to_json_with_redaction(
+                analyze_field.plan, REDACTED_PLAN_SUFFIXES
+            )
+
+
 # global telemetry client
 telemetry = Telemetry(is_enabled="SNOWPARK_CONNECT_DISABLE_TELEMETRY" not in os.environ)
 

snowflake/snowpark_connect/utils/temporary_view_cache.py

@@ -0,0 +1,67 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+from typing import Optional, Tuple
+
+from pyspark.errors import AnalysisException
+
+from snowflake.snowpark_connect.dataframe_container import DataFrameContainer
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
+from snowflake.snowpark_connect.utils.concurrent import SynchronizedDict
+from snowflake.snowpark_connect.utils.context import get_spark_session_id
+
+_temp_views = SynchronizedDict[Tuple[str, str], DataFrameContainer]()
+
+
+def register_temp_view(name: str, df: DataFrameContainer, replace: bool) -> None:
+    normalized_name = _normalize(name)
+    current_session_id = get_spark_session_id()
+    for key in list(_temp_views.keys()):
+        if _normalize(key[0]) == normalized_name and key[1] == current_session_id:
+            if replace:
+                _temp_views.remove(key)
+                break
+            else:
+                exception = AnalysisException(
+                    f"[TEMP_TABLE_OR_VIEW_ALREADY_EXISTS] Cannot create the temporary view `{name}` because it already exists."
+                )
+                attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+                raise exception
+
+    _temp_views[(name, current_session_id)] = df
+
+
+def unregister_temp_view(name: str) -> bool:
+    normalized_name = _normalize(name)
+
+    for key in _temp_views.keys():
+        normalized_key = _normalize(key[0])
+        if normalized_name == normalized_key and key[1] == get_spark_session_id():
+            pop_result = _temp_views.remove(key)
+            return pop_result is not None
+    return False
+
+
+def get_temp_view(name: str) -> Optional[DataFrameContainer]:
+    normalized_name = _normalize(name)
+    for key in _temp_views.keys():
+        normalized_key = _normalize(key[0])
+        if normalized_name == normalized_key and key[1] == get_spark_session_id():
+            return _temp_views.get(key)
+    return None
+
+
+def get_temp_view_normalized_names() -> list[str]:
+    return [
+        _normalize(key[0])
+        for key in _temp_views.keys()
+        if key[1] == get_spark_session_id()
+    ]
+
+
+def _normalize(name: str) -> str:
+    from snowflake.snowpark_connect.config import global_config
+
+    return name if global_config.spark_sql_caseSensitive else name.lower()