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

snowflake/snowpark_connect/date_time_format_mapping.py

@@ -24,6 +24,8 @@
 from pyspark.errors.exceptions.base import DateTimeException
 
 from snowflake.snowpark.types import DataType, StringType
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+from snowflake.snowpark_connect.error.error_utils import attach_custom_error_code
 
 
 # TODO: There are more patterns where spark may throw an error.
@@ -229,9 +231,11 @@ def convert_spark_format_to_snowflake(
     timestamp_input_type: DataType | None = None,
 ):
     if spark_format in {"Y", "w", "W"}:
-        raise DateTimeException(
+        exception = DateTimeException(
             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter."
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT)
+        raise exception
     snowflake_format = ""
     i = 0
     n = len(spark_format)
@@ -299,28 +303,46 @@ def convert_spark_format_to_snowflake(
                     # Spark's 'a' would be at most 1 times
                     is_valid_a_pattern = spark_format[i : i + 2] != char * 2
                     if not is_valid_a_pattern:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "h" | "K" | "k" | "H" | "m" | "s" | "d":
                     # Spark's characters would be at most 2 times
                     is_valid_2_patterns = spark_format[i : i + 3] != char * 3
                     if not is_valid_2_patterns:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "D":
                     # Spark's 'D'' would be at most 3 times
                     is_valid_D_patterns = spark_format[i : i + 4] != char * 4
                     if not is_valid_D_patterns:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "V":
                     # Spark's 'V' for Zone ID requires 'VV'. A single 'V' is invalid.
                     is_valid_vv_pattern = spark_format[i : i + 2] == "VV"
                     if not is_valid_vv_pattern:
-                        raise DateTimeException("Pattern letter count must be 2: V")
+                        exception = DateTimeException(
+                            "Pattern letter count must be 2: V"
+                        )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "O":
                     # Spark's 'O' would be either 1 or 4.
                     is_valid_o_or_oooo_pattern = spark_format[i : i + 2] != "OO" or (
@@ -328,28 +350,52 @@ def convert_spark_format_to_snowflake(
                         and spark_format[i : i + 5] != "OOOOO"
                     )
                     if not is_valid_o_or_oooo_pattern:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             "Pattern letter count must be 1 or 4: O"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "q" | "Q" | "z" | "E":
                     # Spark's characters would be at most 4 times
                     is_valid_4_patterns = spark_format[i : i + 5] != char * 5
                     if not is_valid_4_patterns:
-                        raise DateTimeException(f"Too many pattern letters: {char}")
+                        exception = DateTimeException(
+                            f"Too many pattern letters: {char}"
+                        )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "x" | "X" | "Z":
                     # Spark's 'x' or 'X' or 'z' or 'Z' would be at most 5 times
                     is_valid_xz_pattern = spark_format[i : i + 6] != char * 6
                     if not is_valid_xz_pattern:
-                        raise DateTimeException(f"Too many pattern letters: {char}")
+                        exception = DateTimeException(
+                            f"Too many pattern letters: {char}"
+                        )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "y":
                     # Spark's 'y' would be at most 6 times
                     is_valid_y_pattern = spark_format[i : i + 7] != char * 7
                     if not is_valid_y_pattern:
-                        raise DateTimeException(
+                        exception = DateTimeException(
                             f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                        )
+                        raise exception
                 case "C" | "I":
-                    raise DateTimeException(f"Unknown pattern letter: {char}")
+                    exception = DateTimeException(f"Unknown pattern letter: {char}")
+                    attach_custom_error_code(
+                        exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                    )
+                    raise exception
 
             if (
                 spark_format[i : i + 5] in {"M" * 5, "L" * 5}
@@ -362,9 +408,13 @@ def convert_spark_format_to_snowflake(
                 or spark_format[i : i + 3] in {"kkk", "KKK"}
                 or spark_format[i : i + 10] == "SSSSSSSSSS"
             ):
-                raise DateTimeException(
+                exception = DateTimeException(
                     f"Fail to recognize '{spark_format}' pattern in the DateTimeFormatter."
                 )
+                attach_custom_error_code(
+                    exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                )
+                raise exception
 
             matched_pattern = False
 
@@ -375,7 +425,11 @@ def convert_spark_format_to_snowflake(
                         spark_key
                     ]
                     if isinstance(snowflake_equivalent, _UnsupportedSparkFormatPattern):
-                        raise DateTimeException(snowflake_equivalent.message)
+                        exception = DateTimeException(snowflake_equivalent.message)
+                        attach_custom_error_code(
+                            exception, ErrorCodes.UNSUPPORTED_OPERATION
+                        )
+                        raise exception
                     if snowflake_equivalent is not None:
                         snowflake_format += snowflake_equivalent
                     i += len(spark_key)
@@ -389,7 +443,11 @@ def convert_spark_format_to_snowflake(
                     isinstance(timestamp_input_type, StringType)
                     and char not in snowflake_time_format_separator
                 ):
-                    raise DateTimeException(f"Illegal pattern character: {char}")
+                    exception = DateTimeException(f"Illegal pattern character: {char}")
+                    attach_custom_error_code(
+                        exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT
+                    )
+                    raise exception
 
                 snowflake_format += f'"{char}"'
                 i += 1

snowflake/snowpark_connect/error/error_codes.py

@@ -0,0 +1,50 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+"""
+Error code constants for Snowpark Connect.
+
+This module defines custom error codes that can be attached to exceptions
+and included in gRPC error responses.
+"""
+
+
+class ErrorCodes:
+    """Constants for Snowpark Connect custom error codes."""
+
+    # 1000-1999: Startup related errors
+    MISSING_DATABASE = 1001
+    MISSING_SCHEMA = 1002
+    RESOURCE_INITIALIZATION_FAILED = 1003
+    TCP_PORT_ALREADY_IN_USE = 1004
+    INVALID_SPARK_CONNECT_URL = 1005
+    INVALID_STARTUP_INPUT = 1006
+    INVALID_STARTUP_OPERATION = 1007
+    STARTUP_CONNECTION_FAILED = 1008
+
+    # 2000-2999: Configuration related errors
+    INVALID_CONFIG_VALUE = 2001
+    CONFIG_CHANGE_NOT_ALLOWED = 2002
+    CONFIG_NOT_ENABLED = 2003
+
+    # 3000-3999: User code errors
+    INVALID_SQL_SYNTAX = 3001
+    TYPE_MISMATCH = 3002
+    INVALID_CAST = 3003
+    INVALID_FUNCTION_ARGUMENT = 3004
+    ARRAY_INDEX_OUT_OF_BOUNDS = 3005
+    DIVISION_BY_ZERO = 3006
+    INVALID_INPUT = 3007
+    INVALID_OPERATION = 3008
+    INSUFFICIENT_INPUT = 3009
+
+    # 4000-4999: What we don't support
+    UNSUPPORTED_OPERATION = 4001
+    UNSUPPORTED_TYPE = 4002
+
+    # 5000-5999: Internal errors
+    INTERNAL_ERROR = 5001
+    TABLE_NOT_FOUND = 5002
+    COLUMN_NOT_FOUND = 5003
+    AMBIGUOUS_COLUMN_NAME = 5004

snowflake/snowpark_connect/error/error_utils.py

@@ -12,11 +12,13 @@ https://github.com/apache/spark/blob/master/common/utils/src/main/resources/erro
 import json
 import pathlib
 import re
+import threading
 import traceback
 
 import jpype
 from google.protobuf import any_pb2
 from google.rpc import code_pb2, error_details_pb2, status_pb2
+from pyspark.errors import TempTableAlreadyExistsException
 from pyspark.errors.error_classes import ERROR_CLASSES_MAP
 from pyspark.errors.exceptions.base import (
     AnalysisException,
@@ -35,9 +37,12 @@ from snowflake.core.exceptions import NotFoundError
 
 from snowflake.connector.errors import ProgrammingError
 from snowflake.snowpark.exceptions import SnowparkClientException, SnowparkSQLException
-from snowflake.snowpark_connect.config import global_config
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
 from snowflake.snowpark_connect.error.error_mapping import ERROR_MAPPINGS_JSON
 
+# Thread-local storage for custom error codes when we can't attach them directly to exceptions
+_thread_local = threading.local()
+
 # The JSON string in error_mapping.py is a copy of https://github.com/apache/spark/blob/master/common/utils/src/main/resources/error/error-conditions.json.
 # The file doesn't have to be synced with spark latest main. Just update it when required.
 current_dir = pathlib.Path(__file__).parent.resolve()
@@ -54,8 +59,11 @@ SPARK_PYTHON_TO_JAVA_EXCEPTION = {
     SparkConnectGrpcException: "pyspark.errors.exceptions.connect.SparkConnectGrpcException",
     PythonException: "org.apache.spark.api.python.PythonException",
     UnsupportedOperationException: "java.lang.UnsupportedOperationException",
+    TempTableAlreadyExistsException: "org.apache.spark.sql.catalyst.analysis.TempTableAlreadyExistsException",
 }
 
+TABLE_OR_VIEW_NOT_FOUND_ERROR_CLASS = "TABLE_OR_VIEW_NOT_FOUND"
+
 WINDOW_FUNCTION_ANALYSIS_EXCEPTION_SQL_ERROR_CODE = {1005, 2303}
 ANALYSIS_EXCEPTION_SQL_ERROR_CODE = {
     904,
@@ -79,6 +87,23 @@ invalid_bit_pattern = re.compile(
     r"Invalid bit position: \d+ exceeds the bit (?:upper|lower) limit",
     re.IGNORECASE,
 )
+CREATE_SCHEMA_PATTERN = re.compile(r"create\s+schema", re.IGNORECASE)
+CREATE_TABLE_PATTERN = re.compile(r"create\s+table", re.IGNORECASE)
+
+
+def attach_custom_error_code(exception: Exception, custom_error_code: int) -> Exception:
+    """
+    Attach a custom error code to any exception instance.
+    This allows us to add custom error codes to existing PySpark exceptions.
+    """
+    if not hasattr(exception, "custom_error_code"):
+        try:
+            exception.custom_error_code = custom_error_code
+        except (AttributeError, TypeError):
+            # Some exception types (like Java exceptions) don't allow setting custom attributes
+            # Store the error code in thread-local storage for later retrieval
+            _thread_local.pending_error_code = custom_error_code
+    return exception
 
 
 def contains_udtf_select(sql_string):
@@ -100,20 +125,29 @@ def _get_converted_known_sql_or_custom_exception(
 
     # custom exception
     if "[snowpark_connect::invalid_array_index]" in msg:
-        return ArrayIndexOutOfBoundsException(
+        exception = ArrayIndexOutOfBoundsException(
             message='The index <indexValue> is out of bounds. The array has <arraySize> elements. Use the SQL function `get()` to tolerate accessing element at invalid index and return NULL instead. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error.'
         )
+        attach_custom_error_code(exception, ErrorCodes.ARRAY_INDEX_OUT_OF_BOUNDS)
+        return exception
     if "[snowpark_connect::invalid_index_of_zero]" in msg:
-        return SparkRuntimeException(
+        exception = SparkRuntimeException(
             message="[INVALID_INDEX_OF_ZERO] The index 0 is invalid. An index shall be either < 0 or > 0 (the first element has index 1)."
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+        return exception
     if "[snowpark_connect::invalid_index_of_zero_in_slice]" in msg:
-        return SparkRuntimeException(
+        exception = SparkRuntimeException(
             message="Unexpected value for start in function slice: SQL array indices start at 1."
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+        return exception
+
     invalid_bit = invalid_bit_pattern.search(msg)
     if invalid_bit:
-        return IllegalArgumentException(message=invalid_bit.group(0))
+        exception = IllegalArgumentException(message=invalid_bit.group(0))
+        attach_custom_error_code(exception, ErrorCodes.INVALID_FUNCTION_ARGUMENT)
+        return exception
     match = snowpark_connect_exception_pattern.search(
         ex.message if hasattr(ex, "message") else str(ex)
     )
@@ -125,71 +159,136 @@ def _get_converted_known_sql_or_custom_exception(
             if class_name
             else SparkConnectGrpcException
         )
-        return exception_class(message=message)
+        exception = exception_class(message=message)
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        return exception
 
     if "select with no columns" in msg and contains_udtf_select(query):
         # We try our best to detect if the SQL string contains a UDTF call and the output schema is empty.
-        return PythonException(message=f"[UDTF_RETURN_SCHEMA_MISMATCH] {ex.message}")
+        exception = PythonException(
+            message=f"[UDTF_RETURN_SCHEMA_MISMATCH] {ex.message}"
+        )
+        attach_custom_error_code(exception, ErrorCodes.TYPE_MISMATCH)
+        return exception
 
     # known sql exception
     if ex.sql_error_code not in (100038, 100037, 100035, 100357):
         return None
 
     if "(22018): numeric value" in msg:
-        return NumberFormatException(
+        exception = NumberFormatException(
             message='[CAST_INVALID_INPUT] Correct the value as per the syntax, or change its target type. Use `try_cast` to tolerate malformed input and return NULL instead. If necessary setting "spark.sql.ansi.enabled" to "false" may bypass this error.'
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_CAST)
+        return exception
     if "(22018): boolean value" in msg:
-        return SparkRuntimeException(
+        exception = SparkRuntimeException(
             message='[CAST_INVALID_INPUT] Correct the value as per the syntax, or change its target type. Use `try_cast` to tolerate malformed input and return NULL instead. If necessary setting "spark.sql.ansi.enabled" to "false" may bypass this error.'
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_CAST)
+        return exception
     if "(22007): timestamp" in msg:
-        return AnalysisException(
+        exception = AnalysisException(
             "[DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE] Data type mismatch"
         )
+        attach_custom_error_code(exception, ErrorCodes.TYPE_MISMATCH)
+        return exception
 
     if getattr(ex, "sql_error_code", None) == 100357:
         if re.search(init_multi_args_exception_pattern, msg):
-            return PythonException(
+            exception = PythonException(
                 message=f"[UDTF_EXEC_ERROR] User defined table function encountered an error in the init method {ex.message}"
             )
+            attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+            return exception
         if re.search(terminate_multi_args_exception_pattern, msg):
-            return PythonException(
+            exception = PythonException(
                 message=f"[UDTF_EXEC_ERROR] User defined table function encountered an error in the terminate method: {ex.message}"
             )
+            attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+            return exception
 
         if "failed to split string, provided pattern:" in msg:
-            return IllegalArgumentException(
+            exception = IllegalArgumentException(
                 message=f"Failed to split string using provided pattern. {ex.message}"
             )
+            attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+            return exception
 
         if "100357" in msg and "wrong tuple size for returned value" in msg:
-            return PythonException(
+            exception = PythonException(
                 message=f"[UDTF_RETURN_SCHEMA_MISMATCH] The number of columns in the result does not match the specified schema. {ex.message}"
             )
+            attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+            return exception
 
         if "100357 (p0000): python interpreter error:" in msg:
             if "in eval" in msg:
-                return PythonException(
+                exception = PythonException(
                     message=f"[UDTF_EXEC_ERROR] User defined table function encountered an error in the 'eval' method: error. {ex.message}"
                 )
+                attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+                return exception
 
             if "in terminate" in msg:
-                return PythonException(
+                exception = PythonException(
                     message=f"[UDTF_EXEC_ERROR] User defined table function encountered an error in the 'terminate' method: terminate error. {ex.message}"
                 )
+                attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+                return exception
 
             if "object is not iterable" in msg and contains_udtf_select(query):
-                return PythonException(
+                exception = PythonException(
                     message=f"[UDTF_RETURN_NOT_ITERABLE] {ex.message}"
                 )
+                attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+                return exception
 
-            return PythonException(message=f"{ex.message}")
+            exception = PythonException(message=f"{ex.message}")
+            attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+            return exception
 
     return None
 
 
+def _sanitize_custom_error_message(msg):
+    if "[snowpark_connect::unsupported_operation]" in msg:
+        return (
+            msg.replace("[snowpark_connect::unsupported_operation] ", ""),
+            ErrorCodes.UNSUPPORTED_OPERATION,
+        )
+    if "[snowpark_connect::internal_error]" in msg:
+        return (
+            msg.replace("[snowpark_connect::internal_error] ", ""),
+            ErrorCodes.INTERNAL_ERROR,
+        )
+    if "[snowpark_connect::invalid_operation]" in msg:
+        return (
+            msg.replace("[snowpark_connect::invalid_operation] ", ""),
+            ErrorCodes.INVALID_OPERATION,
+        )
+    if "[snowpark_connect::type_mismatch]" in msg:
+        return (
+            msg.replace("[snowpark_connect::type_mismatch] ", ""),
+            ErrorCodes.TYPE_MISMATCH,
+        )
+    if "[snowpark_connect::invalid_input]" in msg:
+        return (
+            msg.replace("[snowpark_connect::invalid_input] ", ""),
+            ErrorCodes.INVALID_INPUT,
+        )
+    if "[snowpark_connect::unsupported_type]" in msg:
+        return (
+            msg.replace("[snowpark_connect::unsupported_type] ", ""),
+            ErrorCodes.UNSUPPORTED_TYPE,
+        )
+    return msg, None
+
+
 def build_grpc_error_response(ex: Exception) -> status_pb2.Status:
+    # Lazy import to avoid circular dependency
+    from snowflake.snowpark_connect.config import global_config
+
     include_stack_trace = (
         global_config.get("spark.sql.pyspark.jvmStacktrace.enabled")
         if hasattr(global_config, "spark.sql.pyspark.jvmStacktrace.enabled")
@@ -203,6 +302,16 @@ def build_grpc_error_response(ex: Exception) -> status_pb2.Status:
         match ex:
             case SnowparkSQLException():
                 if ex.sql_error_code in ANALYSIS_EXCEPTION_SQL_ERROR_CODE:
+                    # Creation of schema that already exists
+                    if ex.sql_error_code == 2002 and "already exists" in str(ex):
+                        if CREATE_SCHEMA_PATTERN.search(ex.query):
+                            spark_java_classes.append(
+                                "org.apache.spark.sql.catalyst.analysis.NamespaceAlreadyExistsException"
+                            )
+                        elif CREATE_TABLE_PATTERN.search(ex.query):
+                            spark_java_classes.append(
+                                "org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException"
+                            )
                     # Data type mismatch, invalid window function
                     spark_java_classes.append("org.apache.spark.sql.AnalysisException")
                 elif ex.sql_error_code == 100051:
@@ -211,6 +320,7 @@ def build_grpc_error_response(ex: Exception) -> status_pb2.Status:
                         error_class="DIVIDE_BY_ZERO",
                         message_parameters={"config": '"spark.sql.ansi.enabled"'},
                     )
+                    attach_custom_error_code(ex, ErrorCodes.DIVISION_BY_ZERO)
                 elif ex.sql_error_code in (100096, 100040):
                     # Spark seems to want the Java base class instead of org.apache.spark.sql.SparkDateTimeException
                     # which is what should really be thrown
@@ -221,6 +331,9 @@ def build_grpc_error_response(ex: Exception) -> status_pb2.Status:
                     ex = spark_ex
                     spark_java_classes.append(SPARK_PYTHON_TO_JAVA_EXCEPTION[type(ex)])
                 elif ex.sql_error_code == 2043:
+                    spark_java_classes.append(
+                        "org.apache.spark.sql.catalyst.analysis.NoSuchDatabaseException"
+                    )
                     spark_java_classes.append("org.apache.spark.sql.AnalysisException")
                     message = f"does_not_exist: {str(ex)}"
                 else:
@@ -252,13 +365,23 @@ def build_grpc_error_response(ex: Exception) -> status_pb2.Status:
         )
     elif isinstance(ex, PySparkException):
         # pyspark exceptions thrown in sas layer
+
+        error_derived_java_class = []
+        if ex.error_class == TABLE_OR_VIEW_NOT_FOUND_ERROR_CLASS:
+            error_derived_java_class.append(
+                "org.apache.spark.sql.catalyst.analysis.NoSuchTableException"
+            )
+
         classes = type(ex).__mro__
         spark_java_classes = [
             SPARK_PYTHON_TO_JAVA_EXCEPTION[clazz]
             for clazz in classes
             if clazz in SPARK_PYTHON_TO_JAVA_EXCEPTION
         ]
-        metadata = {"classes": json.dumps(spark_java_classes)}
+
+        metadata = {
+            "classes": json.dumps(error_derived_java_class + spark_java_classes)
+        }
         if include_stack_trace:
             metadata["stackTrace"] = "".join(
                 traceback.TracebackException.from_exception(ex).format()
@@ -299,14 +422,40 @@ def build_grpc_error_response(ex: Exception) -> status_pb2.Status:
             domain="snowflake.sas",
         )
 
-    detail = any_pb2.Any()
-    detail.Pack(error_info)
-
     if message is None:
         message = str(ex)
 
+    custom_error_code = None
+
+    # attach error code using visa exception message
+    message, custom_error_code_from_msg = _sanitize_custom_error_message(message)
+
+    # Check if exception already has a custom error code, if not add INTERNAL_ERROR as default
+    if not hasattr(ex, "custom_error_code") or ex.custom_error_code is None:
+        attach_custom_error_code(
+            ex,
+            ErrorCodes.INTERNAL_ERROR
+            if custom_error_code_from_msg is None
+            else custom_error_code_from_msg,
+        )
+
+    # Get the custom error code from the exception or thread-local storage
+    custom_error_code = getattr(ex, "custom_error_code", None) or getattr(
+        _thread_local, "pending_error_code", None
+    )
+
+    # Clear thread-local storage after retrieving the error code
+    if hasattr(_thread_local, "pending_error_code"):
+        delattr(_thread_local, "pending_error_code")
+
+    separator = "==========================================="
+    error_code_added_message = f"\n{separator}\nSNOWPARK CONNECT ERROR CODE: {custom_error_code}\n{separator}\n{message}"
+
+    detail = any_pb2.Any()
+    detail.Pack(error_info)
+
     rich_status = status_pb2.Status(
-        code=code_pb2.INTERNAL, message=message, details=[detail]
+        code=code_pb2.INTERNAL, message=error_code_added_message, details=[detail]
     )
     return rich_status
 

snowflake/snowpark_connect/error/exceptions.py

@@ -2,27 +2,36 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
+from snowflake.snowpark_connect.error.error_codes import ErrorCodes
+
 
 class SnowparkConnectException(Exception):
     """Parent class to all SnowparkConnect related exceptions."""
 
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(self, *args, custom_error_code=None, **kwargs) -> None:
         super().__init__(*args, **kwargs)
+        self.custom_error_code = custom_error_code
 
 
 class MissingDatabase(SnowparkConnectException):
-    def __init__(self) -> None:
+    def __init__(self, custom_error_code=None) -> None:
         super().__init__(
             "No default database found in session",
+            custom_error_code=custom_error_code or ErrorCodes.MISSING_DATABASE,
         )
 
 
 class MissingSchema(SnowparkConnectException):
-    def __init__(self) -> None:
+    def __init__(self, custom_error_code=None) -> None:
         super().__init__(
             "No default schema found in session",
+            custom_error_code=custom_error_code or ErrorCodes.MISSING_SCHEMA,
         )
 
 
 class MaxRetryExceeded(SnowparkConnectException):
-    ...
+    def __init__(
+        self,
+        message="Maximum retry attempts exceeded",
+    ) -> None:
+        super().__init__(message)