feat: implement vended credential refresh

Author: gabeiglio

pyiceberg/catalog/rest/__init__.py

@@ -32,6 +32,11 @@
 from pyiceberg import __version__
 from pyiceberg.catalog import BOTOCORE_SESSION, TOKEN, URI, WAREHOUSE_LOCATION, Catalog, PropertiesUpdateSummary
 from pyiceberg.catalog.rest.auth import AUTH_MANAGER, AuthManager, AuthManagerAdapter, AuthManagerFactory, LegacyOAuth2AuthManager
+from pyiceberg.catalog.rest.credential_provider import (
+    REFRESH_CREDENTIALS_ENABLED,
+    CredentialsProvider,
+    resolve_storage_credentials,
+)
 from pyiceberg.catalog.rest.response import _handle_non_200_response
 from pyiceberg.catalog.rest.scan_planning import (
     FetchScanTasksRequest,
@@ -466,26 +471,6 @@ def _create_session(self) -> Session:
 
         return session
 
-    @staticmethod
-    def _resolve_storage_credentials(storage_credentials: list[StorageCredential], location: str | None) -> Properties:
-        """Resolve the best-matching storage credential by longest prefix match.
-
-        Mirrors the Java implementation in S3FileIO.clientForStoragePath() which iterates
-        over storage credential prefixes and selects the one with the longest match.
-
-        See: https://github.com/apache/iceberg/blob/main/aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIO.java
-        """
-        if not storage_credentials or not location:
-            return {}
-
-        best_match: StorageCredential | None = None
-        for cred in storage_credentials:
-            if location.startswith(cred.prefix):
-                if best_match is None or len(cred.prefix) > len(best_match.prefix):
-                    best_match = cred
-
-        return best_match.config if best_match else {}
-
     def _load_file_io(self, properties: Properties = EMPTY_DICT, location: str | None = None) -> FileIO:
         merged_properties = {**self.properties, **properties}
         if self._auth_manager:
@@ -827,37 +812,50 @@ def add_headers(self, request: PreparedRequest, **kwargs: Any) -> None:  # pylin
 
     def _response_to_table(self, identifier_tuple: tuple[str, ...], table_response: TableResponse) -> Table:
         # Per Iceberg spec: storage-credentials take precedence over config
-        credential_config = self._resolve_storage_credentials(
-            table_response.storage_credentials, table_response.metadata_location
+        credential_config = resolve_storage_credentials(table_response.storage_credentials, table_response.metadata_location)
+        io = self._load_file_io(
+            {**table_response.metadata.properties, **table_response.config, **credential_config},
+            table_response.metadata_location,
         )
+        self._attach_credentials_provider(io, identifier_tuple, table_response.storage_credentials)
         return Table(
             identifier=identifier_tuple,
             metadata_location=table_response.metadata_location,  # type: ignore
             metadata=table_response.metadata,
-            io=self._load_file_io(
-                {**table_response.metadata.properties, **table_response.config, **credential_config},
-                table_response.metadata_location,
-            ),
+            io=io,
             catalog=self,
             config=table_response.config,
         )
 
     def _response_to_staged_table(self, identifier_tuple: tuple[str, ...], table_response: TableResponse) -> StagedTable:
         # Per Iceberg spec: storage-credentials take precedence over config
-        credential_config = self._resolve_storage_credentials(
-            table_response.storage_credentials, table_response.metadata_location
+        credential_config = resolve_storage_credentials(table_response.storage_credentials, table_response.metadata_location)
+        io = self._load_file_io(
+            {**table_response.metadata.properties, **table_response.config, **credential_config},
+            table_response.metadata_location,
         )
+        self._attach_credentials_provider(io, identifier_tuple, table_response.storage_credentials)
         return StagedTable(
             identifier=identifier_tuple,
             metadata_location=table_response.metadata_location,  # type: ignore
             metadata=table_response.metadata,
-            io=self._load_file_io(
-                {**table_response.metadata.properties, **table_response.config, **credential_config},
-                table_response.metadata_location,
-            ),
+            io=io,
             catalog=self,
         )
 
+    def _attach_credentials_provider(
+        self, io: FileIO, identifier: str | Identifier, storage_credentials: list[StorageCredential]
+    ) -> None:
+        """Attach a CredentialsProvider to io if credential refresh is enabled and credentials were vended.
+
+        The refresh callback returns the full LoadCredentialsResponse so the provider can re-run
+        longest-prefix matching against the freshly vended credentials.
+        """
+        if storage_credentials and property_as_bool(self.properties, REFRESH_CREDENTIALS_ENABLED, False):
+            io.set_credentials_provider(
+                CredentialsProvider(storage_credentials, refresh_fn=lambda: self._load_credentials(identifier))
+            )
+
     def _response_to_view(self, identifier_tuple: tuple[str, ...], view_response: ViewResponse) -> View:
         return View(
             identifier=identifier_tuple,
@@ -1124,7 +1122,7 @@ def load_credentials(
     ) -> Properties:
         """Load vended storage credentials and return the best match for a location."""
         credentials_response = self._load_credentials(identifier)
-        return self._resolve_storage_credentials(credentials_response.storage_credentials, location)
+        return resolve_storage_credentials(credentials_response.storage_credentials, location)
 
     @retry(**_RETRY_ARGS)
     @override

pyiceberg/catalog/rest/credential_provider.py

@@ -0,0 +1,112 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable
+from datetime import datetime
+from typing import TYPE_CHECKING
+from urllib.parse import urlparse
+
+from pyiceberg.catalog.rest.scan_planning import StorageCredential
+from pyiceberg.io import S3_SESSION_TOKEN_EXPIRES_AT_MS
+from pyiceberg.typedef import Properties
+from pyiceberg.utils.properties import get_first_property_value
+
+if TYPE_CHECKING:
+    from pyiceberg.catalog.rest import LoadCredentialsResponse
+
+REFRESH_CREDENTIALS_ENABLED = "client.refresh-credentials-enabled"
+
+
+def is_s3_credential_expired(config: Properties, threshold_seconds: int = 300) -> bool:
+    """Return True if the S3 session token expires within threshold_seconds (5 mins)."""
+    if expiry := get_first_property_value(config, S3_SESSION_TOKEN_EXPIRES_AT_MS):
+        expires_at = datetime.fromtimestamp(int(expiry) / 1000)
+        seconds_remaining = (expires_at - datetime.now()).total_seconds()
+        return seconds_remaining < threshold_seconds
+    return False
+
+
+# Per-scheme hooks for detecting whether a resolved credential needs to be refreshed.
+# Other schemes (e.g. gs, abfss) can register here later.
+NEEDS_REFRESH_BY_SCHEME: dict[str, Callable[[Properties], bool]] = {
+    "s3": is_s3_credential_expired,
+    "s3a": is_s3_credential_expired,
+    "s3n": is_s3_credential_expired,
+}
+
+
+def resolve_storage_credentials(storage_credentials: list[StorageCredential], location: str | None) -> Properties:
+    """Resolve the best-matching storage credential by longest prefix match.
+
+    Mirrors the Java implementation in S3FileIO.clientForStoragePath() which iterates
+    over storage credential prefixes and selects the one with the longest match.
+
+    See: https://github.com/apache/iceberg/blob/main/aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIO.java
+    """
+    if not storage_credentials or not location:
+        return {}
+
+    best_match: StorageCredential | None = None
+    for cred in storage_credentials:
+        if location.startswith(cred.prefix):
+            if best_match is None or len(cred.prefix) > len(best_match.prefix):
+                best_match = cred
+
+    return best_match.config if best_match else {}
+
+
+class CredentialsProvider:
+    """Vended-credential refresh and location-based lookup for a REST catalog table."""
+
+    _storage_credentials: list[StorageCredential]
+    _refresh_fn: Callable[[], LoadCredentialsResponse]
+    _needs_refresh_by_scheme: dict[str, Callable[[Properties], bool]]
+    _lock: threading.Lock
+
+    def __init__(
+        self,
+        storage_credentials: list[StorageCredential],
+        refresh_fn: Callable[[], LoadCredentialsResponse],
+        needs_refresh_by_scheme: dict[str, Callable[[Properties], bool]] | None = None,
+    ):
+        self._storage_credentials = storage_credentials
+        self._refresh_fn = refresh_fn
+        self._needs_refresh_by_scheme = (
+            needs_refresh_by_scheme if needs_refresh_by_scheme is not None else NEEDS_REFRESH_BY_SCHEME
+        )
+        self._lock = threading.Lock()
+
+    def _can_refresh(self, location: str) -> bool:
+        scheme = urlparse(location).scheme
+        refresh_by_scheme = self._needs_refresh_by_scheme.get(scheme)
+        config = resolve_storage_credentials(self._storage_credentials, location)
+        return config != {} and refresh_by_scheme is not None and refresh_by_scheme(config)
+
+    def properties_for(self, location: str) -> Properties:
+        """Return the credential properties that apply to the given location, refreshing if needed."""
+        config = resolve_storage_credentials(self._storage_credentials, location)
+
+        if self._can_refresh(location):
+            with self._lock:
+                if self._can_refresh(location):
+                    response = self._refresh_fn()
+                    self._storage_credentials = response.storage_credentials
+                config = resolve_storage_credentials(self._storage_credentials, location)
+
+        return config

pyiceberg/io/__init__.py

@@ -41,6 +41,16 @@
 
 logger = logging.getLogger(__name__)
 
+
+@runtime_checkable
+class CredentialsProviderProtocol(Protocol):
+    """Protocol for objects that can resolve credential properties for a file location."""
+
+    def properties_for(self, location: str) -> Properties:
+        """Return the credential properties that apply to the given location."""
+        ...
+
+
 AWS_PROFILE_NAME = "client.profile-name"
 AWS_REGION = "client.region"
 AWS_ACCESS_KEY_ID = "client.access-key-id"
@@ -54,6 +64,7 @@
 S3_ACCESS_KEY_ID = "s3.access-key-id"
 S3_SECRET_ACCESS_KEY = "s3.secret-access-key"
 S3_SESSION_TOKEN = "s3.session-token"
+S3_SESSION_TOKEN_EXPIRES_AT_MS = "s3.session-token-expires-at-ms"
 S3_REGION = "s3.region"
 S3_RESOLVE_REGION = "s3.resolve-region"
 S3_PROXY_URI = "s3.proxy-uri"
@@ -258,6 +269,7 @@ class FileIO(ABC):
     """A base class for FileIO implementations."""
 
     properties: Properties
+    _credentials_provider: CredentialsProviderProtocol | None = None
 
     def __init__(self, properties: Properties = EMPTY_DICT):
         self.properties = properties
@@ -291,6 +303,18 @@ def delete(self, location: str | InputFile | OutputFile) -> None:
             FileNotFoundError: When the file at the provided location does not exist.
         """
 
+    def set_credentials_provider(self, provider: CredentialsProviderProtocol) -> None:
+        """Inject a credentials provider for refreshing vended storage credentials.
+
+        Backends that support credential refresh (e.g. S3) consult the provider at file-access
+        time and rebuild their underlying filesystem when credentials change. Backends that do
+        not support refresh simply hold the reference without using it.
+
+        Args:
+            provider (CredentialsProviderProtocol): Resolves credential properties for a file location.
+        """
+        self._credentials_provider = provider
+
 
 LOCATION = "location"
 WAREHOUSE = "warehouse"

pyiceberg/io/fsspec.py

@@ -443,7 +443,7 @@ def new_input(self, location: str) -> FsspecInputFile:
             FsspecInputFile: An FsspecInputFile instance for the given location.
         """
         uri = urlparse(location)
-        fs = self._get_fs_from_uri(uri)
+        fs = self._get_fs_from_uri(uri, location)
         return FsspecInputFile(location=location, fs=fs)
 
     @override
@@ -457,7 +457,7 @@ def new_output(self, location: str) -> FsspecOutputFile:
             FsspecOutputFile: An FsspecOutputFile instance for the given location.
         """
         uri = urlparse(location)
-        fs = self._get_fs_from_uri(uri)
+        fs = self._get_fs_from_uri(uri, location)
         return FsspecOutputFile(location=location, fs=fs)
 
     @override
@@ -475,31 +475,44 @@ def delete(self, location: str | InputFile | OutputFile) -> None:
             str_location = location
 
         uri = urlparse(str_location)
-        fs = self._get_fs_from_uri(uri)
+        fs = self._get_fs_from_uri(uri, str_location)
         fs.rm(str_location)
 
-    def _get_fs_from_uri(self, uri: "ParseResult") -> AbstractFileSystem:
-        """Get a filesystem from a parsed URI, using hostname for ADLS account resolution."""
+    def _get_fs_from_uri(self, uri: "ParseResult", location: str) -> AbstractFileSystem:
+        """Get a filesystem from a parsed URI, using hostname for ADLS account resolution.
+
+        When a credentials provider is attached, its resolved (and possibly refreshed) credential
+        properties are folded into the cache key so rotated credentials build a new filesystem.
+        """
+        creds_key: frozenset[tuple[str, str]] = frozenset()
+        if provider := self._credentials_provider:
+            creds_key = frozenset(provider.properties_for(location).items())
         if uri.scheme in _ADLS_SCHEMES:
-            return self.get_fs(uri.scheme, uri.hostname)
-        return self.get_fs(uri.scheme)
+            return self.get_fs(uri.scheme, uri.hostname, creds_key)
+        return self.get_fs(uri.scheme, None, creds_key)
 
-    def get_fs(self, scheme: str, hostname: str | None = None) -> AbstractFileSystem:
+    def get_fs(
+        self, scheme: str, hostname: str | None = None, creds_key: frozenset[tuple[str, str]] = frozenset()
+    ) -> AbstractFileSystem:
         """Get a filesystem for a specific scheme, cached per thread."""
         if not hasattr(self._thread_locals, "get_fs_cached"):
             self._thread_locals.get_fs_cached = lru_cache(self._get_fs)
 
-        return self._thread_locals.get_fs_cached(scheme, hostname)
+        return self._thread_locals.get_fs_cached(scheme, hostname, creds_key)
 
-    def _get_fs(self, scheme: str, hostname: str | None = None) -> AbstractFileSystem:
+    def _get_fs(
+        self, scheme: str, hostname: str | None = None, creds_key: frozenset[tuple[str, str]] = frozenset()
+    ) -> AbstractFileSystem:
         """Get a filesystem for a specific scheme."""
         if scheme not in self._scheme_to_fs:
             raise ValueError(f"No registered filesystem for scheme: {scheme}")
 
+        properties = {**self.properties, **dict(creds_key)}
+
         if scheme in _ADLS_SCHEMES:
-            return _adls(self.properties, hostname)
+            return _adls(properties, hostname)
 
-        return self._scheme_to_fs[scheme](self.properties)
+        return self._scheme_to_fs[scheme](properties)
 
     def __getstate__(self) -> dict[str, Any]:
         """Create a dictionary of the FsSpecFileIO fields used when pickling."""