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

snowflake/snowpark_connect/relation/read/jdbc_read_dbapi.py

@@ -38,6 +38,8 @@ from snowflake.snowpark.types import (
     TimeType,
     _NumericType,
 )
+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.relation.read.utils import (
     DATA_SOURCE_SQL_COMMENT,
     Connection,
@@ -147,9 +149,11 @@ class JdbcDataFrameReader(DataFrameReader):
                     or upper_bound is not None
                     or num_partitions is not None
                 ):
-                    raise ValueError(
+                    exception = ValueError(
                         "when column is not specified, lower_bound, upper_bound, num_partitions are expected to be None"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
                 if table is not None:
                     partitioned_queries = []
                     table_query = f"SELECT * FROM {table}"
@@ -160,24 +164,32 @@ class JdbcDataFrameReader(DataFrameReader):
                 elif query is not None:
                     partitioned_queries = [query]
                 else:
-                    raise ValueError("table or query is not specified")
+                    exception = ValueError("table or query is not specified")
+                    attach_custom_error_code(exception, ErrorCodes.INSUFFICIENT_INPUT)
+                    raise exception
             else:
                 if lower_bound is None or upper_bound is None or num_partitions is None:
-                    raise ValueError(
+                    exception = ValueError(
                         "when column is specified, lower_bound, upper_bound, num_partitions must be specified"
                     )
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
 
                 column_type = None
                 for field in struct_schema.fields:
                     if field.name.lower() == column.lower():
                         column_type = field.datatype
                 if column_type is None:
-                    raise ValueError("Column does not exist")
+                    exception = ValueError("Column does not exist")
+                    attach_custom_error_code(exception, ErrorCodes.COLUMN_NOT_FOUND)
+                    raise exception
 
                 if not isinstance(column_type, _NumericType) and not isinstance(
                     column_type, DateType
                 ):
-                    raise ValueError(f"unsupported type {column_type}")
+                    exception = ValueError(f"unsupported type {column_type}")
+                    attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+                    raise exception
                 spark_column_name = f'"{column}"'
                 partitioned_queries = self._generate_partition(
                     table,
@@ -240,7 +252,11 @@ class JdbcDataFrameReader(DataFrameReader):
                             )
                             query_thread_executor.shutdown(wait=False)
                             upload_thread_executor.shutdown(wait=False)
-                            raise future.result()
+                            exception = future.result()
+                            attach_custom_error_code(
+                                exception, ErrorCodes.INTERNAL_ERROR
+                            )
+                            raise exception
                         else:
                             path = future.result()
                             if not path:
@@ -266,7 +282,11 @@ class JdbcDataFrameReader(DataFrameReader):
                             )
                             query_thread_executor.shutdown(wait=False)
                             upload_thread_executor.shutdown(wait=False)
-                            raise f.result()
+                            exception = f.result()
+                            attach_custom_error_code(
+                                exception, ErrorCodes.INTERNAL_ERROR
+                            )
+                            raise exception
         finally:
             close_connection(conn)
 
@@ -283,7 +303,9 @@ class JdbcDataFrameReader(DataFrameReader):
         elif query is not None:
             sql = f"SELECT * FROM ({query}) WHERE 1=0"
         else:
-            raise ValueError("table or query is not specified")
+            exception = ValueError("table or query is not specified")
+            attach_custom_error_code(exception, ErrorCodes.INSUFFICIENT_INPUT)
+            raise exception
 
         cursor = conn.cursor()
         cursor.execute(sql)
@@ -301,7 +323,11 @@ class JdbcDataFrameReader(DataFrameReader):
             dt = parser.parse(value)
             return int(dt.replace(tzinfo=pytz.UTC).timestamp())
         else:
-            raise TypeError(f"unsupported column type for partition: {column_type}")
+            exception = TypeError(
+                f"unsupported column type for partition: {column_type}"
+            )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+            raise exception
 
     # this function is only used in data source API for SQL server
     def _to_external_value(self, value: Union[int, str, float], column_type: DataType):
@@ -311,7 +337,11 @@ class JdbcDataFrameReader(DataFrameReader):
             # TODO: SNOW-1909315: support timezone
             return datetime.datetime.fromtimestamp(value, tz=pytz.UTC)
         else:
-            raise TypeError(f"unsupported column type for partition: {column_type}")
+            exception = TypeError(
+                f"unsupported column type for partition: {column_type}"
+            )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+            raise exception
 
     def _to_snowpark_type(self, schema: Tuple[tuple]) -> StructType:
         fields = []
@@ -339,7 +369,9 @@ class JdbcDataFrameReader(DataFrameReader):
                 case jaydebeapi.BINARY:
                     field = StructField(name, BinaryType(), is_nullable)
                 case _:
-                    raise ValueError(f"unsupported type: {dbapi_type}")
+                    exception = ValueError(f"unsupported type: {dbapi_type}")
+                    attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+                    raise exception
 
             fields.append(field)
         return StructType(fields)
@@ -359,7 +391,9 @@ class JdbcDataFrameReader(DataFrameReader):
         processed_lower_bound = self._to_internal_value(lower_bound, column_type)
         processed_upper_bound = self._to_internal_value(upper_bound, column_type)
         if processed_lower_bound > processed_upper_bound:
-            raise ValueError("lower_bound cannot be greater than upper_bound")
+            exception = ValueError("lower_bound cannot be greater than upper_bound")
+            attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+            raise exception
 
         if processed_lower_bound == processed_upper_bound or num_partitions <= 1:
             return [select_query]
@@ -665,4 +699,6 @@ def get_jdbc_dialect(url: str) -> JdbcDialect:
     for jdbc_dialect in jdbc_dialects:
         if jdbc_dialect.can_handle(url):
             return jdbc_dialect
-    raise ValueError(f"Unsupported JDBC datasource: {url}")
+    exception = ValueError(f"Unsupported JDBC datasource: {url}")
+    attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+    raise exception

snowflake/snowpark_connect/relation/read/map_read.py

@@ -15,6 +15,8 @@ from snowflake import snowpark
 from snowflake.snowpark.types import StructType
 from snowflake.snowpark_connect.config import global_config
 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.relation.io_utils import (
     convert_file_prefix_path,
     get_compression_for_source_and_options,
@@ -159,12 +161,20 @@ def map_read(
                                 options[DBTABLE_OPTION], session, rel.common.plan_id
                             )
                     case other:
-                        raise SnowparkConnectNotImplementedError(
+                        exception = SnowparkConnectNotImplementedError(
                             f"UNSUPPORTED FORMAT {other} WITH NO PATH"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.UNSUPPORTED_OPERATION
+                        )
+                        raise exception
         case other:
             # TODO: Empty data source
-            raise SnowparkConnectNotImplementedError(f"Unsupported read type: {other}")
+            exception = SnowparkConnectNotImplementedError(
+                f"Unsupported read type: {other}"
+            )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
     return df_cache_map_put_if_absent(
         (get_session_id(), rel.common.plan_id),
@@ -274,9 +284,11 @@ def _read_file(
 
             return map_read_text(rel, schema, session, paths)
         case _:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unsupported format: {read_format}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
 
 
 def _skip_upload(path: str, read_format: str):

snowflake/snowpark_connect/relation/read/map_read_csv.py

@@ -12,6 +12,8 @@ from snowflake import snowpark
 from snowflake.snowpark.dataframe_reader import DataFrameReader
 from snowflake.snowpark.types import StringType, StructField, StructType
 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.relation.read.map_read import CsvReaderConfig
 from snowflake.snowpark_connect.relation.read.metadata_utils import (
     add_filename_metadata_to_reader,
@@ -43,9 +45,11 @@ def map_read_csv(
 
     if rel.read.is_streaming is True:
         # TODO: Structured streaming implementation.
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "Streaming is not supported for CSV files."
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
     else:
         snowpark_options = options.convert_to_snowpark_args()
         parse_header = snowpark_options.get("PARSE_HEADER", False)
@@ -188,14 +192,18 @@ def read_data(
 
     if schema is not None:
         if len(schema.fields) != len(non_metadata_fields):
-            raise Exception(f"csv load from {filename} failed.")
+            exception = Exception(f"csv load from {filename} failed.")
+            attach_custom_error_code(exception, ErrorCodes.INVALID_CAST)
+            raise exception
         if raw_options.get("enforceSchema", "True").lower() == "false":
             for i in range(len(schema.fields)):
                 if (
                     schema.fields[i].name != non_metadata_fields[i].name
                     and f'"{schema.fields[i].name}"' != non_metadata_fields[i].name
                 ):
-                    raise Exception("CSV header does not conform to the schema")
+                    exception = Exception("CSV header does not conform to the schema")
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+                    raise exception
         return df
 
     headers = get_header_names(

snowflake/snowpark_connect/relation/read/map_read_jdbc.py

@@ -9,6 +9,8 @@ import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 from snowflake import snowpark
 from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
 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.relation.read.jdbc_read_dbapi import JdbcDataFrameReader
 from snowflake.snowpark_connect.relation.read.utils import (
     Connection,
@@ -28,7 +30,9 @@ def create_connection(jdbc_options: dict[str, str]) -> Connection:
         return jaydebeapi.connect(driver, url, jdbc_options)
     except Exception as e:
         jpype.detachThreadFromJVM()
-        raise Exception(f"Error connecting JDBC datasource: {e}")
+        exception = Exception(f"Error connecting JDBC datasource: {e}")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception
 
 
 def close_connection(conn: Connection) -> None:
@@ -70,17 +74,23 @@ def map_read_jdbc(
         dbtable = None
 
     if not dbtable and not query:
-        raise ValueError("Include dbtable or query is required option")
+        exception = ValueError("Include dbtable or query is required option")
+        attach_custom_error_code(exception, ErrorCodes.INSUFFICIENT_INPUT)
+        raise exception
 
     if query is not None and dbtable is not None:
-        raise ValueError(
+        exception = ValueError(
             "Not allowed to specify dbtable and query options at the same time"
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+        raise exception
 
     if query is not None and partition_column is not None:
-        raise ValueError(
+        exception = ValueError(
             "Not allowed to specify partitionColumn and query options at the same time"
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+        raise exception
 
     try:
         df = JdbcDataFrameReader(session, jdbc_options).jdbc_read_dbapi(
@@ -105,4 +115,6 @@ def map_read_jdbc(
             snowpark_column_types=[f.datatype for f in df.schema.fields],
         )
     except Exception as e:
-        raise Exception(f"Error accessing JDBC datasource for read: {e}")
+        exception = Exception(f"Error accessing JDBC datasource for read: {e}")
+        attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+        raise exception

snowflake/snowpark_connect/relation/read/map_read_json.py

@@ -28,6 +28,8 @@ from snowflake.snowpark.types import (
     TimestampType,
 )
 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.relation.read.map_read import JsonReaderConfig
 from snowflake.snowpark_connect.relation.read.metadata_utils import (
     add_filename_metadata_to_reader,
@@ -64,9 +66,11 @@ def map_read_json(
 
     if rel.read.is_streaming is True:
         # TODO: Structured streaming implementation.
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "Streaming is not supported for JSON files."
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
     else:
         snowpark_options = options.convert_to_snowpark_args()
         raw_options = rel.read.data_source.options
@@ -363,9 +367,11 @@ def construct_row_by_schema(
                     content.get(col_name, None), sf.datatype, snowpark_options
                 )
         else:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"JSON construct {str(content)} to StructType failed"
             )
+            attach_custom_error_code(exception, ErrorCodes.TYPE_MISMATCH)
+            raise exception
         return result
     elif isinstance(schema, ArrayType):
         result = []

snowflake/snowpark_connect/relation/read/map_read_parquet.py

@@ -22,6 +22,8 @@ from snowflake.snowpark._internal.analyzer.analyzer_utils import (
 from snowflake.snowpark.column import METADATA_FILENAME
 from snowflake.snowpark.types import DataType, DoubleType, IntegerType, StringType
 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.relation.read.metadata_utils import (
     add_filename_metadata_to_reader,
 )
@@ -44,9 +46,11 @@ def map_read_parquet(
     """Read a Parquet file into a Snowpark DataFrame."""
 
     if rel.read.is_streaming is True:
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "Streaming is not supported for Parquet files."
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     snowpark_options = options.convert_to_snowpark_args()
     raw_options = rel.read.data_source.options
@@ -155,10 +159,14 @@ def _discover_partition_columns(
                     if i not in dir_level_to_column_name:
                         dir_level_to_column_name[i] = key
                     elif dir_level_to_column_name[i] != key:
-                        raise ValueError(
+                        exception = ValueError(
                             f"Conflicting partition column names detected: '{dir_level_to_column_name[i]}' and '{key}' "
                             f"at the same directory level"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_OPERATION
+                        )
+                        raise exception
 
                     partition_columns_values[key].add(value)
 
@@ -166,10 +174,12 @@ def _discover_partition_columns(
     for level in sorted(dir_level_to_column_name.keys()):
         col_name = dir_level_to_column_name[level]
         if col_name in seen_columns:
-            raise ValueError(
+            exception = ValueError(
                 f"Found partition column '{col_name}' at multiple directory levels. "
                 f"A partition column can only appear at a single level."
             )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+            raise exception
         seen_columns.add(col_name)
 
     ordered_columns = [

snowflake/snowpark_connect/relation/read/map_read_socket.py

@@ -9,6 +9,8 @@ import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 from snowflake import snowpark
 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.telemetry import (
     SnowparkConnectNotImplementedError,
 )
@@ -30,7 +32,9 @@ def map_read_socket(
         host = options.get("host", None)
         port = options.get("port", None)
         if not host or not port:
-            raise ValueError("Host and port must be provided in options.")
+            exception = ValueError("Host and port must be provided in options.")
+            attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+            raise exception
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
             try:
                 s.connect((host, int(port)))
@@ -56,8 +60,12 @@ def map_read_socket(
                     snowpark_column_names=[snowpark_cname],
                 )
             except OSError as e:
-                raise Exception(f"Error connecting to {host}:{port} - {e}")
+                exception = Exception(f"Error connecting to {host}:{port} - {e}")
+                attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+                raise exception
     else:
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "Socket reads are only supported in streaming mode."
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception

snowflake/snowpark_connect/relation/read/map_read_table.py

@@ -16,8 +16,11 @@ from snowflake.snowpark_connect.column_name_handler import (
     ColumnNameMap,
     make_column_names_snowpark_compatible,
 )
+from snowflake.snowpark_connect.column_qualifier import ColumnQualifier
 from snowflake.snowpark_connect.config import auto_uppercase_non_column_identifiers
 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.relation.read.utils import (
     rename_columns_as_snowflake_standard,
 )
@@ -56,7 +59,7 @@ def post_process_df(
             spark_column_names=true_names,
             snowpark_column_names=snowpark_column_names,
             snowpark_column_types=[f.datatype for f in df.schema.fields],
-            column_qualifiers=[name_parts] * len(true_names)
+            column_qualifiers=[{ColumnQualifier(tuple(name_parts))} for _ in true_names]
             if source_table_name
             else None,
         )
@@ -64,9 +67,11 @@ def post_process_df(
         # Check if this is a table/view not found error
         # Snowflake error codes: 002003 (42S02) - Object does not exist or not authorized
         if hasattr(e, "sql_error_code") and e.sql_error_code == 2003:
-            raise AnalysisException(
+            exception = AnalysisException(
                 f"[TABLE_OR_VIEW_NOT_FOUND] The table or view cannot be found. {source_table_name}"
-            ) from None  # Suppress original exception to reduce message size
+            )
+            attach_custom_error_code(exception, ErrorCodes.INTERNAL_ERROR)
+            raise exception from None  # Suppress original exception to reduce message size
         # Re-raise if it's not a table not found error
         raise
 
@@ -90,8 +95,10 @@ def _get_temporary_view(
         spark_column_names=temp_view.column_map.get_spark_columns(),
         snowpark_column_names=snowpark_column_names,
         column_metadata=temp_view.column_map.column_metadata,
-        column_qualifiers=[split_fully_qualified_spark_name(table_name)]
-        * len(temp_view.column_map.get_spark_columns()),
+        column_qualifiers=[
+            {ColumnQualifier(tuple(split_fully_qualified_spark_name(table_name)))}
+            for _ in range(len(temp_view.column_map.get_spark_columns()))
+        ],
         parent_column_name_map=temp_view.column_map.get_parent_column_name_map(),
     )
 
@@ -118,9 +125,11 @@ def get_table_from_name(
 
     # Verify if recursive view read is not attempted
     if table_name in get_processed_views():
-        raise AnalysisException(
+        exception = AnalysisException(
             f"[RECURSIVE_VIEW] Recursive view `{table_name}` detected (cycle: `{table_name}` -> `{table_name}`)"
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+        raise exception
 
     snowpark_name = ".".join(
         quote_name_without_upper_casing(part)
@@ -159,10 +168,14 @@ def map_read_table(
         and rel.read.data_source.format.lower() == "iceberg"
     ):
         if len(rel.read.data_source.paths) != 1:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 f"Unexpected paths: {rel.read.data_source.paths}"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
         table_identifier = rel.read.data_source.paths[0]
     else:
-        raise ValueError("The relation must have a table identifier.")
+        exception = ValueError("The relation must have a table identifier.")
+        attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+        raise exception
     return get_table_from_name(table_identifier, session, rel.common.plan_id)

snowflake/snowpark_connect/relation/read/map_read_text.py

@@ -8,6 +8,8 @@ import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 from snowflake import snowpark
 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.relation.read.utils import (
     get_spark_column_names_from_snowpark_columns,
     rename_columns_as_snowflake_standard,
@@ -82,9 +84,11 @@ def map_read_text(
     """
     if rel.read.is_streaming is True:
         # TODO: Structured streaming implementation.
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "Streaming is not supported for CSV files."
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     df = read_text(paths[0], schema, session, rel.read.data_source.options)
     if len(paths) > 1:

snowflake/snowpark_connect/relation/read/metadata_utils.py

@@ -16,6 +16,8 @@ from snowflake.snowpark.column import METADATA_FILENAME
 from snowflake.snowpark.functions import col
 from snowflake.snowpark.types import StructField
 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
 
 # Constant for the metadata filename column name
 METADATA_FILENAME_COLUMN = "METADATA$FILENAME"
@@ -129,9 +131,11 @@ def filter_metadata_columns(
     if len(non_metadata_columns) == 0:
         # DataFrame contains only metadata columns (METADATA$FILENAME), no actual data columns remaining.
         # We don't have a way to return an empty dataframe.
-        raise AnalysisException(
+        exception = AnalysisException(
             "[DATAFRAME_MISSING_DATA_COLUMNS] Cannot perform operation on DataFrame that contains no data columns."
         )
+        attach_custom_error_code(exception, ErrorCodes.INVALID_OPERATION)
+        raise exception
 
     filtered_df = result_df.select([col(name) for name in non_metadata_columns])
 

snowflake/snowpark_connect/relation/stage_locator.py

@@ -11,6 +11,8 @@ from s3fs.core import S3FileSystem
 from snowflake import snowpark
 from snowflake.snowpark.session import Session
 from snowflake.snowpark_connect.config import sessions_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.relation.io_utils import (
     get_cloud_from_url,
     parse_azure_url,
@@ -42,9 +44,11 @@ def get_paths_from_stage(
                 rewrite_paths.append(f"{stage_name}/{path}")
             paths = rewrite_paths
         case "gcp":
-            raise AnalysisException(
+            exception = AnalysisException(
                 "You must configure an integration for Google Cloud Storage to perform I/O operations rather than accessing the URL directly. Reference: https://docs.snowflake.com/en/user-guide/data-load-gcs-config"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
         case _:
             filesystem, parsed_path = url_to_fs(paths[0])
             if isinstance(filesystem, S3FileSystem):  # aws

snowflake/snowpark_connect/relation/write/jdbc_write_dbapi.py

@@ -11,6 +11,8 @@ from snowflake import snowpark
 from snowflake.snowpark import DataFrameWriter
 from snowflake.snowpark.dataframe import DataFrame
 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.relation.read import jdbc_read_dbapi
 from snowflake.snowpark_connect.relation.read.jdbc_read_dbapi import JdbcDialect
 from snowflake.snowpark_connect.relation.read.utils import Connection
@@ -65,9 +67,13 @@ class JdbcDataFrameWriter(DataFrameWriter):
                         self._create_table(conn, table, container, jdbc_dialect)
                 case "errorifexists":
                     if table_exist:
-                        raise ValueError(
+                        exception = ValueError(
                             "table is already exist and write mode is ERROR_IF_EXISTS"
                         )
+                        attach_custom_error_code(
+                            exception, ErrorCodes.INVALID_OPERATION
+                        )
+                        raise exception
                     else:
                         self._create_table(conn, table, container, jdbc_dialect)
                 case "overwrite":
@@ -82,7 +88,9 @@ class JdbcDataFrameWriter(DataFrameWriter):
                     else:
                         self._create_table(conn, table, container, jdbc_dialect)
                 case _:
-                    raise ValueError(f"Invalid write mode value{write_mode}")
+                    exception = ValueError(f"Invalid write mode value{write_mode}")
+                    attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+                    raise exception
 
             task_insert_into_data_source_with_retry(
                 input_df,
@@ -141,6 +149,7 @@ class JdbcDataFrameWriter(DataFrameWriter):
                 cursor.execute(sql)
         except Exception as e:
             logger.error(f"failed to drop table {table} from the data source {e}")
+            attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
             raise e
 
     def _create_table(
@@ -189,6 +198,7 @@ class JdbcDataFrameWriter(DataFrameWriter):
                 cursor.execute(sql)
         except Exception as e:
             logger.error(f"failed to create a table {table} from the data source {e}")
+            attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
             raise e
 
 
@@ -218,6 +228,7 @@ def _task_insert_into_data_source(
     except Exception as e:
         logger.debug(f"failed to insert into data source  {e}")
         conn.rollback()
+        attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
         raise e
     finally:
         cursor.close()
@@ -274,6 +285,7 @@ def task_insert_into_data_source_with_retry(
             )
     except Exception as e:
         logger.debug(f"failed to insert into data source  {e}")
+        attach_custom_error_code(e, ErrorCodes.INTERNAL_ERROR)
         raise e
     finally:
         close_connection(conn)
@@ -339,4 +351,8 @@ def convert_sp_to_sql_type(
                 case _:
                     return "TIMESTAMP"
         case _:
-            raise TypeError(f"Unsupported data type: {datatype.__class__.__name__}")
+            exception = TypeError(
+                f"Unsupported data type: {datatype.__class__.__name__}"
+            )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_TYPE)
+            raise exception