Revert "SNOW-2267482 replacing SnowAPI calls with SQL calls (#3942)"

This reverts commit 1dec439.

Author: sfc-gh-aling

src/snowflake/snowpark/catalog.py

@@ -2,19 +2,15 @@
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
 
-from ctypes import ArgumentError
 import re
-from typing import List, Optional, Union, TYPE_CHECKING
-
-from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
-from snowflake.snowpark.exceptions import SnowparkSQLException, NotFoundError
+from typing import List, Optional, Union
 
 try:
+    from snowflake.core import Root  # type: ignore
     from snowflake.core.database import Database  # type: ignore
-    from snowflake.core.database._generated.models import Database as ModelDatabase  # type: ignore
+    from snowflake.core.exceptions import NotFoundError
     from snowflake.core.procedure import Procedure
     from snowflake.core.schema import Schema  # type: ignore
-    from snowflake.core.schema._generated.models import Schema as ModelSchema  # type: ignore
     from snowflake.core.table import Table, TableColumn
     from snowflake.core.user_defined_function import UserDefinedFunction
     from snowflake.core.view import View
@@ -23,30 +19,27 @@
         "Missing optional dependency: 'snowflake.core'."
     ) from e  # pragma: no cover
 
-from snowflake.snowpark._internal.type_utils import (
-    convert_sp_to_sf_type,
-    type_string_to_type_object,
-)
+
+import snowflake.snowpark
+from snowflake.snowpark._internal.type_utils import convert_sp_to_sf_type
 from snowflake.snowpark.functions import lit, parse_json
 from snowflake.snowpark.types import DataType
 
-if TYPE_CHECKING:
-    from snowflake.snowpark.session import Session
-
 
 class Catalog:
     """The Catalog class provides methods to interact with and manage the Snowflake objects.
     It allows users to list, get, and drop various database objects such as databases, schemas, tables,
     views, functions, etc.
     """
 
-    def __init__(self, session: "Session") -> None:
+    def __init__(self, session: "snowflake.snowpark.session.Session") -> None:  # type: ignore
         self._session = session
