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

snowflake/snowpark_connect/server.py

@@ -56,7 +56,11 @@ from snowflake.snowpark_connect.analyze_plan.map_tree_string import map_tree_str
 from snowflake.snowpark_connect.config import route_config_proto
 from snowflake.snowpark_connect.constants import SERVER_SIDE_SESSION_ID
 from snowflake.snowpark_connect.control_server import ControlServicer
-from snowflake.snowpark_connect.error.error_utils import build_grpc_error_response
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import (
+    attach_custom_error_code,
+    build_grpc_error_response,
+)
 from snowflake.snowpark_connect.execute_plan.map_execution_command import (
     map_execution_command,
 )
@@ -96,7 +100,7 @@ from snowflake.snowpark_connect.utils.interrupt import (
     interrupt_queries_with_tag,
     interrupt_query,
 )
-from snowflake.snowpark_connect.utils.profiling import profile_method
+from snowflake.snowpark_connect.utils.profiling import PROFILING_ENABLED, profile_method
 from snowflake.snowpark_connect.utils.session import (
     configure_snowpark_session,
     get_or_create_snowpark_session,
@@ -157,6 +161,7 @@ def _handle_exception(context, e: Exception):
 
     if tcm.TCM_MODE:
         # TODO: SNOW-2009834 gracefully return error back in TCM
+        attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
         raise e
 
     from grpc_status import rpc_status
@@ -374,9 +379,13 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
                         ),
                     )
                 case _:
-                    raise SnowparkConnectNotImplementedError(
+                    exception = SnowparkConnectNotImplementedError(
                         f"ANALYZE PLAN NOT IMPLEMENTED:\n{request}"
                     )
+                    attach_custom_error_code(
+                        exception, ErrorCodes.UNSUPPORTED_OPERATION
+                    )
+                    raise exception
         except Exception as e:
             _handle_exception(context, e)
         finally:
@@ -527,9 +536,13 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
                             ),
                         )
                 case _:
