snowpark-connect 0.30.1__py3-none-any.whl → 0.32.0__py3-none-any.whl

snowflake/snowpark_connect/utils/env_utils.py

@@ -8,6 +8,8 @@ Environment variable utilities for Snowpark Connect.
 
 import os
 
+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.snowpark_connect_logging import logger
 
 
@@ -37,9 +39,11 @@ def get_int_from_env(env_var: str, default: int) -> int:
     """
     # Validate that default is actually an integer
     if not isinstance(default, int):
-        raise TypeError(
+        exception = TypeError(
             f"Default value must be an integer, got {type(default).__name__}: {default}"
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+        raise exception
 
     value = os.getenv(env_var)
     if value is None:

snowflake/snowpark_connect/utils/identifiers.py

@@ -13,6 +13,8 @@ from snowflake.snowpark_connect.config import (
     auto_uppercase_column_identifiers,
     auto_uppercase_non_column_identifiers,
 )
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 
 QUOTED_SPARK_IDENTIFIER = re.compile(r"^`[^`]*(?:``[^`]*)*`$")
 UNQUOTED_SPARK_IDENTIFIER = re.compile(r"^\w+$")
@@ -25,7 +27,9 @@ def unquote_spark_identifier_if_quoted(spark_name: str) -> str:
     if QUOTED_SPARK_IDENTIFIER.match(spark_name):
         return spark_name[1:-1].replace("``", "`")
 
-    raise AnalysisException(f"Invalid name: {spark_name}")
+    exception = AnalysisException(f"Invalid name: {spark_name}")
+    attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+    raise exception
 
 
 def spark_to_sf_single_id_with_unquoting(
@@ -191,7 +195,9 @@ class FQN:
 
     def __eq__(self, other: Any) -> bool:
         if not isinstance(other, FQN):
-            raise AnalysisException(f"{other} is not a valid FQN")
+            exception = AnalysisException(f"{other} is not a valid FQN")
+            attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+            raise exception
         return self.identifier == other.identifier
 
     @classmethod
@@ -209,7 +215,9 @@ class FQN:
         result = re.fullmatch(qualifier_pattern, identifier)
 
         if result is None:
-            raise AnalysisException(f"{identifier} is not a valid identifier")
+            exception = AnalysisException(f"{identifier} is not a valid identifier")
+            attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+            raise exception
 
         unqualified_name = result.group("name")
         if result.group("second_qualifier") is not None:

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,
 )
@@ -21,6 +23,13 @@ from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import telemetry
 from snowflake.snowpark_connect.utils.udf_cache import init_builtin_udf_cache
 
+SKIP_SESSION_CONFIGURATION = False
+
+
+def skip_session_configuration(skip: bool):
+    global SKIP_SESSION_CONFIGURATION
+    SKIP_SESSION_CONFIGURATION = skip
+
 
 # Suppress experimental warnings from snowflake.snowpark logger
 def _filter_experimental_warnings(record):
@@ -55,6 +64,8 @@ def configure_snowpark_session(session: snowpark.Session):
         global_config,
     )
 
+    global SKIP_SESSION_CONFIGURATION
+
     logger.info(f"Configuring session {session}")
 
     telemetry.initialize(session)
@@ -122,9 +133,16 @@ def configure_snowpark_session(session: snowpark.Session):
         "QUERY_TAG": f"'{query_tag}'",
     }
 
-    session.sql(
-        f"ALTER SESSION SET {', '.join([f'{k} = {v}' for k, v in session_params.items()])}"
-    ).collect()
+    # SNOW-2245971: Stored procedures inside Native Apps run as Execute As Owner and hence cannot set session params.
+    if not SKIP_SESSION_CONFIGURATION:
+        session.sql(
+            f"ALTER SESSION SET {', '.join([f'{k} = {v}' for k, v in session_params.items()])}"
+        ).collect()
+    else:
+        session_param_names = ", ".join(session_params.keys())
+        logger.info(
+            f"Skipping Snowpark Connect session configuration as requested. Please make sure following session parameters are set correctly: {session_param_names}"
+        )
 
     # Instrument the snowpark session to use a cache for describe queries.
     instrument_session_for_describe_cache(session)
@@ -191,7 +209,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

@@ -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
 
@@ -346,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

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(