polars-runtime-compat 1.34.0b3__cp39-abi3-macosx_11_0_arm64.whl → 1.34.0b4__cp39-abi3-macosx_11_0_arm64.whl

polars/catalog/unity/client.py

@@ -1,733 +0,0 @@
-from __future__ import annotations
-
-import contextlib
-import importlib
-import os
-import sys
-from typing import TYPE_CHECKING, Any, Literal
-
-from polars._utils.unstable import issue_unstable_warning
-from polars._utils.wrap import wrap_ldf
-from polars.catalog.unity.models import (
-    CatalogInfo,
-    ColumnInfo,
-    NamespaceInfo,
-    TableInfo,
-)
-
-if TYPE_CHECKING:
-    from collections.abc import Generator
-    from datetime import datetime
-
-    import deltalake
-
-    from polars._typing import SchemaDict
-    from polars.catalog.unity.models import DataSourceFormat, TableType
-    from polars.dataframe.frame import DataFrame
-    from polars.io.cloud import (
-        CredentialProviderFunction,
-        CredentialProviderFunctionReturn,
-    )
-    from polars.io.cloud.credential_provider._builder import CredentialProviderBuilder
-    from polars.lazyframe import LazyFrame
-
-with contextlib.suppress(ImportError):
-    from polars._plr import PyCatalogClient
-
-    PyCatalogClient.init_classes(
-        catalog_info_cls=CatalogInfo,
-        namespace_info_cls=NamespaceInfo,
-        table_info_cls=TableInfo,
-        column_info_cls=ColumnInfo,
-    )
-
-
-class Catalog:
-    """
-    Unity catalog client.
-
-    .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-    """
-
-    def __init__(
-        self,
-        workspace_url: str,
-        *,
-        bearer_token: str | None = "auto",
-        require_https: bool = True,
-    ) -> None:
-        """
-        Initialize a catalog client.
-
-        .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-
-        Parameters
-        ----------
-        workspace_url
-            URL of the workspace, or alternatively the URL of the Unity catalog
-            API endpoint.
-        bearer_token
-            Bearer token to authenticate with. This can also be set to:
-
-            * "auto": Automatically retrieve bearer tokens from the environment.
-            * "databricks-sdk": Use the Databricks SDK to retrieve and use the
-              bearer token from the environment.
-        require_https
-            Require the `workspace_url` to use HTTPS.
-        """
-        issue_unstable_warning("`Catalog` functionality is considered unstable.")
-
-        if require_https and not workspace_url.startswith("https://"):
-            msg = (
-                f"a non-HTTPS workspace_url was given ({workspace_url}). To "
-                "allow non-HTTPS URLs, pass require_https=False."
-            )
-            raise ValueError(msg)
-
-        if bearer_token == "databricks-sdk" or (
-            bearer_token == "auto"
-            # For security, in "auto" mode, only retrieve/use the token if:
-            # * We are running inside a Databricks environment
-            # * The `workspace_url` is pointing to Databricks and uses HTTPS
-            and "DATABRICKS_RUNTIME_VERSION" in os.environ
-            and workspace_url.startswith("https://")
-            and (
-                workspace_url.removeprefix("https://")
-                .split("/", 1)[0]
-                .endswith(".cloud.databricks.com")
-            )
-        ):
-            bearer_token = self._get_databricks_token()
-
-        if bearer_token == "auto":
-            bearer_token = None
-
-        self._client = PyCatalogClient.new(workspace_url, bearer_token)
-
-    def list_catalogs(self) -> list[CatalogInfo]:
-        """
-        List the available catalogs.
-
-        .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-        """
-        return self._client.list_catalogs()
-
-    def list_namespaces(self, catalog_name: str) -> list[NamespaceInfo]:
-        """
-        List the available namespaces (unity schema) under the specified catalog.
-
-        .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-
-        Parameters
-        ----------
-        catalog_name
-            Name of the catalog.
-        """
-        return self._client.list_namespaces(catalog_name)
-
-    def list_tables(self, catalog_name: str, namespace: str) -> list[TableInfo]:
-        """
-        List the available tables under the specified schema.
-
-        .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-
-        Parameters
-        ----------
-        catalog_name
-            Name of the catalog.
-        namespace
-            Name of the namespace (unity schema).
-        """
-        return self._client.list_tables(catalog_name, namespace)
-
-    def get_table_info(
-        self, catalog_name: str, namespace: str, table_name: str
-    ) -> TableInfo:
-        """
-        Retrieve the metadata of the specified table.
-
-        .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-
-        Parameters
-        ----------
-        catalog_name
-            Name of the catalog.
-        namespace
-            Name of the namespace (unity schema).
-        table_name
-            Name of the table.
-        """
-        return self._client.get_table_info(catalog_name, namespace, table_name)
-
-    def _get_table_credentials(
-        self, table_id: str, *, write: bool
-    ) -> tuple[dict[str, str] | None, dict[str, str], int]:
-        return self._client.get_table_credentials(table_id=table_id, write=write)
-
-    def scan_table(
-        self,
-        catalog_name: str,
-        namespace: str,
-        table_name: str,
-        *,
-        delta_table_version: int | str | datetime | None = None,
-        delta_table_options: dict[str, Any] | None = None,
-        storage_options: dict[str, Any] | None = None,
-        credential_provider: (
-            CredentialProviderFunction | Literal["auto"] | None
-        ) = "auto",
-        retries: int = 2,
-    ) -> LazyFrame:
-        """
-        Retrieve the metadata of the specified table.
-
-        .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-
-        Parameters
-        ----------
-        catalog_name
-            Name of the catalog.
-        namespace
-            Name of the namespace (unity schema).
-        table_name
-            Name of the table.
-        delta_table_version
-            Version of the table to scan (Deltalake only).
-        delta_table_options
-            Additional keyword arguments while reading a Deltalake table.
-        storage_options
-            Options that indicate how to connect to a cloud provider.
-
-            The cloud providers currently supported are AWS, GCP, and Azure.
-            See supported keys here:
-
-            * `aws <https://docs.rs/object_store/latest/object_store/aws/enum.AmazonS3ConfigKey.html>`_
-            * `gcp <https://docs.rs/object_store/latest/object_store/gcp/enum.GoogleConfigKey.html>`_
-            * `azure <https://docs.rs/object_store/latest/object_store/azure/enum.AzureConfigKey.html>`_
-            * Hugging Face (`hf://`): Accepts an API key under the `token` parameter: \
-            `{'token': '...'}`, or by setting the `HF_TOKEN` environment variable.
-
-            If `storage_options` is not provided, Polars will try to infer the
-            information from environment variables.
-        credential_provider
-            Provide a function that can be called to provide cloud storage
-            credentials. The function is expected to return a dictionary of
-            credential keys along with an optional credential expiry time.
-
-            .. warning::
-                This functionality is considered **unstable**. It may be changed
-                at any point without it being considered a breaking change.
-        retries
-            Number of retries if accessing a cloud instance fails.
-
-        """
-        table_info = self.get_table_info(catalog_name, namespace, table_name)
-        storage_location, data_source_format = _extract_location_and_data_format(
-            table_info, "scan table"
-        )
-
-        credential_provider, storage_options = self._init_credentials(  # type: ignore[assignment]
-            credential_provider,
-            storage_options,
-            table_info,
-            write=False,
-            caller_name="Catalog.scan_table",
-        )
-
-        if data_source_format in ["DELTA", "DELTASHARING"]:
-            from polars.io.delta import scan_delta
-
-            return scan_delta(
-                storage_location,
-                version=delta_table_version,
-                delta_table_options=delta_table_options,
-                storage_options=storage_options,
-                credential_provider=credential_provider,
-            )
-
-        if delta_table_version is not None:
-            msg = (
-                "cannot apply delta_table_version for table of type "
-                f"{data_source_format}"
-            )
-            raise ValueError(msg)
-
-        if delta_table_options is not None:
-            msg = (
-                "cannot apply delta_table_options for table of type "
-                f"{data_source_format}"
-            )
-            raise ValueError(msg)
-
-        if storage_options:
-            storage_options = list(storage_options.items())  # type: ignore[assignment]
-        else:
-            # Handle empty dict input
-            storage_options = None
-
-        return wrap_ldf(
-            self._client.scan_table(
-                catalog_name,
-                namespace,
-                table_name,
-                credential_provider=credential_provider,
-                cloud_options=storage_options,
-                retries=retries,
-            )
-        )
-
-    def write_table(
-        self,
-        df: DataFrame,
-        catalog_name: str,
-        namespace: str,
-        table_name: str,
-        *,
-        delta_mode: Literal[
-            "error", "append", "overwrite", "ignore", "merge"
-        ] = "error",
-        delta_write_options: dict[str, Any] | None = None,
-        delta_merge_options: dict[str, Any] | None = None,
-        storage_options: dict[str, str] | None = None,
-        credential_provider: CredentialProviderFunction
-        | Literal["auto"]
-        | None = "auto",
-    ) -> None | deltalake.table.TableMerger:
-        """
-        Write a DataFrame to a catalog table.
-
-        .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-
-        Parameters
-        ----------
-        df
-            DataFrame to write.
-        catalog_name
-            Name of the catalog.
-        namespace
-            Name of the namespace (unity schema).
-        table_name
-            Name of the table.
-        delta_mode : {'error', 'append', 'overwrite', 'ignore', 'merge'}
-            (For delta tables) How to handle existing data.
-
-            - If 'error', throw an error if the table already exists (default).
-            - If 'append', will add new data.
-            - If 'overwrite', will replace table with new data.
-            - If 'ignore', will not write anything if table already exists.
-            - If 'merge', return a `TableMerger` object to merge data from the DataFrame
-              with the existing data.
-        delta_write_options
-            (For delta tables) Additional keyword arguments while writing a
-            Delta lake Table.
-            See a list of supported write options `here <https://delta-io.github.io/delta-rs/api/delta_writer/#deltalake.write_deltalake>`__.
-        delta_merge_options
-            (For delta tables) Keyword arguments which are required to `MERGE` a
-            Delta lake Table.
-            See a list of supported merge options `here <https://delta-io.github.io/delta-rs/api/delta_table/#deltalake.DeltaTable.merge>`__.
-        storage_options
-            Options that indicate how to connect to a cloud provider.
-
-            The cloud providers currently supported are AWS, GCP, and Azure.
-            See supported keys here:
-
-            * `aws <https://docs.rs/object_store/latest/object_store/aws/enum.AmazonS3ConfigKey.html>`_
-            * `gcp <https://docs.rs/object_store/latest/object_store/gcp/enum.GoogleConfigKey.html>`_
-            * `azure <https://docs.rs/object_store/latest/object_store/azure/enum.AzureConfigKey.html>`_
-            * Hugging Face (`hf://`): Accepts an API key under the `token` parameter: \
-            `{'token': '...'}`, or by setting the `HF_TOKEN` environment variable.
-
-            If `storage_options` is not provided, Polars will try to infer the
-            information from environment variables.
-        credential_provider
-            Provide a function that can be called to provide cloud storage
-            credentials. The function is expected to return a dictionary of
-            credential keys along with an optional credential expiry time.
-
-            .. warning::
-                This functionality is considered **unstable**. It may be changed
-                at any point without it being considered a breaking change.
-        """
-        table_info = self.get_table_info(catalog_name, namespace, table_name)
-        storage_location, data_source_format = _extract_location_and_data_format(
-            table_info, "scan table"
-        )
-
-        credential_provider, storage_options = self._init_credentials(  # type: ignore[assignment]
-            credential_provider,
-            storage_options,
-            table_info,
-            write=True,
-            caller_name="Catalog.write_table",
-        )
-
-        if data_source_format in ["DELTA", "DELTASHARING"]:
-            return df.write_delta(  # type: ignore[misc]
-                storage_location,
-                storage_options=storage_options,
-                credential_provider=credential_provider,
-                mode=delta_mode,
-                delta_write_options=delta_write_options,
-                delta_merge_options=delta_merge_options,
-            )  # type: ignore[call-overload]
-
-        else:
-            msg = (
-                "write_table: table format of "
-                f"{catalog_name}.{namespace}.{table_name} "
-                f"({data_source_format}) is unsupported."
-            )
-            raise NotImplementedError(msg)
-
-    def create_catalog(
-        self,
-        catalog_name: str,
-        *,
-        comment: str | None = None,
-        storage_root: str | None = None,
-    ) -> CatalogInfo:
-        """
-        Create a catalog.
-
-        .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-
-        Parameters
-        ----------
-        catalog_name
-            Name of the catalog.
-        comment
-            Leaves a comment about the catalog.
-        storage_root
-            Base location at which to store the catalog.
-        """
-        return self._client.create_catalog(
-            catalog_name=catalog_name, comment=comment, storage_root=storage_root
-        )
-
-    def delete_catalog(
-        self,
-        catalog_name: str,
-        *,
-        force: bool = False,
-    ) -> None:
-        """
-        Delete a catalog.
-
-        Note that depending on the table type and catalog server, this may not
-        delete the actual data files from storage. For more details, please
-        consult the documentation of the catalog provider you are using.
-
-        .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-
-        Parameters
-        ----------
-        catalog_name
-            Name of the catalog.
-        force
-            Forcibly delete the catalog even if it is not empty.
-        """
-        self._client.delete_catalog(catalog_name=catalog_name, force=force)
-
-    def create_namespace(
-        self,
-        catalog_name: str,
-        namespace: str,
-        *,
-        comment: str | None = None,
-        storage_root: str | None = None,
-    ) -> NamespaceInfo:
-        """
-        Create a namespace (unity schema) in the catalog.
-
-        .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-
-        Parameters
-        ----------
-        catalog_name
-            Name of the catalog.
-        namespace
-            Name of the namespace (unity schema).
-        comment
-            Leaves a comment about the table.
-        storage_root
-            Base location at which to store the namespace.
-        """
-        return self._client.create_namespace(
-            catalog_name=catalog_name,
-            namespace=namespace,
-            comment=comment,
-            storage_root=storage_root,
-        )
-
-    def delete_namespace(
-        self,
-        catalog_name: str,
-        namespace: str,
-        *,
-        force: bool = False,
-    ) -> None:
-        """
-        Delete a namespace (unity schema) in the catalog.
-
-        Note that depending on the table type and catalog server, this may not
-        delete the actual data files from storage. For more details, please
-        consult the documentation of the catalog provider you are using.
-
-        .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-
-        Parameters
-        ----------
-        catalog_name
-            Name of the catalog.
-        namespace
-            Name of the namespace (unity schema).
-        force
-            Forcibly delete the namespace even if it is not empty.
-        """
-        self._client.delete_namespace(
-            catalog_name=catalog_name, namespace=namespace, force=force
-        )
-
-    def create_table(
-        self,
-        catalog_name: str,
-        namespace: str,
-        table_name: str,
-        *,
-        schema: SchemaDict | None,
-        table_type: TableType,
-        data_source_format: DataSourceFormat | None = None,
-        comment: str | None = None,
-        storage_root: str | None = None,
-        properties: dict[str, str] | None = None,
-    ) -> TableInfo:
-        """
-        Create a table in the catalog.
-
-        .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-
-        Parameters
-        ----------
-        catalog_name
-            Name of the catalog.
-        namespace
-            Name of the namespace (unity schema).
-        table_name
-            Name of the table.
-        schema
-            Schema of the table.
-        table_type
-            Type of the table
-        data_source_format
-            Storage format of the table.
-        comment
-            Leaves a comment about the table.
-        storage_root
-            Base location at which to store the table.
-        properties
-            Extra key-value metadata to store.
-        """
-        return self._client.create_table(
-            catalog_name=catalog_name,
-            namespace=namespace,
-            table_name=table_name,
-            schema=schema,
-            table_type=table_type,
-            data_source_format=data_source_format,
-            comment=comment,
-            storage_root=storage_root,
-            properties=list((properties or {}).items()),
-        )
-
-    def delete_table(
-        self,
-        catalog_name: str,
-        namespace: str,
-        table_name: str,
-    ) -> None:
-        """
-        Delete the table stored at this location.
-
-        Note that depending on the table type and catalog server, this may not
-        delete the actual data files from storage. For more details, please
-        consult the documentation of the catalog provider you are using.
-
-        If you would like to perform manual deletions, the storage location of
-        the files can be found using `get_table_info`.
-
-        .. warning::
-            This functionality is considered **unstable**. It may be changed
-            at any point without it being considered a breaking change.
-
-        Parameters
-        ----------
-        catalog_name
-            Name of the catalog.
-        namespace
-            Name of the namespace (unity schema).
-        table_name
-            Name of the table.
-        """
-        self._client.delete_table(
-            catalog_name=catalog_name,
-            namespace=namespace,
-            table_name=table_name,
-        )
-
-    def _init_credentials(
-        self,
-        credential_provider: CredentialProviderFunction | Literal["auto"] | None,
-        storage_options: dict[str, Any] | None,
-        table_info: TableInfo,
-        *,
-        write: bool,
-        caller_name: str,
-    ) -> tuple[
-        CredentialProviderBuilder | None,
-        dict[str, Any] | None,
-    ]:
-        from polars.io.cloud.credential_provider._builder import (
-            CredentialProviderBuilder,
-        )
-
-        if credential_provider != "auto":
-            if credential_provider:
-                return CredentialProviderBuilder.from_initialized_provider(
-                    credential_provider
-                ), storage_options
-            else:
-                return None, storage_options
-
-        verbose = os.getenv("POLARS_VERBOSE") == "1"
-
-        catalog_credential_provider = CatalogCredentialProvider(
-            self, table_info.table_id, write=write
-        )
-
-        try:
-            v = catalog_credential_provider._credentials_iter()
-            storage_update_options = next(v)
-
-            if storage_update_options:
-                storage_options = {**(storage_options or {}), **storage_update_options}
-
-            for _ in v:
-                pass
-
-        except Exception as e:
-            if verbose:
-                table_name = table_info.name
-                table_id = table_info.table_id
-                msg = (
-                    f"error auto-initializing CatalogCredentialProvider: {e!r} "
-                    f"{table_name = } ({table_id = }) ({write = })"
-                )
-                print(msg, file=sys.stderr)
-        else:
-            if verbose:
-                table_name = table_info.name
-                table_id = table_info.table_id
-                msg = (
-                    "auto-selected CatalogCredentialProvider for "
-                    f"{table_name = } ({table_id = })"
-                )
-                print(msg, file=sys.stderr)
-
-            return CredentialProviderBuilder.from_initialized_provider(
-                catalog_credential_provider
-            ), storage_options
-
-        # This should generally not happen, but if using the temporary
-        # credentials API fails for whatever reason, we fallback to our built-in
-        # credential provider resolution.
-
-        from polars.io.cloud.credential_provider._builder import (
-            _init_credential_provider_builder,
-        )
-
-        return _init_credential_provider_builder(
-            "auto", table_info.storage_location, storage_options, caller_name
-        ), storage_options
-
-    @classmethod
-    def _get_databricks_token(cls) -> str:
-        if importlib.util.find_spec("databricks.sdk") is None:
-            msg = "could not get Databricks token: databricks-sdk is not installed"
-            raise ImportError(msg)
-
-        # We code like this to bypass linting
-        m = importlib.import_module("databricks.sdk.core").__dict__
-
-        return m["DefaultCredentials"]()(m["Config"]())()["Authorization"][7:]
-
-
-class CatalogCredentialProvider:
-    """Retrieves credentials from the Unity catalog temporary credentials API."""
-
-    def __init__(self, catalog: Catalog, table_id: str, *, write: bool) -> None:
-        self.catalog = catalog
-        self.table_id = table_id
-        self.write = write
-
-    def __call__(self) -> CredentialProviderFunctionReturn:  # noqa: D102
-        _, (creds, expiry) = self._credentials_iter()
-        return creds, expiry
-
-    def _credentials_iter(
-        self,
-    ) -> Generator[Any]:
-        creds, storage_update_options, expiry = self.catalog._get_table_credentials(
-            self.table_id, write=self.write
-        )
-
-        yield storage_update_options
-
-        if not creds:
-            table_id = self.table_id
-            msg = (
-                "did not receive credentials from temporary credentials API for "
-                f"{table_id = }"
-            )
-            raise Exception(msg)  # noqa: TRY002
-
-        yield creds, expiry
-
-
-def _extract_location_and_data_format(
-    table_info: TableInfo, operation: str
-) -> tuple[str, DataSourceFormat]:
-    if table_info.storage_location is None:
-        msg = f"cannot {operation}: no storage_location found"
-        raise ValueError(msg)
-
-    if table_info.data_source_format is None:
-        msg = f"cannot {operation}: no data_source_format found"
-        raise ValueError(msg)
-
-    return table_info.storage_location, table_info.data_source_format