snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/relation/catalogs/snowflake_catalog.py

@@ -0,0 +1,467 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import re
+import typing
+
+import pandas
+import pyspark.sql.connect.proto.common_pb2 as common_proto
+import pyspark.sql.connect.proto.types_pb2 as types_proto
+from snowflake.core.exceptions import NotFoundError
+
+from snowflake import snowpark
+from snowflake.snowpark import functions
+from snowflake.snowpark._internal.analyzer.analyzer_utils import (
+    quote_name_without_upper_casing,
+    unquote_if_quoted,
+)
+from snowflake.snowpark.functions import lit
+from snowflake.snowpark.types import BooleanType, StringType
+from snowflake.snowpark_connect import column_name_handler
+from snowflake.snowpark_connect.config import auto_uppercase_ddl, global_config
+from snowflake.snowpark_connect.relation.catalogs.abstract_spark_catalog import (
+    AbstractSparkCatalog,
+    _get_current_snowflake_schema,
+    _process_multi_layer_database,
+    _process_multi_layer_identifier,
+)
+from snowflake.snowpark_connect.type_mapping import proto_to_snowpark_type
+from snowflake.snowpark_connect.utils.attribute_handling import (
+    split_fully_qualified_spark_name,
+)
+from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
+from snowflake.snowpark_connect.utils.telemetry import (
+    SnowparkConnectNotImplementedError,
+)
+from snowflake.snowpark_connect.utils.udf_cache import cached_udf
+
+
+def _normalize_identifier(identifier: str | None) -> str | None:
+    if identifier is None:
+        return None
+    return identifier.upper() if auto_uppercase_ddl() else identifier
+
+
+def sf_quote(name: str | None) -> str | None:
+    if name is None:
+        return None
+    return quote_name_without_upper_casing(_normalize_identifier(name))
+
+
+class SnowflakeCatalog(AbstractSparkCatalog):
+    def __init__(self) -> None:
+        super().__init__(name="spark_catalog", description=None)
+
+    def listDatabases(
+        self,
+        pattern: str | None = None,
+    ) -> pandas.DataFrame:
+        """List all databases accessible in Snowflake with an optional name to filter by."""
+
+        if pattern == "":
+            return pandas.DataFrame([])
+
+        # This pattern is case-sensitive while our SAS implementation is not
+        catalog, sf_database, sf_schema = _process_multi_layer_database(pattern)
+        sf_schema = sf_schema.replace("*", ".*")
+        if catalog is not None and self != catalog:
+            raise SnowparkConnectNotImplementedError(
+                "Calling into another catalog is not currently supported"
+            )
+        sp_catalog = get_or_create_snowpark_session().catalog
+
+        dbs = sp_catalog.list_schemas(
+            database=sf_quote(sf_database),
+            pattern=_normalize_identifier(sf_schema),
+        )
+        names: list[str] = list()
+        catalogs: list[str] = list()
+        descriptions: list[str | None] = list()
+        locationUris: list[str] = list()
+        for db in dbs:
+            name = unquote_if_quoted(db.name)
+            if name == "INFORMATION_SCHEMA" and global_config._get_config_setting(
+                "spark.Catalog.databaseFilterInformationSchema"
+            ):
+                continue
+            names.append(name)
+            catalogs.append(self.name)
+            descriptions.append(db.comment)
+            locationUris.append(f"snowflake://{name}")
+        return pandas.DataFrame(
+            {
+                "name": names,
+                "catalog": catalogs,
+                "description": descriptions,
+                "locationUri": locationUris,
+            }
+        )
+
+    def getDatabase(
+        self,
+        spark_dbName: str,
+    ) -> pandas.DataFrame:
+        """Listing a single database that's accessible in Snowflake."""
+        catalog, sf_database, sf_schema = _process_multi_layer_database(spark_dbName)
+        if catalog is not None and self != catalog:
+            raise SnowparkConnectNotImplementedError(
+                "Calling into another catalog is not currently supported"
+            )
+        sp_catalog = get_or_create_snowpark_session().catalog
+
+        db = sp_catalog.get_schema(
+            schema=sf_quote(sf_schema), database=sf_quote(sf_database)
+        )
+
+        name = unquote_if_quoted(db.name)
+        return pandas.DataFrame(
+            {
+                "name": [name],
+                "catalog": [self.name],
+                "description": [db.comment],
+                "locationUri": [f"snowflake://{name}"],
+            }
+        )
+
+    def databaseExists(
+        self,
+        spark_dbName: str,
+    ) -> pandas.DataFrame:
+        """Whether a database with provided name exists in Snowflake."""
+        try:
+            self.getDatabase(spark_dbName)
+            exists = True
+        except NotFoundError:
+            exists = False
+        return pandas.DataFrame({"exists": [exists]})
+
+    def listTables(
+        self,
+        spark_dbName: str | None = None,
+        pattern: str | None = None,
+    ) -> pandas.DataFrame:
+        """Listing all tables/views accessible in Snowflake, optionally filterable on database, schema, and a pattern for the table names."""
+        if spark_dbName is not None:
+            catalog, sf_database, sf_schema = _process_multi_layer_database(
+                spark_dbName
+            )
+            if catalog is not None and self != catalog:
+                raise SnowparkConnectNotImplementedError(
+                    "Calling into another catalog is not currently supported"
+                )
+        else:
+            catalog = sf_database = sf_schema = None
+
+        tables = self._list_objects(
+            object_name="TABLES",
+            database=sf_quote(sf_database),
+            schema=sf_quote(sf_schema),
+            pattern=_normalize_identifier(pattern),
+        )
+        views = self._list_objects(
+            object_name="VIEWS",
+            database=sf_quote(sf_database),
+            schema=sf_quote(sf_schema),
+            pattern=_normalize_identifier(pattern),
+        )
+        names: list[str] = list()
+        catalogs: list[str] = list()
+        namespaces: list[list[str | None]] = list()
+        names: list[str] = list()
+        descriptions: list[str | None] = list()
+        table_types: list[str | None] = list()
+        is_temporaries: list[bool] = list()
+        for o in tables:
+            names.append(unquote_if_quoted(o[1]))
+            catalogs.append(self.name)
+            namespaces.append([unquote_if_quoted(o[3])])
+            descriptions.append(o[5] if o[5] else None)
+            table_types.append("PERMANENT" if o[4] == "TABLE" else o[4])
+            is_temporaries.append(o[4] == "TEMPORARY")
+        for o in views:
+            names.append(unquote_if_quoted(o[1]))
+            catalogs.append(self.name)
+            namespaces.append([unquote_if_quoted(o[4])])
+            descriptions.append(o[6] if o[6] else None)
+            table_types.append("PERMANENT")
+            is_temporaries.append(False)
+        return pandas.DataFrame(
+            {
+                "name": names,
+                "catalog": catalogs,
+                "namespace": namespaces,
+                "description": descriptions,
+                "tableType": table_types,
+                "isTemporary": is_temporaries,
+            }
+        )
+
+    def _list_objects(
+        self,
+        *,
+        object_name: str,
+        database: typing.Optional[str],
+        schema: typing.Optional[str],
+        pattern: typing.Optional[str] = None,
+    ):
+        session = get_or_create_snowpark_session()
+        if not database:
+            database = session.catalog.get_current_database()
+        if not schema:
+            schema = session.catalog.get_current_schema()
+        df = get_or_create_snowpark_session().sql(
+            f"SHOW {object_name} IN {database}.{schema}"
+        )
+        if pattern:
+
+            @cached_udf(
+                input_types=[StringType(), StringType()], return_type=BooleanType()
+            )
+            def python_regex_filter(pattern: str, input: str) -> bool:
+                return bool(re.match(pattern, input))
+
+            df = df.filter(python_regex_filter(lit(pattern), df['"name"']))
+
+        return df.collect()
+
+    def getTable(
+        self,
+        spark_tableName: str,
+    ) -> pandas.DataFrame:
+        """Listing a single table/view with provided name that's accessible in Snowflake."""
+        sp_catalog = get_or_create_snowpark_session().catalog
+        catalog, sf_database, sf_schema, table_name = _process_multi_layer_identifier(
+            spark_tableName
+        )
+        if catalog is not None and self != catalog:
+            raise SnowparkConnectNotImplementedError(
+                "Calling into another catalog is not currently supported"
+            )
+
+        table = sp_catalog.get_table(
+            database=sf_quote(sf_database),
+            schema=sf_quote(sf_schema),
+            table_name=sf_quote(table_name),
+        )
+
+        return pandas.DataFrame(
+            {
+                "name": [unquote_if_quoted(table.name)],
+                "catalog": [self.name],
+                "namespace": [[unquote_if_quoted(table.schema_name)]],
+                "description": [table.comment],
+                "tableType": [table.kind],
+                "isTemporary": [table.kind == "TEMPORARY"],
+            }
+        )
+
+    def tableExists(
+        self,
+        spark_tableName: str,
+        spark_dbName: str | None,
+    ) -> pandas.DataFrame:
+        """Whether a table/view with provided name exists in Snowflake, optionally filterable with dbName.
+        If no database is specified, first try to treat tableName as a multi-layer-namespace identifier
+        (or fully qualified name), then try tableName as a normal table name in the current database if necessary.
+        Argument dbName is not actually implemented yet while we figure out how to map databases from Spark to Snowflake.
+        """
+        table_mli = spark_tableName
+        if spark_dbName:
+            table_mli = f"{spark_dbName}.{table_mli}"
+
+        try:
+            self.getTable(table_mli)
+            exists = True
+        except NotFoundError:
+            exists = False
+        return pandas.DataFrame({"exists": [exists]})
+
+    def listColumns(
+        self,
+        spark_tableName: str,
+        spark_dbName: typing.Optional[str] = None,
+    ) -> pandas.DataFrame:
+        """List all columns in a table/view, optionally database name filter can be provided."""
+        sp_catalog = get_or_create_snowpark_session().catalog
+        if spark_dbName is None:
+            catalog, sf_database, sf_schema, sf_table = _process_multi_layer_identifier(
+                spark_tableName
+            )
+            if catalog is not None and self != catalog:
+                raise SnowparkConnectNotImplementedError(
+                    "Calling into another catalog is not currently supported"
+                )
+            columns = sp_catalog.list_columns(
+                database=sf_quote(sf_database),
+                schema=sf_quote(sf_schema),
+                table_name=sf_quote(sf_table),
+            )
+        else:
+            columns = sp_catalog.list_columns(
+                schema=sf_quote(spark_dbName),
+                table_name=sf_quote(spark_tableName),
+            )
+        names: list[str] = list()
+        descriptions: list[str | None] = list()
+        data_types: list[str] = list()
+        nullables: list[bool] = list()
+        is_partitions: list[bool] = list()
+        is_buckets: list[bool] = list()
+        for column in columns:
+            names.append(unquote_if_quoted(column.name))
+            descriptions.append(column.comment)
+            data_types.append(column.datatype)
+            nullables.append(bool(column.nullable))
+            is_partitions.append(False)
+            is_buckets.append(False)
+
+        return pandas.DataFrame(
+            {
+                "name": names,
+                "description": descriptions,
+                "dataType": data_types,
+                "nullable": nullables,
+                "isPartition": is_partitions,
+                "isBucket": is_buckets,
+            }
+        )
+
+    def currentDatabase(self) -> pandas.DataFrame:
+        """Get the currently used database's name."""
+        db_name = _get_current_snowflake_schema()
+        assert db_name is not None, "current database could not be confirmed"
+        return pandas.DataFrame({"current_database": [unquote_if_quoted(db_name)]})
+
+    def setCurrentDatabase(
+        self,
+        spark_dbName: str,
+    ) -> pandas.DataFrame:
+        """Set the currently used database's name."""
+        sp_catalog = get_or_create_snowpark_session().catalog
+        sp_catalog.setCurrentSchema(sf_quote(spark_dbName))
+        return pandas.DataFrame({"current_database": [spark_dbName]})
+
+    def dropGlobalTempView(
+        self,
+        spark_view_name: str,
+    ) -> snowpark.DataFrame:
+        session = get_or_create_snowpark_session()
+        schema = global_config.spark_sql_globalTempDatabase
+        result_df = session.sql(
+            "drop view if exists identifier(?)",
+            params=[f"{sf_quote(schema)}.{sf_quote(spark_view_name)}"],
+        )
+        result_df = result_df.select(
+            functions.contains('"status"', functions.lit("successfully dropped")).alias(
+                "value"
+            )
+        )
+        columns = ["value"]
+        return column_name_handler.with_column_map(
+            result_df, columns, columns, [BooleanType()]
+        )
+
+    def dropTempView(
+        self,
+        spark_view_name: str,
+    ) -> snowpark.DataFrame:
+        """Drop the current temporary view."""
+        session = get_or_create_snowpark_session()
+        result = session.sql(
+            "drop view if exists identifier(?)",
+            params=[sf_quote(spark_view_name)],
+        ).collect()
+        view_was_dropped = (
+            len(result) == 1 and "successfully dropped" in result[0]["status"]
+        )
+        result_df = session.createDataFrame([(view_was_dropped,)], schema=["value"])
+        columns = ["value"]
+        return column_name_handler.with_column_map(
+            result_df, columns, columns, [BooleanType()]
+        )
+
+    def createTable(
+        self,
+        tableName: str,
+        path: str,
+        source: str,
+        schema: types_proto.DataType,
+        description: str,
+        **options: typing.Any,
+    ) -> snowpark.DataFrame:
+        """Create either an external, or a managed table.
+
+        If path is supplied in which the data for this table exists. When path is specified, an external table is
+        created from the data at the given path. Otherwise a managed table is created.
+
+        In case a managed table is being created, schema is required.
+        """
+        # TODO: support fully-qualified tableName
+        if source == "":
+            source = global_config.get("spark.sql.sources.default")
+        if source not in ("csv", "json", "avro", "parquet", "orc", "xml"):
+            raise SnowparkConnectNotImplementedError(
+                f"Source '{source}' is not currently supported by Catalog.createTable. "
+                "Maybe default value through 'spark.sql.sources.default' should be set."
+            )
+        if path != "":
+            # External table creation is not supported currently.
+            raise SnowparkConnectNotImplementedError(
+                "External table creation is not supported currently."
+            )
+
+        session = get_or_create_snowpark_session()
+        # Managed table
+        if schema.ByteSize() == 0:
+            raise SnowparkConnectNotImplementedError(
+                f"Unable to infer schema for {source.upper()}. It must be specified manually.",
+            )
+        sp_schema = proto_to_snowpark_type(schema)
+        columns = [c.name for c in schema.struct.fields]
+        table_name_parts = split_fully_qualified_spark_name(tableName)
+        qualifiers = [table_name_parts for _ in columns]
+        column_types = [f.datatype for f in sp_schema.fields]
+        return column_name_handler.with_column_map(
+            session.createDataFrame([], sp_schema),
+            columns,
+            columns,
+            column_types,
+            column_qualifiers=qualifiers,
+        )
+
+    def isCached(self, spark_tableName: str) -> pandas.DataFrame:
+        """Whether a table is cached by us locally.
+
+        Check whether a table exists and then delegate to the local cache.
+        """
+        self.getTable(spark_tableName)
+        return super().isCached(spark_tableName)
+
+    def cacheTable(
+        self,
+        spark_tableName: str,
+        storageLevel: common_proto.StorageLevel | None = None,
+    ) -> pandas.DataFrame:
+        """Cache a table, or view locally.
+
+        Check whether a table exists and then delegate to the local cache.
+        """
+        self.getTable(spark_tableName)
+        return super().cacheTable(spark_tableName, storageLevel)
+
+    def uncacheTable(self, spark_tableName: str) -> pandas.DataFrame:
+        """Uncache a table, or view locally.
+
+        Check whether a table exists and then delegate to the local cache.
+        """
+        self.getTable(spark_tableName)
+        return super().uncacheTable(spark_tableName)
+
+    def refreshTable(self, spark_tableName: str) -> pandas.DataFrame:
+        """Refresh a table, or view locally.
+
+        Check whether a table exists and then delegate to the local cache.
+        """
+        self.getTable(spark_tableName)
+        return super().refreshTable(spark_tableName)

