feat(python): Add Python-side caching for credentials and provider auto-initialization (#23736)

Author: nameexhaustion

py-polars/polars/io/cloud/_utils.py

@@ -1,11 +1,53 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Any
+from typing import Any, Generic, TypeVar
 
 from polars._utils.various import is_path_or_str_sequence
 from polars.io.partition import PartitionMaxSize
 
+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)
+
+
+class ZeroHashWrap(Generic[T]):
+    """Wrapper that always hashes to 0 and always returns True for __eq__."""
+
+    def __init__(self, value: T) -> None:
+        self._value = value
+
+    def get(self) -> T:
+        return self._value
+
+    def __eq__(self, _other: object) -> bool:
+        return True
+
+    def __hash__(self) -> int:
+        return 0
+
 
 def _first_scan_path(
     source: Any,

py-polars/polars/io/cloud/credential_provider/_builder.py

@@ -1,23 +1,30 @@
 from __future__ import annotations
 
 import abc
-from typing import TYPE_CHECKING, Any, Literal
+import os
+from functools import lru_cache
+from typing import TYPE_CHECKING, Any, Callable, Literal, Union
 
 import polars._utils.logging
 from polars._utils.logging import eprint, verbose
 from polars._utils.unstable import issue_unstable_warning
+from polars.io.cloud._utils import NoPickleOption, ZeroHashWrap
 from polars.io.cloud.credential_provider._providers import (
     CredentialProvider,
     CredentialProviderAWS,
     CredentialProviderAzure,
+    CredentialProviderFunction,
     CredentialProviderGCP,
+    UserProvidedGCPToken,
 )
 
 if TYPE_CHECKING:
-    from polars.io.cloud.credential_provider._providers import (
-        CredentialProviderFunction,
-        CredentialProviderFunctionReturn,
-    )
+    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(
@@ -42,6 +49,10 @@
     ]
 )
 
+CredentialProviderBuilderReturn: TypeAlias = Union[
+    CredentialProvider, CredentialProviderFunction, None
+]
+
 
 class CredentialProviderBuilder:
     """
@@ -70,7 +81,7 @@ def __init__(
     # Note: The rust-side expects this exact function name.
     def build_credential_provider(
         self,
-    ) -> CredentialProvider | CredentialProviderFunction | None:
+    ) -> CredentialProviderBuilderReturn:
         """Instantiate a credential provider from configuration."""
         verbose = polars._utils.logging.verbose()
 
@@ -154,38 +165,91 @@ def provider_repr(self) -> str:
         return repr(self.credential_provider)
 
 
+AUTO_INIT_LRU_CACHE: (
+    Callable[
+        [bytes, ZeroHashWrap[Callable[[], CredentialProviderBuilderReturn]]],
+        CredentialProviderBuilderReturn,
+    ]
+    | None
+) = None
+
+
+def _auto_init_with_cache(
+    get_cache_key_func: Callable[[], bytes],
+    build_provider_func: ZeroHashWrap[Callable[[], CredentialProviderBuilderReturn]],
+) -> CredentialProviderBuilderReturn:
+    global AUTO_INIT_LRU_CACHE
+
+    if (
+        maxsize := int(os.getenv("POLARS_CREDENTIAL_PROVIDER_BUILDER_CACHE_SIZE", 8))
+    ) <= 0:
+        AUTO_INIT_LRU_CACHE = None
+
+        return build_provider_func.get()()
+
+    if AUTO_INIT_LRU_CACHE is None:
+        if verbose():
+            eprint(f"Create credential provider AutoInit LRU cache ({maxsize = })")
+
+        @lru_cache(maxsize=maxsize)
+        def cache(
+            _cache_key: bytes,
+            build_provider_func: ZeroHashWrap[
+                Callable[[], CredentialProviderBuilderReturn]
+            ],
+        ) -> CredentialProviderBuilderReturn:
+            return build_provider_func.get()()
+
+        AUTO_INIT_LRU_CACHE = cache
+
+    return AUTO_INIT_LRU_CACHE(
+        get_cache_key_func(),
+        build_provider_func,
+    )
+
+
 # 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) -> Any:
+    def __call__(self) -> CredentialProviderFunction | None:
         # This is used for credential_provider="auto", which allows for
         # ImportErrors.
         try:
-            return self.cls(**self.kw)
+            return _auto_init_with_cache(
+                self.get_or_init_cache_key,
+                ZeroHashWrap(lambda: self.cls(**self.kw)),
+            )
         except ImportError as e:
             if verbose():
                 eprint(f"failed to auto-initialize {self.provider_repr}: {e!r}")
 
         return None
 
-    @property
-    def provider_repr(self) -> str:
-        return self.cls.__name__
+    def get_or_init_cache_key(self) -> bytes:
+        cache_key = self._cache_key.get()
 
+        if cache_key is None:
+            import hashlib
+            import pickle
 
-class UserProvidedGCPToken(CredentialProvider):
-    """User-provided GCP token in storage_options."""
+            hash = hashlib.sha256(pickle.dumps(self))
+            self._cache_key.set(hash.digest())
+            cache_key = self._cache_key.get()
+            assert isinstance(cache_key, bytes)
 
-    def __init__(self, token: str) -> None:
-        self.token = token
+            if verbose():
+                eprint(f"{self!r}: AutoInit cache key: {hash.hexdigest()}")
+
+        return cache_key
 
-    def __call__(self) -> CredentialProviderFunctionReturn:
-        """Fetches the credentials."""
-        return {"bearer_token": self.token}, None
+    @property
+    def provider_repr(self) -> str:
+        return self.cls.__name__
 
 
 def _init_credential_provider_builder(

py-polars/polars/io/cloud/credential_provider/_providers.py

@@ -8,10 +8,18 @@
 import sys
 import zoneinfo
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Callable, Optional, TypedDict, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Optional,
+    TypedDict,
+    Union,
+)
 
 import polars._utils.logging
 from polars._utils.logging import eprint, verbose
+from polars.io.cloud._utils import NoPickleOption
 
 if TYPE_CHECKING:
     if sys.version_info >= (3, 10):
@@ -56,9 +64,53 @@ class CredentialProvider(abc.ABC):
             at any point without it being considered a breaking change.
     """
 
