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

snowflake/snowpark_connect/relation/write/map_write.py

@@ -214,27 +214,71 @@ def map_write(request: proto_base.ExecutePlanRequest):
             )
             snowpark_table_name = _spark_to_snowflake(table_name)
 
-            if write_mode == "overwrite":
-                if check_snowflake_table_existance(snowpark_table_name, session):
-                    session.sql(f"DELETE FROM {snowpark_table_name}").collect()
-                    write_mode = "append"
-
-            if write_mode in (None, "", "overwrite"):
-                create_iceberg_table(
-                    snowpark_table_name=snowpark_table_name,
-                    location=write_op.options.get("location", None),
-                    schema=input_df.schema,
-                    snowpark_session=session,
-                )
-                write_mode = "append"
-
-            _validate_schema_and_get_writer(
-                input_df, write_mode, snowpark_table_name
-            ).saveAsTable(
-                table_name=snowpark_table_name,
-                mode=write_mode,
-                column_order=_column_order_for_write,
-            )
+            match write_mode:
+                case None | "error" | "errorifexists":
+                    if check_snowflake_table_existence(snowpark_table_name, session):
+                        raise AnalysisException(
+                            f"Table {snowpark_table_name} already exists"
+                        )
+                    create_iceberg_table(
+                        snowpark_table_name=snowpark_table_name,
+                        location=write_op.options.get("location", None),
+                        schema=input_df.schema,
+                        snowpark_session=session,
+                    )
+                    _validate_schema_and_get_writer(
+                        input_df, "append", snowpark_table_name
+                    ).saveAsTable(
+                        table_name=snowpark_table_name,
+                        mode="append",
+                        column_order=_column_order_for_write,
+                    )
+                case "append":
+                    _validate_schema_and_get_writer(
+                        input_df, "append", snowpark_table_name
+                    ).saveAsTable(
+                        table_name=snowpark_table_name,
+                        mode="append",
+                        column_order=_column_order_for_write,
+                    )
+                case "ignore":
+                    if not check_snowflake_table_existence(
+                        snowpark_table_name, session
+                    ):
+                        create_iceberg_table(
+                            snowpark_table_name=snowpark_table_name,
+                            location=write_op.options.get("location", None),
+                            schema=input_df.schema,
+                            snowpark_session=session,
+                        )
+                        _validate_schema_and_get_writer(
+                            input_df, "append", snowpark_table_name
+                        ).saveAsTable(
+                            table_name=snowpark_table_name,
+                            mode="append",
+                            column_order=_column_order_for_write,
+                        )
+                case "overwrite":
+                    if check_snowflake_table_existence(snowpark_table_name, session):
+                        session.sql(f"DELETE FROM {snowpark_table_name}").collect()
+                    else:
+                        create_iceberg_table(
+                            snowpark_table_name=snowpark_table_name,
+                            location=write_op.options.get("location", None),
+                            schema=input_df.schema,
+                            snowpark_session=session,
+                        )
+                    _validate_schema_and_get_writer(
+                        input_df, "append", snowpark_table_name
+                    ).saveAsTable(
+                        table_name=snowpark_table_name,
+                        mode="append",
+                        column_order=_column_order_for_write,
+                    )
+                case _:
+                    raise SnowparkConnectNotImplementedError(
+                        f"Write mode {write_mode} is not supported"
+                    )
         case _:
             snowpark_table_name = _spark_to_snowflake(write_op.table.table_name)
 
@@ -299,14 +343,14 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
         commands_proto.WriteOperationV2.MODE_OVERWRITE,
         commands_proto.WriteOperationV2.MODE_APPEND,
     ):
-        if not check_snowflake_table_existance(snowpark_table_name, session):
+        if not check_snowflake_table_existence(snowpark_table_name, session):
             raise AnalysisException(
                 f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found. "
                 f"Verify the spelling and correctness of the schema and catalog.\n"
             )
 
     if write_op.provider.lower() == "iceberg":