snowflake/snowpark_connect/relation/catalogs/utils.py

@@ -0,0 +1,51 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+from collections import defaultdict
+
+from snowflake.connector.errors import ProgrammingError
+from snowflake.snowpark_connect.relation.catalogs import CATALOGS, SNOWFLAKE_CATALOG
+from snowflake.snowpark_connect.relation.catalogs.abstract_spark_catalog import (
+    AbstractSparkCatalog,
+)
+from snowflake.snowpark_connect.utils.session import get_or_create_snowpark_session
+
+CURRENT_CATALOG = SNOWFLAKE_CATALOG
+CURRENT_CATALOG_NAME: str | None = "spark_catalog"
+CATALOG_TEMP_OBJECTS: defaultdict[
+    str | None, set[tuple[str | None, str | None, str]]
+] = defaultdict(set)
+
+
+def get_current_catalog() -> AbstractSparkCatalog:
+    return CURRENT_CATALOG
+
+
+def set_current_catalog(catalog_name: str | None) -> AbstractSparkCatalog:
+    global CURRENT_CATALOG_NAME
+
+    # Validate input parameters to match PySpark behavior
+    if catalog_name is None:
+        raise ValueError("Catalog name cannot be None")
+    if catalog_name == "":
+        raise ValueError(
+            "Catalog '' plugin class not found: spark.sql.catalog. is not defined"
+        )
+
+    CURRENT_CATALOG_NAME = catalog_name
+    if catalog_name in CATALOGS:
+        return CATALOGS[catalog_name]
+
+    sf_catalog = get_or_create_snowpark_session().catalog
+    try:
+        sf_catalog.setCurrentDatabase(catalog_name if catalog_name is not None else "")
+        return get_current_catalog()
+    except ProgrammingError as e:
+        raise Exception(
+            f"Catalog '{catalog_name}' plugin class not found: spark.sql.catalog.{catalog_name} is not defined"
+        ) from e
+
+
+def _get_current_temp_objects() -> set[tuple[str | None, str | None, str]]:
+    return CATALOG_TEMP_OBJECTS[CURRENT_CATALOG_NAME]