+        self._root = Root(session)
         self._python_regex_udf = None
 
     def _parse_database(
         self,
-        database: object,
+        database: Optional[Union[str, Database]],
         model_obj: Optional[
             Union[str, Schema, Table, View, Procedure, UserDefinedFunction]
         ] = None,
@@ -73,7 +66,7 @@ def _parse_database(
 
     def _parse_schema(
         self,
-        schema: object,
+        schema: Optional[Union[str, Schema]],
         model_obj: Optional[
             Union[str, Table, View, Procedure, UserDefinedFunction]
         ] = None,
@@ -173,28 +166,11 @@ def list_databases(
             pattern: the python regex pattern of name to match. Defaults to None.
             like: the sql style pattern for name to match. Default to None.
         """
-        like_str = f"LIKE '{like}'" if like else ""
-        df = self._session.sql(f"SHOW AS RESOURCE DATABASES {like_str}")
+        iter = self._root.databases.iter(like=like)
         if pattern:
-            # initialize udf
-            self._initialize_regex_udf()
-            assert self._python_regex_udf is not None  # pyright
-
-            # The result of SHOW AS RESOURCE query is a json string which contains
-            # key 'name' to store the name of the object. We parse json for the returned
-            # result and apply the filter on name.
-            df = df.filter(
-                self._python_regex_udf(
-                    lit(pattern), parse_json('"As Resource"')["name"]
-                )
-            )
+            iter = filter(lambda x: re.match(pattern, x.name), iter)
 
-        return list(
-            map(
-                lambda row: Database._from_model(ModelDatabase.from_json(str(row[0]))),
-                df.collect(),
-            )
-        )
+        return list(iter)
 
     def list_schemas(
         self,
@@ -212,28 +188,10 @@ def list_schemas(
             like: the sql style pattern for name to match. Default to None.
         """
         db_name = self._parse_database(database)
-        like_str = f"LIKE '{like}'" if like else ""
-        df = self._session.sql(f"SHOW AS RESOURCE SCHEMAS {like_str} IN {db_name}")
+        iter = self._root.databases[db_name].schemas.iter(like=like)
         if pattern:
-            # initialize udf
-            self._initialize_regex_udf()
-            assert self._python_regex_udf is not None  # pyright
-
-            # The result of SHOW AS RESOURCE query is a json string which contains
-            # key 'name' to store the name of the object. We parse json for the returned
-            # result and apply the filter on name.
-            df = df.filter(
-                self._python_regex_udf(
-                    lit(pattern), parse_json('"As Resource"')["name"]
-                )
-            )
-
-        return list(
-            map(
-                lambda row: Schema._from_model(ModelSchema.from_json(str(row[0]))),
-                df.collect(),
-            )
-        )
+            iter = filter(lambda x: re.match(pattern, x.name), iter)
+        return list(iter)
 
     def list_tables(
         self,
@@ -371,27 +329,14 @@ def get_current_schema(self) -> Optional[str]:
 
     def get_database(self, database: str) -> Database:
         """Name of the database to get"""
-        try:
-            return self.list_databases(like=unquote_if_quoted(database))[0]
-        except IndexError:
-            raise NotFoundError(f"Database with name {database} could not be found")
+        return self._root.databases[database].fetch()
 
     def get_schema(
         self, schema: str, *, database: Optional[Union[str, Database]] = None
     ) -> Schema:
         """Name of the schema to get."""
         db_name = self._parse_database(database)
-        try:
-            return self.list_schemas(database=db_name, like=unquote_if_quoted(schema))[
-                0
-            ]
-        except (
-            IndexError,  # schema with this name doesn't exist
-            SnowparkSQLException,  # database in which we are looking doesn't exist
-        ):
-            raise NotFoundError(
-                f"Schema with name {schema} could not be found in database '{db_name}'"
-            )
+        return self._root.databases[db_name].schemas[schema].fetch()
 
     def get_table(
         self,
@@ -410,16 +355,12 @@ def get_table(
         """
         db_name = self._parse_database(database)
         schema_name = self._parse_schema(schema)
-        try:
-            return self.listTables(
-                database=db_name,
-                schema=schema_name,
-                like=unquote_if_quoted(table_name),
-            )[0]
-        except IndexError:
-            raise NotFoundError(
-                f"Table with name {table_name} could not be found in schema '{db_name}.{schema_name}'"
-            )
+        return (
+            self._root.databases[db_name]
+            .schemas[schema_name]
+            .tables[table_name]
+            .fetch()
+        )
 
     def get_view(
         self,
@@ -438,16 +379,9 @@ def get_view(
         """
         db_name = self._parse_database(database)
         schema_name = self._parse_schema(schema)
-        try:
-            return self.list_views(
-                database=db_name,
-                schema=schema_name,
-                like=unquote_if_quoted(view_name),
-            )[0]
-        except IndexError:
-            raise NotFoundError(
-                f"View with name {view_name} could not be found in schema '{db_name}.{schema_name}'"
-            )
+        return (
+            self._root.databases[db_name].schemas[schema_name].views[view_name].fetch()
+        )
 
     def get_procedure(
         self,
@@ -469,19 +403,12 @@ def get_procedure(
         db_name = self._parse_database(database)
         schema_name = self._parse_schema(schema)
         procedure_id = self._parse_function_or_procedure(procedure_name, arg_types)
-
-        try:
-            procedures = self._session.sql(
-                f"DESCRIBE AS RESOURCE PROCEDURE {db_name}.{schema_name}.{procedure_id}"
-            ).collect()
-            return Procedure.from_json(str(procedures[0][0]))
-        except (
-            IndexError,  # when sql returned no results
-            SnowparkSQLException,  # when database, or schema doesn't exist
-        ):
-            raise NotFoundError(
-                f"Procedure with name {procedure_name} and arguments {arg_types} could not be found in schema '{db_name}.{schema_name}'"
-            )
+        return (
+            self._root.databases[db_name]
+            .schemas[schema_name]
+            .procedures[procedure_id]
+            .fetch()
+        )
 
     def get_user_defined_function(
         self,
@@ -504,19 +431,12 @@ def get_user_defined_function(
         db_name = self._parse_database(database)
         schema_name = self._parse_schema(schema)
         function_id = self._parse_function_or_procedure(udf_name, arg_types)
-
-        try:
-            procedures = self._session.sql(
-                f"DESCRIBE AS RESOURCE FUNCTION {db_name}.{schema_name}.{function_id}"
-            ).collect()
-            return UserDefinedFunction.from_json(str(procedures[0][0]))
-        except (
-            IndexError,  # when sql returned no results
-            SnowparkSQLException,  # when database, or schema doesn't exist
-        ):
-            raise NotFoundError(
-                f"Function with name {udf_name} and arguments {arg_types} could not be found in schema '{db_name}.{schema_name}'"
-            )
+        return (
+            self._root.databases[db_name]
+            .schemas[schema_name]
+            .user_defined_functions[function_id]
+            .fetch()
+        )
 
     # set methods
     def set_current_database(self, database: Union[str, Database]) -> None:
@@ -546,7 +466,7 @@ def database_exists(self, database: Union[str, Database]) -> bool:
         """
         db_name = self._parse_database(database)
         try:
-            self.get_database(db_name)
+            self._root.databases[db_name].fetch()
             return True
         except NotFoundError:
             return False
@@ -567,7 +487,7 @@ def schema_exists(
         db_name = self._parse_database(database, schema)
         schema_name = self._parse_schema(schema)
         try:
-            self.get_schema(schema=schema_name, database=db_name)
+            self._root.databases[db_name].schemas[schema_name].fetch()
             return True
         except NotFoundError:
             return False
@@ -591,7 +511,9 @@ def table_exists(
         schema_name = self._parse_schema(schema, table)
         table_name = table if isinstance(table, str) else table.name
         try:
-            self.get_table(table_name=table_name, database=db_name, schema=schema_name)
+            self._root.databases[db_name].schemas[schema_name].tables[
+                table_name
+            ].fetch()
             return True
         except NotFoundError:
             return False
@@ -615,7 +537,7 @@ def view_exists(
         schema_name = self._parse_schema(schema, view)
         view_name = view if isinstance(view, str) else view.name
         try:
-            self.get_view(view_name=view_name, database=db_name, schema=schema_name)
+            self._root.databases[db_name].schemas[schema_name].views[view_name].fetch()
             return True
         except NotFoundError:
             return False
@@ -637,24 +559,14 @@ def procedure_exists(
             database: database name or ``Database`` object. Defaults to None.
             schema: schema name or ``Schema`` object. Defaults to None.
         """
+        db_name = self._parse_database(database, procedure)
+        schema_name = self._parse_schema(schema, procedure)
+        procedure_id = self._parse_function_or_procedure(procedure, arg_types)
+
         try:
-            if isinstance(procedure, Procedure):
-                if arg_types is not None or database is not None or schema is not None:
-                    raise ArgumentError(
-                        "When provided procedure is a Procedure class no other arguments can be provided"
-                    )
-                database = procedure.database_name
-                schema = procedure.schema_name
-                arg_types = [
-                    type_string_to_type_object(a.datatype) for a in procedure.arguments
-                ]
-                procedure = procedure.name
-            self.get_procedure(
-                procedure_name=procedure,
-                arg_types=arg_types,
-                database=database,
-                schema=schema,
-            )
+            self._root.databases[db_name].schemas[schema_name].procedures[
+                procedure_id
+            ].fetch()
             return True
         except NotFoundError:
             return False
@@ -678,24 +590,14 @@ def user_defined_function_exists(
             database: database name or ``Database`` object. Defaults to None.
             schema: schema name or ``Schema`` object. Defaults to None.
         """
+        db_name = self._parse_database(database, udf)
+        schema_name = self._parse_schema(schema, udf)
+        function_id = self._parse_function_or_procedure(udf, arg_types)
+
         try:
-            if isinstance(udf, UserDefinedFunction):
-                if arg_types is not None or database is not None or schema is not None:
-                    raise ArgumentError(
-                        "When provided udf is a UserDefinedFunction class no other arguments can be provided"
-                    )
-                database = udf.database_name
-                schema = udf.schema_name
-                arg_types = [
-                    type_string_to_type_object(a.datatype) for a in udf.arguments
-                ]
-                udf = udf.name
-            self.get_user_defined_function(
-                udf_name=udf,
-                arg_types=arg_types,
-                database=database,
-                schema=schema,
-            )
+            self._root.databases[db_name].schemas[schema_name].user_defined_functions[
+                function_id
+            ].fetch()
             return True
         except NotFoundError:
             return False
@@ -708,7 +610,7 @@ def drop_database(self, database: Union[str, Database]) -> None:
             database: database name or ``Database`` object.
         """
         db_name = self._parse_database(database)
-        self._session.sql(f"DROP DATABASE {db_name}").collect()
+        self._root.databases[db_name].drop()
 
     def drop_schema(
         self,
@@ -725,7 +627,7 @@ def drop_schema(
         """
         db_name = self._parse_database(database, schema)
         schema_name = self._parse_schema(schema)
-        self._session.sql(f"DROP SCHEMA {db_name}.{schema_name}").collect()
+        self._root.databases[db_name].schemas[schema_name].drop()
 
     def drop_table(
         self,
@@ -746,7 +648,7 @@ def drop_table(
         schema_name = self._parse_schema(schema, table)
         table_name = table if isinstance(table, str) else table.name
 
-        self._session.sql(f"DROP TABLE {db_name}.{schema_name}.{table_name}").collect()
+        self._root.databases[db_name].schemas[schema_name].tables[table_name].drop()
 
     def drop_view(
         self,
@@ -767,7 +669,7 @@ def drop_view(
         schema_name = self._parse_schema(schema, view)
         view_name = view if isinstance(view, str) else view.name
 
-        self._session.sql(f"DROP VIEW {db_name}.{schema_name}.{view_name}").collect()
+        self._root.databases[db_name].schemas[schema_name].views[view_name].drop()
 
     # aliases
     listDatabases = list_databases