polars-runtime-compat 1.34.0b2__cp39-abi3-win_arm64.whl

polars/io/cloud/_utils.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Generic, TypeVar
+
+from polars._typing import PartitioningScheme
+from polars._utils.various import is_path_or_str_sequence
+
+T = TypeVar("T")
+
+
+class NoPickleOption(Generic[T]):
+    """
+    Wrapper that does not pickle the wrapped value.
+
+    This wrapper will unpickle to contain a None. Used for cached values.
+    """
+
+    def __init__(self, opt_value: T | None = None) -> None:
+        self._opt_value = opt_value
+
+    def get(self) -> T | None:
+        return self._opt_value
+
+    def set(self, value: T | None) -> None:
+        self._opt_value = value
+
+    def __getstate__(self) -> tuple[()]:
+        # Needs to return not-None for `__setstate__()` to be called
+        return ()
+
+    def __setstate__(self, _state: tuple[()]) -> None:
+        NoPickleOption.__init__(self)
+
+
+def _first_scan_path(
+    source: Any,
+) -> str | Path | None:
+    if isinstance(source, (str, Path)):
+        return source
+    elif is_path_or_str_sequence(source) and source:
+        return source[0]
+    elif isinstance(source, PartitioningScheme):
+        return source._base_path
+
+    return None
+
+
+def _get_path_scheme(path: str | Path) -> str | None:
+    path_str = str(path)
+    i = path_str.find("://")
+
+    return path_str[:i] if i >= 0 else None
+
+
+def _is_aws_cloud(*, scheme: str, first_scan_path: str) -> bool:
+    if any(scheme == x for x in ["s3", "s3a"]):
+        return True
+
+    if scheme == "http" or scheme == "https":
+        bucket_end = first_scan_path.find(".s3.")
+        region_end = first_scan_path.find(".amazonaws.com/", bucket_end + 4)
+
+        if (
+            first_scan_path.find("/", len(scheme) + 3, region_end) > 0
+            or "?" in first_scan_path
+        ):
+            return False
+
+        return 0 < bucket_end < region_end
+
+    return False
+
+
+def _is_azure_cloud(scheme: str) -> bool:
+    return any(scheme == x for x in ["az", "azure", "adl", "abfs", "abfss"])
+
+
+def _is_gcp_cloud(scheme: str) -> bool:
+    return any(scheme == x for x in ["gs", "gcp", "gcs"])

polars/io/cloud/credential_provider/__init__.py

@@ -0,0 +1,17 @@
+from polars.io.cloud.credential_provider._providers import (
+    CredentialProvider,
+    CredentialProviderAWS,
+    CredentialProviderAzure,
+    CredentialProviderFunction,
+    CredentialProviderFunctionReturn,
+    CredentialProviderGCP,
+)
+
+__all__ = [
+    "CredentialProvider",
+    "CredentialProviderAWS",
+    "CredentialProviderAzure",
+    "CredentialProviderFunction",
+    "CredentialProviderFunctionReturn",
+    "CredentialProviderGCP",
+]

polars/io/cloud/credential_provider/_builder.py

@@ -0,0 +1,520 @@
+from __future__ import annotations
+
+import abc
+import os
+import threading
+from typing import TYPE_CHECKING, Any, Callable, Literal, Union
+
+import polars._utils.logging
+from polars._utils.cache import LRUCache
+from polars._utils.logging import eprint, verbose
+from polars._utils.unstable import issue_unstable_warning
+from polars.io.cloud._utils import NoPickleOption
+from polars.io.cloud.credential_provider._providers import (
+    CachedCredentialProvider,
+    CachingCredentialProvider,
+    CredentialProvider,
+    CredentialProviderAWS,
+    CredentialProviderAzure,
+    CredentialProviderFunction,
+    CredentialProviderGCP,
+    UserProvidedGCPToken,
+)
+
+if TYPE_CHECKING:
+    import sys
+
+    if sys.version_info >= (3, 10):
+        from typing import TypeAlias
+    else:
+        from typing_extensions import TypeAlias
+
+# https://docs.rs/object_store/latest/object_store/enum.ClientConfigKey.html
+OBJECT_STORE_CLIENT_OPTIONS: frozenset[str] = frozenset(
+    [
+        "allow_http",
+        "allow_invalid_certificates",
+        "connect_timeout",
+        "default_content_type",
+        "http1_only",
+        "http2_only",
+        "http2_keep_alive_interval",
+        "http2_keep_alive_timeout",
+        "http2_keep_alive_while_idle",
+        "http2_max_frame_size",
+        "pool_idle_timeout",
+        "pool_max_idle_per_host",
+        "proxy_url",
+        "proxy_ca_certificate",
+        "proxy_excludes",
+        "timeout",
+        "user_agent",
+    ]
+)
+
+CredentialProviderBuilderReturn: TypeAlias = Union[
+    CredentialProvider, CredentialProviderFunction, None
+]
+
+
+class CredentialProviderBuilder:
+    """
+    Builds credential providers.
+
+    This is used to defer credential provider initialization to happen at
+    `collect()` rather than immediately during query construction. This makes
+    the behavior predictable when queries are sent to another environment for
+    execution.
+    """
+
+    def __init__(
+        self,
+        credential_provider_init: CredentialProviderBuilderImpl,
+    ) -> None:
+        """
+        Initialize configuration for building a credential provider.
+
+        Parameters
+        ----------
+        credential_provider_init
+            Initializer function that returns a credential provider.
+        """
+        self.credential_provider_init = credential_provider_init
+
+    # Note: The rust-side expects this exact function name.
+    def build_credential_provider(
+        self,
+        clear_cached_credentials: bool = False,  # noqa: FBT001
+    ) -> CredentialProviderBuilderReturn:
+        """
+        Instantiate a credential provider from configuration.
+
+        Parameters
+        ----------
+        clear_cached_credentials
+            If the built provider is an instance of `CachingCredentialProvider`,
+            clears any cached credentials on that object.
+        """
+        verbose = polars._utils.logging.verbose()
+
+        if verbose:
+            eprint(
+                "[CredentialProviderBuilder]: Begin initialize "
+                f"{self.credential_provider_init!r} "
+                f"{clear_cached_credentials = }"
+            )
+
+        v = self.credential_provider_init()
+
+        if verbose:
+            if v is not None:
+                eprint(
+                    f"[CredentialProviderBuilder]: Initialized {v!r} "
+                    f"from {self.credential_provider_init!r}"
+                )
+            else:
+                eprint(
+                    f"[CredentialProviderBuilder]: No provider initialized "
+                    f"from {self.credential_provider_init!r}"
+                )
+
+        if clear_cached_credentials and isinstance(v, CachingCredentialProvider):
+            v.clear_cached_credentials()
+
+            if verbose:
+                eprint(
+                    f"[CredentialProviderBuilder]: Clear cached credentials for {v!r}"
+                )
+
+        return v
+
+    @classmethod
+    def from_initialized_provider(
+        cls, credential_provider: CredentialProviderFunction
+    ) -> CredentialProviderBuilder:
+        """Initialize with an already constructed provider."""
+        return cls(InitializedCredentialProvider(credential_provider))
+
+    def __getstate__(self) -> Any:
+        state = self.credential_provider_init
+
+        if verbose():
+            eprint(f"[CredentialProviderBuilder]: __getstate__(): {state = !r} ")
+
+        return state
+
+    def __setstate__(self, state: Any) -> None:
+        self.credential_provider_init = state
+
+        if verbose():
+            eprint(f"[CredentialProviderBuilder]: __setstate__(): {self = !r}")
+
+    def __repr__(self) -> str:
+        return f"CredentialProviderBuilder({self.credential_provider_init!r})"
+
+
+class CredentialProviderBuilderImpl(abc.ABC):
+    @abc.abstractmethod
+    def __call__(self) -> CredentialProviderFunction | None:
+        pass
+
+    @property
+    @abc.abstractmethod
+    def provider_repr(self) -> str:
+        """Used for logging."""
+
+    def __repr__(self) -> str:
+        provider_repr = self.provider_repr
+        builder_name = type(self).__name__
+
+        return f"{provider_repr} @ {builder_name}"
+
+
+# Wraps an already initialized credential provider into the builder interface.
+# Used for e.g. user-provided credential providers.
+class InitializedCredentialProvider(CredentialProviderBuilderImpl):
+    """Wraps an already initialized credential provider."""
+
+    def __init__(self, credential_provider: CredentialProviderFunction) -> None:
+        self.credential_provider = credential_provider
+
+    def __call__(self) -> CredentialProviderBuilderReturn:
+        if isinstance(self.credential_provider, CachingCredentialProvider):
+            return self.credential_provider
+
+        # We use the cache by keying the entry as the address of the object
+        # provided by the user.
+        return _build_with_cache(
+            lambda: id(self.credential_provider),
+            lambda: CachedCredentialProvider(self.credential_provider),
+        )
+
+    @property
+    def provider_repr(self) -> str:
+        return repr(self.credential_provider)
+
+
+# The keys of this can be:
+# * int: Object address of a user-passed credential provider
+# * bytes: Hash of an AutoInit configuration
+BUILT_PROVIDERS_LRU_CACHE: (
+    LRUCache[int | bytes, CredentialProviderBuilderReturn] | None
+) = None
+BUILT_PROVIDERS_LRU_CACHE_LOCK: threading.RLock = threading.RLock()
+
+
+def _build_with_cache(
+    get_cache_key_func: Callable[[], int | bytes],
+    build_provider_func: Callable[[], CredentialProviderBuilderReturn],
+) -> CredentialProviderBuilderReturn:
+    global BUILT_PROVIDERS_LRU_CACHE
+
+    if (
+        max_items := int(
+            os.getenv(
+                "POLARS_CREDENTIAL_PROVIDER_BUILDER_CACHE_SIZE",
+                8,
+            )
+        )
+    ) <= 0:
+        if BUILT_PROVIDERS_LRU_CACHE_LOCK.acquire(blocking=False):
+            BUILT_PROVIDERS_LRU_CACHE = None
+            BUILT_PROVIDERS_LRU_CACHE_LOCK.release()
+
+        return build_provider_func()
+
+    verbose = polars._utils.logging.verbose()
+
+    with BUILT_PROVIDERS_LRU_CACHE_LOCK:
+        if BUILT_PROVIDERS_LRU_CACHE is None:
+            if verbose:
+                eprint(f"Create built credential providers LRU cache ({max_items = })")
+
+            BUILT_PROVIDERS_LRU_CACHE = LRUCache(max_items)
+
+        cache_key = get_cache_key_func()
+
+        try:
+            provider = BUILT_PROVIDERS_LRU_CACHE[cache_key]
+
+            if verbose:
+                eprint(
+                    f"Loaded credential provider from cache: {provider!r} {cache_key = }"
+                )
+        except KeyError:
+            provider = build_provider_func()
+            BUILT_PROVIDERS_LRU_CACHE[cache_key] = provider
+
+            if verbose:
+                eprint(
+                    f"Added new credential provider to cache: {provider!r} {cache_key = }"
+                )
+
+        return provider
+
+
+# Represents an automatic initialization configuration. This is created for
+# credential_provider="auto".
+class AutoInit(CredentialProviderBuilderImpl):
+    def __init__(self, cls: Any, **kw: Any) -> None:
+        self.cls = cls
+        self.kw = kw
+        self._cache_key: NoPickleOption[bytes] = NoPickleOption()
+
+    def __call__(self) -> CredentialProviderFunction | None:
+        # This is used for credential_provider="auto", which allows for
+        # ImportErrors.
+        try:
+            return _build_with_cache(
+                self.get_or_init_cache_key,
+                lambda: self.cls(**self.kw),
+            )
+        except ImportError as e:
+            if verbose():
+                eprint(f"failed to auto-initialize {self.provider_repr}: {e!r}")
+
+        return None
+
+    def get_or_init_cache_key(self) -> bytes:
+        cache_key = self._cache_key.get()
+
+        if cache_key is None:
+            cache_key = self.get_cache_key_impl()
+            self._cache_key.set(cache_key)
+
+            if verbose():
+                eprint(f"{self!r}: AutoInit cache key: {cache_key.hex()}")
+
+        return cache_key
+
+    def get_cache_key_impl(self) -> bytes:
+        import hashlib
+        import pickle
+
+        hash = hashlib.sha256(pickle.dumps(self))
+        return hash.digest()[:16]
+
+    @property
+    def provider_repr(self) -> str:
+        return self.cls.__name__
+
+
+DEFAULT_CREDENTIAL_PROVIDER: CredentialProviderFunction | Literal["auto"] | None = (
+    "auto"
+)
+
+
+def _init_credential_provider_builder(
+    credential_provider: CredentialProviderFunction
+    | CredentialProviderBuilder
+    | Literal["auto"]
+    | None,
+    source: Any,
+    storage_options: dict[str, Any] | None,
+    caller_name: str,
+) -> CredentialProviderBuilder | None:
+    def f() -> CredentialProviderBuilder | None:
+        # Note: The behavior of this function should depend only on the function
+        # parameters. Any environment-specific behavior should take place inside
+        # instantiated credential providers.
+
+        from polars.io.cloud._utils import (
+            _first_scan_path,
+            _get_path_scheme,
+            _is_aws_cloud,
+            _is_azure_cloud,
+            _is_gcp_cloud,
+        )
+
+        if credential_provider is None:
+            return None
+
+        if isinstance(credential_provider, CredentialProviderBuilder):
+            # This happens when the catalog client auto-inits and passes it to
+            # scan/write_delta, which calls us again.
+            return credential_provider
+
+        if credential_provider != "auto":
+            msg = f"the `credential_provider` parameter of `{caller_name}` is considered unstable."
+            issue_unstable_warning(msg)
+
+            return CredentialProviderBuilder.from_initialized_provider(
+                credential_provider
+            )
+
+        if DEFAULT_CREDENTIAL_PROVIDER is None:
+            return None
+
+        if (first_scan_path := _first_scan_path(source)) is None:
+            return None
+
+        if (scheme := _get_path_scheme(first_scan_path)) is None:
+            return None
+
+        def get_default_credential_provider() -> CredentialProviderBuilder | None:
+            return (
+                CredentialProviderBuilder.from_initialized_provider(
+                    DEFAULT_CREDENTIAL_PROVIDER
+                )
+                if DEFAULT_CREDENTIAL_PROVIDER != "auto"
+                else None
+            )
+
+        if _is_azure_cloud(scheme):
+            tenant_id = None
+            storage_account = None
+
+            if storage_options is not None:
+                for k, v in storage_options.items():
+                    k = k.lower()
+
+                    # https://docs.rs/object_store/latest/object_store/azure/enum.AzureConfigKey.html
+                    if k in {
+                        "azure_storage_tenant_id",
+                        "azure_storage_authority_id",
+                        "azure_tenant_id",
+                        "azure_authority_id",
+                        "tenant_id",
+                        "authority_id",
+                    }:
+                        tenant_id = v
+                    elif k in {"azure_storage_account_name", "account_name"}:
+                        storage_account = v
+                    elif k in {"azure_use_azure_cli", "use_azure_cli"}:
+                        continue
+                    elif k in OBJECT_STORE_CLIENT_OPTIONS:
+                        continue
+                    else:
+                        # We assume some sort of access key was given, so we
+                        # just dispatch to the rust side.
+                        return None
+
+            storage_account = (
+                # Prefer the one embedded in the path
+                CredentialProviderAzure._extract_adls_uri_storage_account(
+                    str(first_scan_path)
+                )
+                or storage_account
+            )
+
+            if (default := get_default_credential_provider()) is not None:
+                return default
+
+            return CredentialProviderBuilder(
+                AutoInit(
+                    CredentialProviderAzure,
+                    tenant_id=tenant_id,
+                    _storage_account=storage_account,
+                )
+            )
+
+        elif _is_aws_cloud(scheme=scheme, first_scan_path=str(first_scan_path)):
+            region = None
+            profile = None
+            default_region = None
+            unhandled_key = None
+            has_endpoint_url = False
+
+            if storage_options is not None:
+                for k, v in storage_options.items():
+                    k = k.lower()
+
+                    # https://docs.rs/object_store/latest/object_store/aws/enum.AmazonS3ConfigKey.html
+                    if k in {"aws_region", "region"}:
+                        region = v
+                    elif k in {"aws_default_region", "default_region"}:
+                        default_region = v
+                    elif k in {"aws_profile", "profile"}:
+                        profile = v
+                    elif k in {
+                        "aws_endpoint",
+                        "aws_endpoint_url",
+                        "endpoint",
+                        "endpoint_url",
+                    }:
+                        has_endpoint_url = True
+                    elif k in {"aws_request_payer", "request_payer"}:
+                        continue
+                    elif k in OBJECT_STORE_CLIENT_OPTIONS:
+                        continue
+                    else:
+                        # We assume this is some sort of access key
+                        unhandled_key = k
+
+            if unhandled_key is not None:
+                if profile is not None:
+                    msg = (
+                        "unsupported: cannot combine aws_profile with "
+                        f"{unhandled_key} in storage_options"
+                    )
+                    raise ValueError(msg)
+
+            if (
+                unhandled_key is None
+                and (default := get_default_credential_provider()) is not None
+            ):
+                return default
+
+            return CredentialProviderBuilder(
+                AutoInit(
+                    CredentialProviderAWS,
+                    profile_name=profile,
+                    region_name=region or default_region,
+                    _auto_init_unhandled_key=unhandled_key,
+                    _storage_options_has_endpoint_url=has_endpoint_url,
+                )
+            )
+
+        elif _is_gcp_cloud(scheme):
+            token = None
+            unhandled_key = None
+
+            if storage_options is not None:
+                for k, v in storage_options.items():
+                    k = k.lower()
+
+                    # https://docs.rs/object_store/latest/object_store/gcp/enum.GoogleConfigKey.html
+                    if k in {"token", "bearer_token"}:
+                        token = v
+                    elif k in {
+                        "google_bucket",
+                        "google_bucket_name",
+                        "bucket",
+                        "bucket_name",
+                    }:
+                        continue
+                    elif k in OBJECT_STORE_CLIENT_OPTIONS:
+                        continue
+                    else:
+                        # We assume some sort of access key was given, so we
+                        # just dispatch to the rust side.
+                        unhandled_key = k
+
+            if unhandled_key is not None:
+                if token is not None:
+                    msg = (
+                        "unsupported: cannot combine token with "
+                        f"{unhandled_key} in storage_options"
+                    )
+                    raise ValueError(msg)
+
+                return None
+
+            if token is not None:
+                return CredentialProviderBuilder(
+                    InitializedCredentialProvider(UserProvidedGCPToken(token))
+                )
+
+            if (default := get_default_credential_provider()) is not None:
+                return default
+
+            return CredentialProviderBuilder(AutoInit(CredentialProviderGCP))
+
+        return None
+
+    credential_provider_init = f()
+
+    if verbose():
+        eprint(f"_init_credential_provider_builder(): {credential_provider_init = !r}")
+
+    return credential_provider_init