-        if write_mode == "overwrite" and check_snowflake_table_existance(
+        if write_mode == "overwrite" and check_snowflake_table_existence(
             snowpark_table_name, session
         ):
             session.sql(f"DELETE FROM {snowpark_table_name}").collect()
@@ -584,7 +628,7 @@ def _truncate_directory(directory_path: Path) -> None:
             shutil.rmtree(file)
 
 
-def check_snowflake_table_existance(
+def check_snowflake_table_existence(
     snowpark_table_name: str,
     snowpark_session: snowpark.Session,
 ):

snowflake/snowpark_connect/server.py

@@ -981,6 +981,7 @@ def start_session(
     stop_event: threading.Event = None,
     snowpark_session: Optional[snowpark.Session] = None,
     connection_parameters: Optional[Dict[str, str]] = None,
+    max_grpc_message_size: int = _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE,
 ) -> threading.Thread | None:
     """
     Starts Spark Connect server connected to Snowflake. No-op if the Server is already running.
@@ -1003,6 +1004,14 @@ def start_session(
                                 provided, the `snowpark_session` parameter must be None.
     """
     try:
+        # Changing the value of our global variable based on the grpc message size provided by the user.
+        global _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE
+        _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE = max_grpc_message_size
+
+        from pyspark.sql.connect.client import ChannelBuilder
+
+        ChannelBuilder.MAX_MESSAGE_LENGTH = max_grpc_message_size
+
         if os.environ.get("SPARK_ENV_LOADED"):
             raise RuntimeError(
                 "Snowpark Connect cannot be run inside of a Spark environment"

snowflake/snowpark_connect/type_mapping.py

@@ -324,6 +324,8 @@ def cast_to_match_snowpark_type(
             return str(content)
         case snowpark.types.VariantType:
             return str(content)
+        case snowpark.types.TimestampType:
+            return str(content)
         case _:
             raise SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type in casting: {data_type}"
@@ -779,6 +781,8 @@ def map_simple_types(simple_type: str) -> snowpark.types.DataType:
             return snowpark.types.TimestampType()
         case "timestamp_ntz":
             return snowpark.types.TimestampType(snowpark.types.TimestampTimeZone.NTZ)
+        case "timestamp_ltz":
+            return snowpark.types.TimestampType(snowpark.types.TimestampTimeZone.LTZ)
         case "day_time_interval":
             # this is not a column type in snowflake so there won't be a dataframe column
             # with this, for now this type won't make any sense

snowflake/snowpark_connect/utils/describe_query_cache.py

@@ -131,21 +131,14 @@ def instrument_session_for_describe_cache(session: snowpark.Session):
             logger.debug(f"DDL detected, clearing describe query cache: '{query}'")
             cache.clear()
 
-    def report_query(qid: str, is_internal: bool) -> None:
-        if is_internal:
-            telemetry.report_internal_query()
-        elif qid:
-            telemetry.report_query_id(qid)
-
     def wrap_execute(wrapped_fn):
         def fn(query: str, **kwargs):
             update_cache_for_query(query)
-            is_internal = kwargs.get("_is_internal", False)
             try:
                 result = wrapped_fn(query, **kwargs)
-                report_query(result.sfqid, is_internal)
+                telemetry.report_query(result, **kwargs)
             except Exception as e:
-                report_query(e.sfqid, is_internal)
+                telemetry.report_query(e, **kwargs)
                 raise e
             return result
 

snowflake/snowpark_connect/utils/session.py

@@ -181,7 +181,3 @@ def set_query_tags(spark_tags: Sequence[str]) -> None:
 
     if spark_tags_str != snowpark_session.query_tag:
         snowpark_session.query_tag = spark_tags_str
-
-
-def get_python_udxf_import_files(session: snowpark.Session) -> str:
-    return ",".join([file for file in [*session._python_files, *session._import_files]])

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