snowflake/snowpark_connect/relation/io_utils.py

@@ -0,0 +1,76 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+from urllib.parse import urlparse
+
+CLOUD_PREFIX_TO_CLOUD = {
+    "abfss": "azure",
+    "wasbs": "azure",
+}
+
+
+def get_cloud_from_url(
+    url: str,
+):
+    """
+    Google Cloud Storage : 'gcs://<bucket>[/<path>/]'
+    Microsoft Azure : 'azure://<account>.blob.core.windows.net/<container>[/<path>/]'
+    """
+    url_parts = url.split("://")
+    if len(url_parts) < 2:
+        return None
+
+    cloud_domain = url_parts[0].lower()
+    return CLOUD_PREFIX_TO_CLOUD.get(cloud_domain, cloud_domain)
+
+
+def split_url_paths(url):
+    return url.split("://")[1].split("/")
+
+
+def parse_gcp_url(url: str):
+    """
+
+    Args:
+        url: 'gcs://<bucket>[/<path>/]'
+
+    """
+    path_parts = split_url_paths(url)
+    bucket_name = path_parts[0]
+    path = "/".join(path_parts[1:])
+    return bucket_name, path
+
+
+def parse_azure_url(url: str):
+    """
+
+    Args:
+        url: 'azure://<account>.blob.core.windows.net/<container>[/<path>/]'
+        url: "abfss://<container-name>@<storage-account-name>.dfs.core.windows.net/<path-to-data>"
+
+    """
+    path_parts = split_url_paths(url)
+    prefix = path_parts[0].split(".")[0].split("@")
+    account = prefix[1]
+    bucket_name = prefix[0]
+    path = "/".join(path_parts[1:])
+    return account, bucket_name, path
+
+
+def is_cloud_path(path: str) -> bool:
+    return (
+        path.startswith("@")  # Snowflake Stage
+        or path.startswith("s3://")
+        or path.startswith("s3a://")  # AWS S3
+        or path.startswith("azure://")
+        or path.startswith("abfss://")
+        or path.startswith("wasbs://")  # Azure
+        or path.startswith("gcs://")  # GCP
+    )
+
+
+def convert_file_prefix_path(path: str) -> str:
+    if path.startswith("file:/"):
+        return urlparse(path).path
+    return path