snowpark-connect 0.30.0__py3-none-any.whl → 0.31.0__py3-none-any.whl

snowflake/snowpark_connect/utils/pandas_udtf_utils.py

@@ -13,6 +13,10 @@ import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
 from snowflake.snowpark.types import IntegerType, PandasDataFrameType, StructType
 
+# Removed error imports to avoid UDF serialization issues
+# from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+# from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
+
 
 def get_map_in_arrow_udtf(
     user_function: Callable,
@@ -60,14 +64,14 @@ def get_map_in_arrow_udtf(
                 result_iterator, "__iter__"
             ):
                 raise RuntimeError(
-                    f"snowpark_connect::UDF_RETURN_TYPE Return type of the user-defined function should be "
+                    f"[snowpark_connect::type_mismatch] Return type of the user-defined function should be "
                     f"iterator of pyarrow.RecordBatch, but is {type(result_iterator).__name__}"
                 )
 
             for batch in result_iterator:
                 if not isinstance(batch, pa.RecordBatch):
                     raise RuntimeError(
-                        f"snowpark_connect::UDF_RETURN_TYPE Return type of the user-defined function should "
+                        f"[snowpark_connect::type_mismatch] Return type of the user-defined function should "
                         f"be iterator of pyarrow.RecordBatch, but is iterator of {type(batch).__name__}"
                     )
                 if batch.num_rows > 0:
@@ -121,7 +125,7 @@ def create_pandas_udtf(
                 result_iterator, "__iter__"
             ):
                 raise RuntimeError(
-                    f"snowpark_connect::UDF_RETURN_TYPE Return type of the user-defined function should be "
+                    f"[snowpark_connect::type_mismatch] Return type of the user-defined function should be "
                     f"iterator of pandas.DataFrame, but is {type(result_iterator).__name__}"
                 )
 
@@ -140,7 +144,7 @@ def create_pandas_udtf(
                     if column not in self.output_column_original_names
                 ]
                 raise RuntimeError(
-                    f"[RESULT_COLUMNS_MISMATCH_FOR_PANDAS_UDF] Column names of the returned pandas.DataFrame do not match specified schema. Missing: {', '.join(sorted(missing_columns))}. Unexpected: {', '.join(sorted(unexpected_columns))}"
+                    f"[snowpark_connect::invalid_operation] [RESULT_COLUMNS_MISMATCH_FOR_PANDAS_UDF] Column names of the returned pandas.DataFrame do not match specified schema. Missing: {', '.join(sorted(missing_columns))}. Unexpected: {', '.join(sorted(unexpected_columns))}"
                     "."
                 )
             reordered_df = output_df[self.output_column_original_names]

snowflake/snowpark_connect/utils/profiling.py

@@ -10,6 +10,7 @@
 
 import cProfile
 import functools
+import inspect
 import os
 from datetime import datetime
 from typing import Any, Callable
@@ -35,13 +36,29 @@ def profile_method(method: Callable) -> Callable:
         profile_filename = f"{PROFILE_OUTPUT_DIR}/{method_name}_{timestamp}.prof"
 
         profiler = cProfile.Profile()
-        profiler.enable()
-
-        try:
-            result = method(*args, **kwargs)
-            return result
-        finally:
-            profiler.disable()
-            profiler.dump_stats(profile_filename)
+
+        if inspect.isgeneratorfunction(method):
+
+            def profiled_generator():
+                profiler.enable()
+                try:
+                    generator = method(*args, **kwargs)
+                    for item in generator:
+                        profiler.disable()
+                        yield item
+                        profiler.enable()
+                finally:
+                    profiler.disable()
+                    profiler.dump_stats(profile_filename)
+
+            return profiled_generator()
+        else:
+            profiler.enable()
+            try:
+                result = method(*args, **kwargs)
+                return result
+            finally:
+                profiler.disable()
+                profiler.dump_stats(profile_filename)
 
     return wrapper

snowflake/snowpark_connect/utils/scala_udf_utils.py

@@ -22,6 +22,8 @@ from typing import List, Union
 
 import snowflake.snowpark.types as snowpark_type
 import snowflake.snowpark_connect.includes.python.pyspark.sql.connect.proto.types_pb2 as types_proto
+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.resources_initializer import RESOURCE_PATH
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.udf_utils import (
@@ -467,7 +469,9 @@ def map_type_to_scala_type(
         case snowpark_type.VariantType:
             return "Variant"
         case _:
-            raise ValueError(f"Unsupported Snowpark type: {t}")
+            exception = ValueError(f"Unsupported Snowpark type: {t}")
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+            raise exception
 
 
 def map_type_to_snowflake_type(
@@ -533,7 +537,9 @@ def map_type_to_snowflake_type(
         case snowpark_type.VariantType:
             return "VARIANT"
         case _:
-            raise ValueError(f"Unsupported Snowpark type: {t}")
+            exception = ValueError(f"Unsupported Snowpark type: {t}")
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+            raise exception
 
 
 def cast_scala_map_args_from_given_type(
@@ -573,7 +579,9 @@ def cast_scala_map_args_from_given_type(
             case snowpark_type.TimestampType | "timestamp" | "timestamp_ntz":
                 return "java.sql.Timestamp.valueOf({arg_name})"
             case _:
-                raise ValueError(f"Unsupported Snowpark type: {t}")
+                exception = ValueError(f"Unsupported Snowpark type: {t}")
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+                raise exception
 
     if (is_snowpark_type and isinstance(input_type, snowpark_type.MapType)) or (
         not is_snowpark_type and input_type.WhichOneof("kind") == "map"

snowflake/snowpark_connect/utils/session.py

@@ -11,6 +11,8 @@ from snowflake import snowpark
 from snowflake.snowpark.exceptions import SnowparkClientException
 from snowflake.snowpark.session import _get_active_session
 from snowflake.snowpark_connect.constants import DEFAULT_CONNECTION_NAME
+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.describe_query_cache import (
     instrument_session_for_describe_cache,
 )
@@ -120,7 +122,6 @@ def configure_snowpark_session(session: snowpark.Session):
         "PYTHON_SNOWPARK_USE_SCOPED_TEMP_OBJECTS": "false",  # this is required for creating udfs from sproc
         "ENABLE_STRUCTURED_TYPES_IN_SNOWPARK_CONNECT_RESPONSE": "true",
         "QUERY_TAG": f"'{query_tag}'",
-        "FEATURE_INTERVAL_TYPES": "enabled",
     }
 
     session.sql(
@@ -192,7 +193,9 @@ def set_query_tags(spark_tags: Sequence[str]) -> None:
     """Sets Snowpark session query_tag value to the tag from the Spark request."""
 
     if any("," in tag for tag in spark_tags):
-        raise ValueError("Tags cannot contain ','.")
+        exception = ValueError("Tags cannot contain ','.")
+        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+        raise exception
 
     # TODO: Tags might not be set correctly in parallel workloads or multi-threaded code.
     snowpark_session = get_or_create_snowpark_session()

snowflake/snowpark_connect/utils/telemetry.py

@@ -11,8 +11,8 @@ 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
@@ -27,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
 
@@ -57,6 +58,7 @@ class TelemetryType(Enum):
 
 class EventType(Enum):
     SERVER_STARTED = "scos_server_started"
+    WARNING = "scos_warning"
 
 
 # global labels
@@ -106,7 +108,16 @@ REDACTED_PLAN_SUFFIXES = [
 ]
 
 
-def _basic_telemetry_data() -> Dict:
+@dataclass
+class TelemetryMessage:
+    """Container for telemetry messages in the processing queue."""
+
+    message: dict
+    timestamp: int
+    is_warning: bool
+
+
+def _basic_telemetry_data() -> dict:
     return {
         **STATIC_TELEMETRY_DATA,
         TelemetryField.KEY_EVENT_ID.value: str(uuid.uuid4()),
@@ -123,9 +134,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
 
@@ -334,6 +347,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
@@ -528,8 +546,8 @@ 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
 
@@ -541,14 +559,56 @@ class Telemetry:
         }
         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 = {
+                **_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")
@@ -566,13 +626,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()
 
@@ -585,7 +648,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:
@@ -594,7 +657,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).
@@ -675,7 +738,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 = {}

snowflake/snowpark_connect/utils/temporary_view_cache.py

@@ -7,6 +7,8 @@ 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_session_id
 
@@ -22,9 +24,11 @@ def register_temp_view(name: str, df: DataFrameContainer, replace: bool) -> None
                 _temp_views.remove(key)
                 break
             else:
-                raise AnalysisException(
+                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
 

snowflake/snowpark_connect/utils/udf_cache.py

@@ -118,7 +118,7 @@ def cached_udaf(
 
     if class_type is None:
         raise ValueError(
-            "Type must be provided for cached_udaf. UDAF contains multiple functions hence it has to be represented by a type. Functions are not supported."
+            "[snowpark_connect::internal_error] Type must be provided for cached_udaf. UDAF contains multiple functions hence it has to be represented by a type. Functions are not supported."
         )
     else:
         # return udaf
@@ -379,7 +379,9 @@ def register_cached_java_udf(
 
                 with zipfile.ZipFile(zip_path, "r") as zip_ref:
                     if jar_path_in_zip not in zip_ref.namelist():
-                        raise FileNotFoundError(f"{jar_path_in_zip} not found")
+                        raise FileNotFoundError(
+                            f"[snowpark_connect::invalid_input] {jar_path_in_zip} not found"
+                        )
                     zip_ref.extract(jar_path_in_zip, temp_dir)
 
                 jar_path = f"{temp_dir}/{jar_path_in_zip}"
@@ -388,7 +390,7 @@ def register_cached_java_udf(
 
             if upload_result[0].status != "UPLOADED":
                 raise RuntimeError(
-                    f"Failed to upload JAR with UDF definitions to stage: {upload_result[0].message}"
+                    f"[snowpark_connect::internal_error] Failed to upload JAR with UDF definitions to stage: {upload_result[0].message}"
                 )
 
         udf_is_cached = function_name in cache

snowflake/snowpark_connect/utils/udf_helper.py

@@ -18,6 +18,8 @@ from snowflake.snowpark import Column, Session
 from snowflake.snowpark.types import DataType, _parse_datatype_json_value
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.config import global_config
+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.expression.map_expression import (
     map_single_column_expression,
 )
@@ -233,19 +235,25 @@ def _check_supported_udf(
 
             session = get_or_create_snowpark_session()
             if udf_proto.java_udf.class_name not in session._cached_java_udfs:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Can not load class {udf_proto.java_udf.class_name}"
                 )
+                attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+                raise exception
             else:
-                raise ValueError(
+                exception = ValueError(
                     "Function type java_udf not supported for common inline user-defined function"
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+                raise exception
         case "scalar_scala_udf":
             pass
         case _ as function_type:
-            raise ValueError(
+            exception = ValueError(
                 f"Function type {function_type} not supported for common inline user-defined function"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def _aggregate_function_check(
@@ -253,9 +261,11 @@ def _aggregate_function_check(
 ):
     name, is_aggregate_function = get_is_aggregate_function()
     if not udf_proto.deterministic and name != "default" and is_aggregate_function:
-        raise AnalysisException(
+        exception = AnalysisException(
             f"[AGGREGATE_FUNCTION_WITH_NONDETERMINISTIC_EXPRESSION] Non-deterministic expression {name}({udf_proto.function_name}) should not appear in the arguments of an aggregate function."
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+        raise exception
 
 
 def _join_checks(snowpark_udf_arg_names: list[str]):
@@ -282,23 +292,27 @@ def _join_checks(snowpark_udf_arg_names: list[str]):
             and is_left_evaluable
             and is_right_evaluable
         ):
-            raise AnalysisException(
+            exception = AnalysisException(
                 f"Detected implicit cartesian product for {is_evaluating_join_condition[0]} join between logical plans. \n"
                 f"Join condition is missing or trivial. \n"
                 f"Either: use the CROSS JOIN syntax to allow cartesian products between those relations, or; "
                 f"enable implicit cartesian products by setting the configuration variable spark.sql.crossJoin.enabled=True."
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+            raise exception
         if (
             is_evaluating_join_condition[0] != "INNER"
             and is_evaluating_join_condition[1]
             and is_left_evaluable
             and is_right_evaluable
         ):
-            raise AnalysisException(
+            exception = AnalysisException(
                 f"[UNSUPPORTED_FEATURE.PYTHON_UDF_IN_ON_CLAUSE] The feature is not supported: "
                 f"Python UDF in the ON clause of a {is_evaluating_join_condition[0]} JOIN. "
                 f"In case of an INNNER JOIN consider rewriting to a CROSS JOIN with a WHERE clause."
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def infer_snowpark_arguments(

snowflake/snowpark_connect/utils/udf_utils.py

@@ -103,7 +103,7 @@ class ProcessCommonInlineUserDefinedFunction:
                 )
             case _:
                 raise ValueError(
-                    f"Function type {self._function_type} not supported for common inline user-defined function"
+                    f"[snowpark_connect::unsupported_operation] Function type {self._function_type} not supported for common inline user-defined function"
                 )
 
     @property
@@ -112,7 +112,7 @@ class ProcessCommonInlineUserDefinedFunction:
             return self._snowpark_udf_args
         else:
             raise ValueError(
-                "Column mapping is not provided, cannot get snowpark udf args"
+                "[snowpark_connect::internal_error] Column mapping is not provided, cannot get snowpark udf args"
             )
 
     @property
@@ -121,7 +121,7 @@ class ProcessCommonInlineUserDefinedFunction:
             return self._snowpark_udf_arg_names
         else:
             raise ValueError(
-                "Column mapping is not provided, cannot get snowpark udf arg names"
+                "[snowpark_connect::internal_error] Column mapping is not provided, cannot get snowpark udf arg names"
             )
 
     def _create_python_udf(self):
@@ -291,5 +291,5 @@ class ProcessCommonInlineUserDefinedFunction:
                 return create_scala_udf(self)
             case _:
                 raise ValueError(
-                    f"Function type {self._function_type} not supported for common inline user-defined function"
+                    f"[snowpark_connect::unsupported_operation] Function type {self._function_type} not supported for common inline user-defined function"
                 )

snowflake/snowpark_connect/utils/udtf_helper.py

@@ -16,6 +16,8 @@ import snowflake.snowpark_connect.tcm as tcm
 from snowflake import snowpark
 from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
 from snowflake.snowpark.types import DataType, StructType, _parse_datatype_json_value
+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.type_mapping import proto_to_snowpark_type
 from snowflake.snowpark_connect.utils import pandas_udtf_utils, udtf_utils
 from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
@@ -37,7 +39,9 @@ def udtf_check(
     udtf_proto: relation_proto.CommonInlineUserDefinedTableFunction,
 ) -> None:
     if udtf_proto.WhichOneof("function") != "python_udtf":
-        raise ValueError(f"Not python udtf {udtf_proto.function}")
+        exception = ValueError(f"Not python udtf {udtf_proto.function}")
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
 
 def require_creating_udtf_in_sproc(