-                    raise ValueError(
+                    exception = ValueError(
                         f"Unexpected payload type in AddArtifacts: {request.WhichOneof('payload')}"
                     )
+                    attach_custom_error_code(
+                        exception, ErrorCodes.UNSUPPORTED_OPERATION
+                    )
+                    raise exception
 
         for name, data in cache_data.items():
             _try_handle_local_relation(name, bytes(data))
@@ -635,9 +648,13 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
                 case proto_base.InterruptRequest.InterruptType.INTERRUPT_TYPE_OPERATION_ID:
                     interrupted_ids = interrupt_query(request.operation_id)
                 case _:
-                    raise SnowparkConnectNotImplementedError(
+                    exception = SnowparkConnectNotImplementedError(
                         f"INTERRUPT NOT IMPLEMENTED:\n{request}"
                     )
+                    attach_custom_error_code(
+                        exception, ErrorCodes.UNSUPPORTED_OPERATION
+                    )
+                    raise exception
 
             return proto_base.InterruptResponse(
                 session_id=request.session_id,
@@ -655,9 +672,11 @@ class SnowflakeConnectServicer(proto_base_grpc.SparkConnectServiceServicer):
         continue. If there is a ResultComplete, the client should use ReleaseExecute with
         """
         logger.info("ReattachExecute")
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "Spark client has detached, please resubmit request. In a future version, the server will be support the reattach."
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     def ReleaseExecute(self, request: proto_base.ReleaseExecuteRequest, context):
         """Release an reattachable execution, or parts thereof.
@@ -760,8 +779,11 @@ def _serve(
 
         ChannelBuilder.MAX_MESSAGE_LENGTH = grpc_max_msg_size
 
+        # cProfile doesn't work correctly with multiple threads
+        max_workers = 1 if PROFILING_ENABLED else 10
+
         server = grpc.server(
-            futures.ThreadPoolExecutor(max_workers=10), options=server_options
+            futures.ThreadPoolExecutor(max_workers=max_workers), options=server_options
         )
         control_servicer = ControlServicer(session)
         proto_base_grpc.add_SparkConnectServiceServicer_to_server(
@@ -791,6 +813,7 @@ def _serve(
             )
         else:
             logger.error("Error starting up Snowpark Connect server", exc_info=True)
+        attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
         raise e
     finally:
         # flush the telemetry queue if possible
@@ -808,7 +831,9 @@ def _set_remote_url(remote_url: str):
     elif parsed_url.scheme == "unix":
         _server_url = remote_url.split("/;")[0]
     else:
-        raise RuntimeError(f"Invalid Snowpark Connect URL: {remote_url}")
+        exception = RuntimeError(f"Invalid Snowpark Connect URL: {remote_url}")
+        attach_custom_error_code(exception, ErrorCodes.INVALID_SPARK_CONNECT_URL)
+        raise exception
 
 
 def _set_server_tcp_port(server_port: int):
@@ -822,7 +847,9 @@ def _check_port_is_free(port: int) -> None:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
         s.settimeout(1)
         if s.connect_ex(("127.0.0.1", port)) == 0:
-            raise RuntimeError(f"TCP port {port} is already in use")
+            exception = RuntimeError(f"TCP port {port} is already in use")
+            attach_custom_error_code(exception, ErrorCodes.TCP_PORT_ALREADY_IN_USE)
+            raise exception
 
 
 def _set_server_unix_domain_socket(path: str):
@@ -834,14 +861,18 @@ def _set_server_unix_domain_socket(path: str):
 def get_server_url() -> str:
     global _server_url
     if not _server_url:
-        raise RuntimeError("Server URL not set")
+        exception = RuntimeError("Server URL not set")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
     return _server_url
 
 
 def get_client_url() -> str:
     global _client_url
     if not _client_url:
-        raise RuntimeError("Client URL not set")
+        exception = RuntimeError("Client URL not set")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
     return _client_url
 
 
@@ -871,12 +902,14 @@ class UnixDomainSocketChannelBuilder(ChannelBuilder):
         if url is None:
             url = get_client_url()
         if url[:6] != "unix:/" or len(url) < 7:
-            raise PySparkValueError(
+            exception = PySparkValueError(
                 error_class="INVALID_CONNECT_URL",
                 message_parameters={
                     "detail": "The URL must start with 'unix://'. Please update the URL to follow the correct format, e.g., 'unix://unix_domain_socket_path'.",
                 },
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_SPARK_CONNECT_URL)
+            raise exception
 
         # Rewrite the URL to use http as the scheme so that we can leverage
         # Python's built-in parser to parse URL parameters
@@ -919,7 +952,7 @@ class UnixDomainSocketChannelBuilder(ChannelBuilder):
             for p in parts:
                 kv = p.split("=")
                 if len(kv) != 2:
-                    raise PySparkValueError(
+                    exception = PySparkValueError(
                         error_class="INVALID_CONNECT_URL",
                         message_parameters={
                             "detail": f"Parameter '{p}' should be provided as a "
@@ -927,6 +960,10 @@ class UnixDomainSocketChannelBuilder(ChannelBuilder):
                             f"the parameter to follow the correct format, e.g., 'key=value'.",
                         },
                     )
+                    attach_custom_error_code(
+                        exception, ErrorCodes.INVALID_SPARK_CONNECT_URL
+                    )
+                    raise exception
                 self.params[kv[0]] = urllib.parse.unquote(kv[1])
 
         netloc = self.url.netloc.split(":")
@@ -942,7 +979,7 @@ class UnixDomainSocketChannelBuilder(ChannelBuilder):
             self.host = netloc[0]
             self.port = int(netloc[1])
         else:
-            raise PySparkValueError(
+            exception = PySparkValueError(
                 error_class="INVALID_CONNECT_URL",
                 message_parameters={
                     "detail": f"Target destination '{self.url.netloc}' should match the "
@@ -950,6 +987,8 @@ class UnixDomainSocketChannelBuilder(ChannelBuilder):
                     f"the correct format, e.g., 'hostname:port'.",
                 },
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_SPARK_CONNECT_URL)
+            raise exception
 
     # We override this to enable compatibility with Spark 4.0
     host = None
@@ -988,9 +1027,11 @@ def start_jvm():
         if tcm.TCM_MODE:
             # No-op if JVM is already started in TCM mode
             return
-        raise RuntimeError(
+        exception = RuntimeError(
             "JVM must not be running when starting the Spark Connect server"
         )
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
     pyspark_jars = (
         pathlib.Path(snowflake.snowpark_connect.__file__).parent / "includes/jars"
@@ -1065,14 +1106,18 @@ def start_session(
         _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE = max_grpc_message_size
 
         if os.environ.get("SPARK_ENV_LOADED"):
-            raise RuntimeError(
+            exception = RuntimeError(
                 "Snowpark Connect cannot be run inside of a Spark environment"
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_STARTUP_OPERATION)
+            raise exception
         if connection_parameters is not None:
             if snowpark_session is not None:
-                raise ValueError(
+                exception = ValueError(
                     "Only specify one of snowpark_session and connection_parameters"
                 )
+                attach_custom_error_code(exception, ErrorCodes.INVALID_STARTUP_INPUT)
+                raise exception
             snowpark_session = snowpark.Session.builder.configs(
                 connection_parameters
             ).create()
@@ -1084,9 +1129,11 @@ def start_session(
             return
 
         if len(list(filter(None, [remote_url, tcp_port, unix_domain_socket]))) > 1:
-            raise RuntimeError(
+            exception = RuntimeError(
                 "Can only set at most one of remote_url, tcp_port, and unix_domain_socket"
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_STARTUP_INPUT)
+            raise exception
 
         url_from_env = os.environ.get("SPARK_REMOTE", None)
         if remote_url:
@@ -1124,7 +1171,11 @@ def start_session(
             server_thread.start()
             _server_running.wait()
             if _server_error:
-                raise RuntimeError("Snowpark Connect session failed to start")
+                exception = RuntimeError("Snowpark Connect session failed to start")
+                attach_custom_error_code(
+                    exception, ErrorCodes.STARTUP_CONNECTION_FAILED
+                )
+                raise exception
             return server_thread
         else:
             # Launch in the foreground.
@@ -1132,6 +1183,7 @@ def start_session(
     except Exception as e:
         _reset_server_run_state()
         logger.error(e, exc_info=True)
+        attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
         raise e
 
 
@@ -1165,6 +1217,7 @@ def get_session(url: Optional[str] = None, conf: SparkConf = None) -> SparkSessi
     except Exception as e:
         _reset_server_run_state()
         logger.error(e, exc_info=True)
+        attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
         raise e
 
 

snowflake/snowpark_connect/type_mapping.py

@@ -29,12 +29,17 @@ from snowflake.snowpark_connect.constants import (
 from snowflake.snowpark_connect.date_time_format_mapping import (
     convert_spark_format_to_snowflake,
 )
+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.literal import get_literal_field_and_name
 from snowflake.snowpark_connect.expression.map_sql_expression import (
     _INTERVAL_DAYTIME_PATTERN_RE,
     _INTERVAL_YEARMONTH_PATTERN_RE,
 )
-from snowflake.snowpark_connect.utils.context import get_is_evaluating_sql
+from snowflake.snowpark_connect.utils.context import (
+    get_is_evaluating_sql,
+    get_jpype_jclass_lock,
+)
 from snowflake.snowpark_connect.utils.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import (
     SnowparkConnectNotImplementedError,
@@ -61,12 +66,14 @@ SNOWPARK_TYPE_NAME_TO_PYSPARK_TYPE_NAME = {
 
 @cache
 def _get_struct_type_class():
-    return jpype.JClass("org.apache.spark.sql.types.StructType")
+    with get_jpype_jclass_lock():
+        return jpype.JClass("org.apache.spark.sql.types.StructType")
 
 
 @cache
 def get_python_sql_utils_class():
-    return jpype.JClass("org.apache.spark.sql.api.python.PythonSQLUtils")
+    with get_jpype_jclass_lock():
+        return jpype.JClass("org.apache.spark.sql.api.python.PythonSQLUtils")
 
 
 def _parse_ddl_with_spark_scala(ddl_string: str) -> pyspark.sql.types.DataType:
@@ -291,9 +298,11 @@ def snowpark_to_proto_type(
                 )
             }
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type: {data_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def cast_to_match_snowpark_type(
@@ -333,7 +342,9 @@ def cast_to_match_snowpark_type(
                 with suppress(TypeError):
                     date = datetime.strptime(content, format)
                     return date
-            raise ValueError(f"Date casting error for {str(content)}")
+            exception = ValueError(f"Date casting error for {str(content)}")
+            attach_custom_error_code(exception, ErrorCodes.INVALID_CAST)
+            raise exception
         case snowpark.types.ShortType:
             return int(content)
         case snowpark.types.StringType:
@@ -363,9 +374,11 @@ def cast_to_match_snowpark_type(
         case snowpark.types.DayTimeIntervalType:
             return str(content)
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type in casting: {data_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def snowpark_to_iceberg_type(data_type: snowpark.types.DataType) -> str:
@@ -398,9 +411,11 @@ def snowpark_to_iceberg_type(data_type: snowpark.types.DataType) -> str:
         case snowpark.types.TimestampType:
             return "timestamp"
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type for iceber: {data_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def proto_to_snowpark_type(
@@ -487,9 +502,11 @@ def map_snowpark_types_to_pyarrow_types(
                     )
                 )
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Unsupported arrow type {pa_type} for snowpark ArrayType."
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+                raise exception
         case snowpark.types.BinaryType:
             return pa.binary()
         case snowpark.types.BooleanType:
@@ -530,9 +547,11 @@ def map_snowpark_types_to_pyarrow_types(
                     ),
                 )
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Unsupported arrow type {pa_type} for snowpark MapType."
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+                raise exception
         case snowpark.types.NullType:
             return pa.string()
         case snowpark.types.ShortType:
@@ -557,9 +576,11 @@ def map_snowpark_types_to_pyarrow_types(
                     ]
                 )
             else:
-                raise AnalysisException(
+                exception = AnalysisException(
                     f"Unsupported arrow type {pa_type} for snowpark StructType."
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+                raise exception
         case snowpark.types.TimestampType:
             unit = pa_type.unit
             tz = pa_type.tz
@@ -576,9 +597,11 @@ def map_snowpark_types_to_pyarrow_types(
             # Return string type so formatted intervals are preserved in display
             return pa.string()
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported snowpark data type: {snowpark_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def map_pyarrow_to_snowpark_types(pa_type: pa.DataType) -> snowpark.types.DataType:
@@ -648,9 +671,11 @@ def map_pyarrow_to_snowpark_types(pa_type: pa.DataType) -> snowpark.types.DataTy
     elif pa.types.is_null(pa_type):
         return snowpark.types.NullType()
     else:
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             f"Unsupported PyArrow data type: {pa_type}"
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
 
 def map_pyspark_types_to_snowpark_types(
@@ -736,9 +761,11 @@ def map_pyspark_types_to_snowpark_types(
         return snowpark.types.DayTimeIntervalType(
             type_to_map.startField, type_to_map.endField
         )
-    raise SnowparkConnectNotImplementedError(
+    exception = SnowparkConnectNotImplementedError(
         f"Unsupported spark data type: {type_to_map}"
     )
+    attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+    raise exception
 
 
 def map_snowpark_to_pyspark_types(
@@ -811,7 +838,11 @@ def map_snowpark_to_pyspark_types(
         return pyspark.sql.types.DayTimeIntervalType(
             type_to_map.start_field, type_to_map.end_field
         )
-    raise SnowparkConnectNotImplementedError(f"Unsupported data type: {type_to_map}")
+    exception = SnowparkConnectNotImplementedError(
+        f"Unsupported data type: {type_to_map}"
+    )
+    attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+    raise exception
 
 
 def map_simple_types(simple_type: str) -> snowpark.types.DataType:
@@ -866,9 +897,11 @@ def map_simple_types(simple_type: str) -> snowpark.types.DataType:
                 precision = int(simple_type.split("(")[1].split(",")[0])
                 scale = int(simple_type.split(",")[1].split(")")[0])
                 return snowpark.types.DecimalType(precision, scale)
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported simple type: {simple_type}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def map_json_schema_to_snowpark(
@@ -1009,9 +1042,11 @@ def map_spark_timestamp_format_expression(
             lit_value, _ = get_literal_field_and_name(arguments.literal)
             return convert_spark_format_to_snowflake(lit_value, timestamp_input_type)
         case other:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported expression type {other} in timestamp format argument"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def map_spark_number_format_expression(
@@ -1030,9 +1065,11 @@ def map_spark_number_format_expression(
             case "literal":
                 lit_value, _ = get_literal_field_and_name(arguments.literal)
             case other:
-                raise SnowparkConnectNotImplementedError(
+                exception = SnowparkConnectNotImplementedError(
                     f"Unsupported expression type {other} in number format argument"
                 )
+                attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+                raise exception
 
     return _map_spark_to_snowflake_number_format(lit_value)
 

snowflake/snowpark_connect/utils/context.py

@@ -2,10 +2,12 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
+import os
 import re
+import threading
 from contextlib import contextmanager
 from contextvars import ContextVar
-from typing import Mapping, Optional
+from typing import Iterator, Mapping, Optional
 
 import pyspark.sql.connect.proto.expressions_pb2 as expressions_proto
 
@@ -25,6 +27,7 @@ _is_evaluating_sql = ContextVar[bool]("_is_evaluating_sql", default=False)
 _is_evaluating_join_condition = ContextVar(
     "_is_evaluating_join_condition", default=("default", False, [], [])
 )
+_is_processing_order_by = ContextVar[bool]("_is_processing_order_by", default=False)
 
 _sql_aggregate_function_count = ContextVar[int](
     "_contains_aggregate_function", default=0
@@ -56,6 +59,23 @@ _is_in_pivot = ContextVar[bool]("_is_in_pivot", default=False)
 _is_in_udtf_context = ContextVar[bool]("_is_in_udtf_context", default=False)
 _accessing_temp_object = ContextVar[bool]("_accessing_temp_object", default=False)
 
+# Thread-safe lock for JPype JClass creation to prevent access violations
+_jpype_jclass_lock = threading.Lock()
+
+
+@contextmanager
+def get_jpype_jclass_lock() -> Iterator[None]:
+    """
+    Context manager that acquires the JPype JClass lock on Windows platforms.
+    On non-Windows (os.name != 'nt'), it yields without acquiring the lock.
+    """
+    if os.name == "nt":
+        with _jpype_jclass_lock:
+            yield
+    else:
+        yield
+
+
 # Lateral Column Alias helpers
 # We keep a thread-local mapping from alias name -> TypedColumn that is
 # populated incrementally while the projection list is being processed.
@@ -207,6 +227,27 @@ def push_evaluating_sql_scope():
         _is_evaluating_sql.set(prev)
 
 
+def get_is_processing_order_by() -> bool:
+    """
+    Gets the value of _is_processing_order_by for the current context, defaults to False.
+    """
+    return _is_processing_order_by.get()
+
+
+@contextmanager
+def push_processing_order_by_scope():
+    """
+    Context manager that sets a flag indicating if ORDER BY expressions are being evaluated.
+    This enables optimizations like reusing already-computed UDF columns.
+    """
+    prev = _is_processing_order_by.get()
+    try:
+        _is_processing_order_by.set(True)
+        yield
+    finally:
+        _is_processing_order_by.set(prev)
+
+
 def get_is_evaluating_join_condition() -> tuple[str, bool, list, list]:
     """
     Gets the value of _is_evaluating_join_condition for the current context, defaults to False.

snowflake/snowpark_connect/utils/describe_query_cache.py

@@ -12,6 +12,8 @@ from typing import Any
 from snowflake import snowpark
 from snowflake.connector.cursor import ResultMetadataV2
 from snowflake.snowpark._internal.server_connection import ServerConnection
+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.snowpark_connect_logging import logger
 from snowflake.snowpark_connect.utils.telemetry import telemetry
@@ -148,6 +150,7 @@ def instrument_session_for_describe_cache(session: snowpark.Session):
                 telemetry.report_query(result, **kwargs)
             except Exception as e:
                 telemetry.report_query(e, **kwargs)
+                attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
                 raise e
             return result
 

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: