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

snowflake/snowpark_connect/relation/catalogs/abstract_spark_catalog.py

@@ -11,6 +11,8 @@ import pyspark.sql.connect.proto.types_pb2 as types_proto
 
 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.error.exceptions import MissingDatabase, MissingSchema
 from snowflake.snowpark_connect.utils.identifiers import (
     split_fully_qualified_spark_name,
@@ -39,28 +41,40 @@ class AbstractSparkCatalog(ABC):
         description: str,
         **options: typing.Any,
     ) -> DataFrameContainer:
-        raise SnowparkConnectNotImplementedError("createTable is not implemented")
+        exception = SnowparkConnectNotImplementedError("createTable is not implemented")
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     @abstractmethod
     def listDatabases(
         self,
         pattern: str | None = None,
     ) -> pandas.DataFrame:
-        raise SnowparkConnectNotImplementedError("listDatabases is not implemented")
+        exception = SnowparkConnectNotImplementedError(
+            "listDatabases is not implemented"
+        )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     @abstractmethod
     def getDatabase(
         self,
         spark_dbName: str,
     ) -> pandas.DataFrame:
-        raise SnowparkConnectNotImplementedError("getDatabase is not implemented")
+        exception = SnowparkConnectNotImplementedError("getDatabase is not implemented")
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     @abstractmethod
     def databaseExists(
         self,
         spark_dbName: str,
     ) -> pandas.DataFrame:
-        raise SnowparkConnectNotImplementedError("databaseExists is not implemented")
+        exception = SnowparkConnectNotImplementedError(
+            "databaseExists is not implemented"
+        )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     @abstractmethod
     def listTables(
@@ -68,14 +82,18 @@ class AbstractSparkCatalog(ABC):
         spark_dbName: str | None = None,
         pattern: str | None = None,
     ) -> pandas.DataFrame:
-        raise SnowparkConnectNotImplementedError("listTables is not implemented")
+        exception = SnowparkConnectNotImplementedError("listTables is not implemented")
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     @abstractmethod
     def getTable(
         self,
         spark_tableName: str,
     ) -> pandas.DataFrame:
-        raise SnowparkConnectNotImplementedError("getTable is not implemented")
+        exception = SnowparkConnectNotImplementedError("getTable is not implemented")
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     @abstractmethod
     def tableExists(
@@ -83,7 +101,9 @@ class AbstractSparkCatalog(ABC):
         spark_tableName: str,
         spark_dbName: str | None,
     ) -> pandas.DataFrame:
-        raise SnowparkConnectNotImplementedError("tableExists is not implemented")
+        exception = SnowparkConnectNotImplementedError("tableExists is not implemented")
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     @abstractmethod
     def listColumns(
@@ -91,36 +111,50 @@ class AbstractSparkCatalog(ABC):
         spark_tableName: str,
         spark_dbName: str | None = None,
     ) -> pandas.DataFrame:
-        raise SnowparkConnectNotImplementedError("listColumns is not implemented")
+        exception = SnowparkConnectNotImplementedError("listColumns is not implemented")
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     @abstractmethod
     def currentDatabase(self) -> pandas.DataFrame:
-        raise SnowparkConnectNotImplementedError("currentDatabase is not implemented")
+        exception = SnowparkConnectNotImplementedError(
+            "currentDatabase is not implemented"
+        )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     @abstractmethod
     def setCurrentDatabase(
         self,
         spark_dbName: str,
     ) -> pandas.DataFrame:
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "setCurrentDatabase is not implemented"
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     @abstractmethod
     def dropGlobalTempView(
         self,
         spark_view_name: str,
     ) -> DataFrameContainer:
-        raise SnowparkConnectNotImplementedError(
+        exception = SnowparkConnectNotImplementedError(
             "dropGlobalTempView is not implemented"
         )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     @abstractmethod
     def dropTempView(
         self,
         spark_view_name: str,
     ) -> DataFrameContainer:
-        raise SnowparkConnectNotImplementedError("dropTempView is not implemented")
+        exception = SnowparkConnectNotImplementedError(
+            "dropTempView is not implemented"
+        )
+        attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+        raise exception
 
     def cacheTable(
         self,
@@ -135,9 +169,11 @@ class AbstractSparkCatalog(ABC):
             spark_tableName
         )
         if catalog is not None and self != catalog:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 "Calling into another catalog is not currently supported"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
         if sf_database is None:
             sf_database = _get_current_snowflake_database()
         if sf_schema is None:
@@ -168,9 +204,11 @@ class AbstractSparkCatalog(ABC):
             spark_tableName
         )
         if catalog is not None and self != catalog:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 "Calling into another catalog is not currently supported"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
         if sf_database is None:
             sf_database = _get_current_snowflake_database()
         if sf_schema is None:
@@ -194,9 +232,11 @@ class AbstractSparkCatalog(ABC):
             spark_tableName
         )
         if catalog is not None and self != catalog:
-            raise SnowparkConnectNotImplementedError(
+            exception = SnowparkConnectNotImplementedError(
                 "Calling into another catalog is not currently supported"
             )
+            attach_custom_error_code(exception, ErrorCodes.UNSUPPORTED_OPERATION)
+            raise exception
         if sf_database is None:
             sf_database = _get_current_snowflake_database()
         if sf_schema is None:
@@ -249,7 +289,11 @@ def _process_multi_layer_database(
             else:
                 return None, c, d
         case _:
-            raise ValueError(f"Unexpected database identifier format: {spark_mli}")
+            exception = ValueError(
+                f"Unexpected database identifier format: {spark_mli}"
+            )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+            raise exception
 
 
 def _process_multi_layer_identifier(
@@ -283,5 +327,9 @@ def _process_multi_layer_identifier(
             else:
                 snowflake_database, snowflake_schema, snowflake_obj = d, s, t
         case _:
-            raise ValueError(f"Unexpected table/view identifier format: {spark_mli}")
+            exception = ValueError(
+                f"Unexpected table/view identifier format: {spark_mli}"
+            )
+            attach_custom_error_code(exception, ErrorCodes.INVALID_INPUT)
+            raise exception
     return spark_catalog, snowflake_database, snowflake_schema, snowflake_obj