-    @abc.abstractmethod
+    def __init__(self) -> None:
+        self._cached_credentials: NoPickleOption[CredentialProviderFunctionReturn] = (
+            NoPickleOption()
+        )
+        self._has_logged_use_cache = False
+
+        if verbose():
+            eprint(
+                f"[{type(self).__name__} @ {hex(id(self))}]: CredentialProvider.__init__()"
+            )
+
     def __call__(self) -> CredentialProviderFunctionReturn:
         """Fetches the credentials."""
+        if os.getenv("POLARS_DISABLE_PYTHON_CREDENTIAL_CACHING") == "1":
+            self._cached_credentials.set(None)
+            return self.retrieve_credentials_impl()
+
+        if not isinstance(getattr(self, "_cached_credentials", None), NoPickleOption):
+            msg = (
+                f"[{type(self).__name__} @ {hex(id(self))}]: `_cached_credentials` attribute "
+                "not found. This can happen if a subclass forgets to call "
+                f"super().__init__() ({type(self) = })"
+            )
+            raise AttributeError(msg)  # noqa: TRY004
+
+        cached = self._cached_credentials.get()
+
+        if cached is None or (
+            (expiry := cached[1]) is not None
+            and expiry <= int(datetime.now().timestamp())
+        ):
+            self._cached_credentials.set(self.retrieve_credentials_impl())
+            self._has_logged_use_cache = False
+            cached = self._cached_credentials.get()
+            assert cached is not None
+
+        elif verbose() and not self._has_logged_use_cache:
+            expiry = cached[1]
+            eprint(
+                f"[{type(self).__name__} @ {hex(id(self))}]: Using cached credentials ({expiry = })"
+            )
+            self._has_logged_use_cache = True
+
+        return cached
+
+    @abc.abstractmethod
+    def retrieve_credentials_impl(self) -> CredentialProviderFunctionReturn: ...
 
 
 class CredentialProviderAWS(CredentialProvider):
@@ -95,14 +147,16 @@ def __init__(
         msg = "`CredentialProviderAWS` functionality is considered unstable"
         issue_unstable_warning(msg)
 
+        super().__init__()
+
         self._ensure_module_availability()
         self.profile_name = profile_name
         self.region_name = region_name
         self.assume_role = assume_role
         self._auto_init_unhandled_key = _auto_init_unhandled_key
         self._storage_options_has_endpoint_url = _storage_options_has_endpoint_url
 
-    def __call__(self) -> CredentialProviderFunctionReturn:
+    def retrieve_credentials_impl(self) -> CredentialProviderFunctionReturn:
         """Fetch the credentials for the configured profile name."""
         assert not self._auto_init_unhandled_key
 
@@ -253,6 +307,8 @@ def __init__(
         msg = "`CredentialProviderAzure` functionality is considered unstable"
         issue_unstable_warning(msg)
 
+        super().__init__()
+
         self.account_name = _storage_account
         self.scopes = (
             scopes if scopes is not None else ["https://storage.azure.com/.default"]
@@ -283,7 +339,7 @@ def __init__(
                 f"{self.scopes = } "
             )
 
-    def __call__(self) -> CredentialProviderFunctionReturn:
+    def retrieve_credentials_impl(self) -> CredentialProviderFunctionReturn:
         """Fetch the credentials."""
         if (
             v := self._try_get_azure_storage_account_credential_if_permitted()
@@ -438,6 +494,8 @@ def __init__(
         msg = "`CredentialProviderGCP` functionality is considered unstable"
         issue_unstable_warning(msg)
 
+        super().__init__()
+
         self._ensure_module_availability()
 
         import google.auth
@@ -462,7 +520,7 @@ def __init__(
         )
         self.creds = creds
 
-    def __call__(self) -> CredentialProviderFunctionReturn:
+    def retrieve_credentials_impl(self) -> CredentialProviderFunctionReturn:
         """Fetch the credentials."""
         import google.auth.transport.requests
 
@@ -487,6 +545,20 @@ def _ensure_module_availability(cls) -> None:
             raise ImportError(msg)
 
 
+class UserProvidedGCPToken(CredentialProvider):
+    """User-provided GCP token in storage_options."""
+
+    def __init__(self, token: str) -> None:
+        self.token = token
+
+    def __call__(self) -> CredentialProviderFunctionReturn:
+        return self.retrieve_credentials_impl()
+
+    def retrieve_credentials_impl(self) -> CredentialProviderFunctionReturn:
+        """Fetches the credentials."""
+        return {"bearer_token": self.token}, None
+
+
 def _get_credentials_from_provider_expiry_aware(
     credential_provider: CredentialProviderFunction,
 ) -> dict[str, str] | None: