From 010d6bc1a3c0b514decb6a185cf47d20a38dcb9e Mon Sep 17 00:00:00 2001 From: Gabriel Igliozzi Date: Wed, 5 Aug 2026 12:20:16 +0200 Subject: [PATCH] feat: implement vended credential refresh --- pyiceberg/catalog/rest/__init__.py | 64 +++-- pyiceberg/catalog/rest/credential_provider.py | 112 ++++++++ pyiceberg/io/__init__.py | 23 ++ pyiceberg/io/fsspec.py | 31 ++- pyiceberg/io/pyarrow.py | 126 +++++---- tests/catalog/test_credential_provider.py | 248 ++++++++++++++++++ tests/catalog/test_rest.py | 86 +++++- tests/io/test_credential_refresh.py | 97 +++++++ 8 files changed, 687 insertions(+), 100 deletions(-) create mode 100644 pyiceberg/catalog/rest/credential_provider.py create mode 100644 tests/catalog/test_credential_provider.py create mode 100644 tests/io/test_credential_refresh.py diff --git a/pyiceberg/catalog/rest/__init__.py b/pyiceberg/catalog/rest/__init__.py index 97a71b429b..ddecaeaef3 100644 --- a/pyiceberg/catalog/rest/__init__.py +++ b/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 diff --git a/pyiceberg/catalog/rest/credential_provider.py b/pyiceberg/catalog/rest/credential_provider.py new file mode 100644 index 0000000000..7b0cd27ab7 --- /dev/null +++ b/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 diff --git a/pyiceberg/io/__init__.py b/pyiceberg/io/__init__.py index 6ff1674740..229eb621b8 100644 --- a/pyiceberg/io/__init__.py +++ b/pyiceberg/io/__init__.py @@ -54,6 +54,15 @@ def _is_local_path(path: str) -> bool: return drive != "" +@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" @@ -67,6 +76,7 @@ def _is_local_path(path: str) -> bool: 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" @@ -271,6 +281,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 @@ -304,6 +315,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" diff --git a/pyiceberg/io/fsspec.py b/pyiceberg/io/fsspec.py index b28ffd0699..b357eea165 100644 --- a/pyiceberg/io/fsspec.py +++ b/pyiceberg/io/fsspec.py @@ -480,29 +480,44 @@ def delete(self, location: str | InputFile | OutputFile) -> None: fs.rm(str_location) def _get_fs_from_uri(self, uri: "ParseResult", location: str = "") -> AbstractFileSystem: - """Get a filesystem from a parsed URI, using hostname for ADLS account resolution.""" + """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. + """ if _is_local_path(location): return self.get_fs("file") + + 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.""" diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index c36f1639d9..f5f75cf37b 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -394,12 +394,24 @@ def to_input_file(self) -> PyArrowFile: class PyArrowFileIO(FileIO): - fs_by_scheme: Callable[[str, str | None], FileSystem] + fs_by_scheme: Callable[..., FileSystem] def __init__(self, properties: Properties = EMPTY_DICT): - self.fs_by_scheme: Callable[[str, str | None], FileSystem] = lru_cache(self._initialize_fs) + self.fs_by_scheme: Callable[..., FileSystem] = lru_cache(self._initialize_fs) super().__init__(properties=properties) + def _fs_for_location(self, location: str, scheme: str, netloc: str | None = None) -> FileSystem: + """Return the filesystem for a location, folding vended credentials into the cache key. + + When a credentials provider is attached, its resolved (and possibly refreshed) credential + properties become part of the 'fs_by_scheme' cache key, so rotated credentials build a + new filesystem while unchanged credentials reuse the cached one. + """ + creds_key: frozenset[tuple[str, str]] = frozenset() + if provider := self._credentials_provider: + creds_key = frozenset(provider.properties_for(location).items()) + return self.fs_by_scheme(scheme, netloc, creds_key) + @staticmethod def parse_location(location: str, properties: Properties = EMPTY_DICT) -> tuple[str, str, str]: r"""Return (scheme, netloc, path) for the given location. @@ -424,22 +436,26 @@ def parse_location(location: str, properties: Properties = EMPTY_DICT) -> tuple[ else: return uri.scheme, uri.netloc, f"{uri.netloc}{uri.path}" - def _initialize_fs(self, scheme: str, netloc: str | None = None) -> FileSystem: + def _initialize_fs( + self, scheme: str, netloc: str | None = None, creds_key: frozenset[tuple[str, str]] = frozenset() + ) -> FileSystem: """Initialize FileSystem for different scheme.""" + props = {**self.properties, **dict(creds_key)} + if scheme in {"oss"}: - return self._initialize_oss_fs() + return self._initialize_oss_fs(props) elif scheme in {"s3", "s3a", "s3n"}: - return self._initialize_s3_fs(netloc) + return self._initialize_s3_fs(netloc, props) elif scheme in {"hdfs", "viewfs"}: return self._initialize_hdfs_fs(scheme, netloc) elif scheme in {"gs", "gcs"}: - return self._initialize_gcs_fs() + return self._initialize_gcs_fs(props) elif scheme in {"abfs", "abfss", "wasb", "wasbs"}: - return self._initialize_azure_fs() + return self._initialize_azure_fs(props) elif scheme in {"file"}: return self._initialize_local_fs() @@ -447,45 +463,45 @@ def _initialize_fs(self, scheme: str, netloc: str | None = None) -> FileSystem: else: raise ValueError(f"Unrecognized filesystem type in URI: {scheme}") - def _initialize_oss_fs(self) -> FileSystem: + def _initialize_oss_fs(self, properties: Properties) -> FileSystem: from pyarrow.fs import S3FileSystem client_kwargs: dict[str, Any] = { - "endpoint_override": self.properties.get(S3_ENDPOINT), - "access_key": get_first_property_value(self.properties, S3_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID), - "secret_key": get_first_property_value(self.properties, S3_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY), - "session_token": get_first_property_value(self.properties, S3_SESSION_TOKEN, AWS_SESSION_TOKEN), - "region": get_first_property_value(self.properties, S3_REGION, AWS_REGION), - "force_virtual_addressing": property_as_bool(self.properties, S3_FORCE_VIRTUAL_ADDRESSING, True), + "endpoint_override": properties.get(S3_ENDPOINT), + "access_key": get_first_property_value(properties, S3_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID), + "secret_key": get_first_property_value(properties, S3_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY), + "session_token": get_first_property_value(properties, S3_SESSION_TOKEN, AWS_SESSION_TOKEN), + "region": get_first_property_value(properties, S3_REGION, AWS_REGION), + "force_virtual_addressing": property_as_bool(properties, S3_FORCE_VIRTUAL_ADDRESSING, True), } - if proxy_uri := self.properties.get(S3_PROXY_URI): + if proxy_uri := properties.get(S3_PROXY_URI): client_kwargs["proxy_options"] = proxy_uri - if connect_timeout := self.properties.get(S3_CONNECT_TIMEOUT): + if connect_timeout := properties.get(S3_CONNECT_TIMEOUT): client_kwargs["connect_timeout"] = float(connect_timeout) - if request_timeout := self.properties.get(S3_REQUEST_TIMEOUT): + if request_timeout := properties.get(S3_REQUEST_TIMEOUT): client_kwargs["request_timeout"] = float(request_timeout) - if role_arn := get_first_property_value(self.properties, S3_ROLE_ARN, AWS_ROLE_ARN): + if role_arn := get_first_property_value(properties, S3_ROLE_ARN, AWS_ROLE_ARN): client_kwargs["role_arn"] = role_arn - if session_name := get_first_property_value(self.properties, S3_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME): + if session_name := get_first_property_value(properties, S3_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME): client_kwargs["session_name"] = session_name - if s3_anonymous := self.properties.get(S3_ANONYMOUS): + if s3_anonymous := properties.get(S3_ANONYMOUS): client_kwargs["anonymous"] = strtobool(s3_anonymous) return S3FileSystem(**client_kwargs) - def _initialize_s3_fs(self, netloc: str | None) -> FileSystem: + def _initialize_s3_fs(self, netloc: str | None, properties: Properties) -> FileSystem: from pyarrow.fs import S3FileSystem - provided_region = get_first_property_value(self.properties, S3_REGION, AWS_REGION) + provided_region = get_first_property_value(properties, S3_REGION, AWS_REGION) # Do this when we don't provide the region at all, or when we explicitly enable it - if provided_region is None or property_as_bool(self.properties, S3_RESOLVE_REGION, False) is True: + if provided_region is None or property_as_bool(properties, S3_RESOLVE_REGION, False) is True: # Resolve region from netloc(bucket), fallback to user-provided region # Only supported by buckets hosted by S3 bucket_region = _cached_resolve_s3_region(bucket=netloc) or provided_region @@ -498,42 +514,42 @@ def _initialize_s3_fs(self, netloc: str | None) -> FileSystem: bucket_region = provided_region client_kwargs: dict[str, Any] = { - "endpoint_override": self.properties.get(S3_ENDPOINT), - "access_key": get_first_property_value(self.properties, S3_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID), - "secret_key": get_first_property_value(self.properties, S3_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY), - "session_token": get_first_property_value(self.properties, S3_SESSION_TOKEN, AWS_SESSION_TOKEN), + "endpoint_override": properties.get(S3_ENDPOINT), + "access_key": get_first_property_value(properties, S3_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID), + "secret_key": get_first_property_value(properties, S3_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY), + "session_token": get_first_property_value(properties, S3_SESSION_TOKEN, AWS_SESSION_TOKEN), "region": bucket_region, } - if proxy_uri := self.properties.get(S3_PROXY_URI): + if proxy_uri := properties.get(S3_PROXY_URI): client_kwargs["proxy_options"] = proxy_uri - if connect_timeout := self.properties.get(S3_CONNECT_TIMEOUT): + if connect_timeout := properties.get(S3_CONNECT_TIMEOUT): client_kwargs["connect_timeout"] = float(connect_timeout) - if request_timeout := self.properties.get(S3_REQUEST_TIMEOUT): + if request_timeout := properties.get(S3_REQUEST_TIMEOUT): client_kwargs["request_timeout"] = float(request_timeout) - if role_arn := get_first_property_value(self.properties, S3_ROLE_ARN, AWS_ROLE_ARN): + if role_arn := get_first_property_value(properties, S3_ROLE_ARN, AWS_ROLE_ARN): client_kwargs["role_arn"] = role_arn - if session_name := get_first_property_value(self.properties, S3_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME): + if session_name := get_first_property_value(properties, S3_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME): client_kwargs["session_name"] = session_name - if self.properties.get(S3_FORCE_VIRTUAL_ADDRESSING) is not None: - client_kwargs["force_virtual_addressing"] = property_as_bool(self.properties, S3_FORCE_VIRTUAL_ADDRESSING, False) + if properties.get(S3_FORCE_VIRTUAL_ADDRESSING) is not None: + client_kwargs["force_virtual_addressing"] = property_as_bool(properties, S3_FORCE_VIRTUAL_ADDRESSING, False) - if (retry_strategy_impl := self.properties.get(S3_RETRY_STRATEGY_IMPL)) and ( + if (retry_strategy_impl := properties.get(S3_RETRY_STRATEGY_IMPL)) and ( retry_instance := _import_retry_strategy(retry_strategy_impl) ): client_kwargs["retry_strategy"] = retry_instance - if s3_anonymous := self.properties.get(S3_ANONYMOUS): + if s3_anonymous := properties.get(S3_ANONYMOUS): client_kwargs["anonymous"] = strtobool(s3_anonymous) return S3FileSystem(**client_kwargs) - def _initialize_azure_fs(self) -> FileSystem: + def _initialize_azure_fs(self, properties: Properties) -> FileSystem: # https://arrow.apache.org/docs/python/generated/pyarrow.fs.AzureFileSystem.html from packaging import version @@ -548,32 +564,32 @@ def _initialize_azure_fs(self) -> FileSystem: client_kwargs: dict[str, str] = {} - if account_name := self.properties.get(ADLS_ACCOUNT_NAME): + if account_name := properties.get(ADLS_ACCOUNT_NAME): client_kwargs["account_name"] = account_name - if account_key := self.properties.get(ADLS_ACCOUNT_KEY): + if account_key := properties.get(ADLS_ACCOUNT_KEY): client_kwargs["account_key"] = account_key - if blob_storage_authority := self.properties.get(ADLS_BLOB_STORAGE_AUTHORITY): + if blob_storage_authority := properties.get(ADLS_BLOB_STORAGE_AUTHORITY): client_kwargs["blob_storage_authority"] = blob_storage_authority - if dfs_storage_authority := self.properties.get(ADLS_DFS_STORAGE_AUTHORITY): + if dfs_storage_authority := properties.get(ADLS_DFS_STORAGE_AUTHORITY): client_kwargs["dfs_storage_authority"] = dfs_storage_authority - if blob_storage_scheme := self.properties.get(ADLS_BLOB_STORAGE_SCHEME): + if blob_storage_scheme := properties.get(ADLS_BLOB_STORAGE_SCHEME): client_kwargs["blob_storage_scheme"] = blob_storage_scheme - if dfs_storage_scheme := self.properties.get(ADLS_DFS_STORAGE_SCHEME): + if dfs_storage_scheme := properties.get(ADLS_DFS_STORAGE_SCHEME): client_kwargs["dfs_storage_scheme"] = dfs_storage_scheme - if sas_token := self.properties.get(ADLS_SAS_TOKEN): + if sas_token := properties.get(ADLS_SAS_TOKEN): client_kwargs["sas_token"] = sas_token - if client_id := self.properties.get(ADLS_CLIENT_ID): + if client_id := properties.get(ADLS_CLIENT_ID): client_kwargs["client_id"] = client_id - if client_secret := self.properties.get(ADLS_CLIENT_SECRET): + if client_secret := properties.get(ADLS_CLIENT_SECRET): client_kwargs["client_secret"] = client_secret - if tenant_id := self.properties.get(ADLS_TENANT_ID): + if tenant_id := properties.get(ADLS_TENANT_ID): client_kwargs["tenant_id"] = tenant_id # Validate that all three are provided together for ClientSecretCredential @@ -607,17 +623,17 @@ def _initialize_hdfs_fs(self, scheme: str, netloc: str | None) -> FileSystem: return HadoopFileSystem(**hdfs_kwargs) - def _initialize_gcs_fs(self) -> FileSystem: + def _initialize_gcs_fs(self, properties: Properties) -> FileSystem: from pyarrow.fs import GcsFileSystem gcs_kwargs: dict[str, Any] = {} - if access_token := self.properties.get(GCS_TOKEN): + if access_token := properties.get(GCS_TOKEN): gcs_kwargs["access_token"] = access_token - if expiration := self.properties.get(GCS_TOKEN_EXPIRES_AT_MS): + if expiration := properties.get(GCS_TOKEN_EXPIRES_AT_MS): gcs_kwargs["credential_token_expiration"] = millis_to_datetime(int(expiration)) - if bucket_location := self.properties.get(GCS_DEFAULT_LOCATION): + if bucket_location := properties.get(GCS_DEFAULT_LOCATION): gcs_kwargs["default_bucket_location"] = bucket_location - if endpoint := self.properties.get(GCS_SERVICE_HOST): + if endpoint := properties.get(GCS_SERVICE_HOST): url_parts = urlparse(endpoint) gcs_kwargs["scheme"] = url_parts.scheme gcs_kwargs["endpoint_override"] = url_parts.netloc @@ -639,7 +655,7 @@ def new_input(self, location: str) -> PyArrowFile: """ scheme, netloc, path = self.parse_location(location, self.properties) return PyArrowFile( - fs=self.fs_by_scheme(scheme, netloc), + fs=self._fs_for_location(location, scheme, netloc), location=location, path=path, buffer_size=int(self.properties.get(BUFFER_SIZE, ONE_MEGABYTE)), @@ -657,7 +673,7 @@ def new_output(self, location: str) -> PyArrowFile: """ scheme, netloc, path = self.parse_location(location, self.properties) return PyArrowFile( - fs=self.fs_by_scheme(scheme, netloc), + fs=self._fs_for_location(location, scheme, netloc), location=location, path=path, buffer_size=int(self.properties.get(BUFFER_SIZE, ONE_MEGABYTE)), @@ -679,7 +695,7 @@ def delete(self, location: str | InputFile | OutputFile) -> None: """ str_location = location.location if isinstance(location, (InputFile, OutputFile)) else location scheme, netloc, path = self.parse_location(str_location, self.properties) - fs = self.fs_by_scheme(scheme, netloc) + fs = self._fs_for_location(str_location, scheme, netloc) try: fs.delete_file(path) diff --git a/tests/catalog/test_credential_provider.py b/tests/catalog/test_credential_provider.py new file mode 100644 index 0000000000..02a377090b --- /dev/null +++ b/tests/catalog/test_credential_provider.py @@ -0,0 +1,248 @@ +# 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. + +import threading +import time +from unittest.mock import MagicMock + +from pyiceberg.catalog.rest import LoadCredentialsResponse +from pyiceberg.catalog.rest.credential_provider import CredentialsProvider, is_s3_credential_expired, resolve_storage_credentials +from pyiceberg.catalog.rest.scan_planning import StorageCredential +from pyiceberg.typedef import Properties + +BASE_CREDENTIAL = StorageCredential( + prefix="s3://warehouse/", + config={ + "s3.access-key-id": "initial-key", + "s3.secret-access-key": "initial-secret", + "s3.session-token": "initial-token", + }, +) + +LOCATION = "s3://warehouse/database/table/metadata/00001.json" + + +def _expiry_ms_in(seconds: float) -> str: + return str(int((time.time() + seconds) * 1000)) + + +def test_no_expiry_is_treated_as_static() -> None: + assert is_s3_credential_expired({"s3.session-token": "token"}) is False + + +def test_far_expiry_is_not_expired() -> None: + assert is_s3_credential_expired({"s3.session-token-expires-at-ms": _expiry_ms_in(3600)}) is False + + +def test_near_expiry_is_expired() -> None: + assert is_s3_credential_expired({"s3.session-token-expires-at-ms": _expiry_ms_in(60)}) is True + + +def test_past_expiry_is_expired() -> None: + assert is_s3_credential_expired({"s3.session-token-expires-at-ms": _expiry_ms_in(-60)}) is True + + +def test_threshold_boundary() -> None: + # Just inside the threshold -> expired; well outside -> not expired + assert is_s3_credential_expired({"s3.session-token-expires-at-ms": _expiry_ms_in(100)}, threshold_seconds=300) is True + assert is_s3_credential_expired({"s3.session-token-expires-at-ms": _expiry_ms_in(600)}, threshold_seconds=300) is False + + +def _near_expiry_credential() -> StorageCredential: + near_expiry_ms = str(int((time.time() + 60) * 1000)) + return StorageCredential( + prefix="s3://warehouse/", + config={**BASE_CREDENTIAL.config, "s3.session-token-expires-at-ms": near_expiry_ms}, + ) + + +def _far_expiry_credential() -> StorageCredential: + far_expiry_ms = str(int((time.time() + 3600) * 1000)) + return StorageCredential( + prefix="s3://warehouse/", + config={**BASE_CREDENTIAL.config, "s3.session-token-expires-at-ms": far_expiry_ms}, + ) + + +def test_resolve_storage_credentials_longest_prefix_wins() -> None: + credentials = [ + StorageCredential(prefix="s3://warehouse/", config={"s3.access-key-id": "short-prefix-key"}), + StorageCredential(prefix="s3://warehouse/database/table", config={"s3.access-key-id": "long-prefix-key"}), + ] + assert resolve_storage_credentials(credentials, LOCATION) == {"s3.access-key-id": "long-prefix-key"} + + +def test_resolve_storage_credentials_empty() -> None: + assert resolve_storage_credentials([], "s3://warehouse/foo") == {} + assert resolve_storage_credentials([], None) == {} + + +def test_properties_for_multiple_prefixes_longest_match_wins() -> None: + credentials = [ + StorageCredential(prefix="s3://warehouse/", config={"s3.access-key-id": "short-prefix-key"}), + StorageCredential(prefix="s3://warehouse/database/table", config={"s3.access-key-id": "long-prefix-key"}), + ] + provider = CredentialsProvider(credentials, refresh_fn=MagicMock()) + assert provider.properties_for(LOCATION) == {"s3.access-key-id": "long-prefix-key"} + + +def test_properties_for_no_expiry_returns_static_creds_without_refresh() -> None: + refresh_fn = MagicMock() + provider = CredentialsProvider([BASE_CREDENTIAL], refresh_fn=refresh_fn) + + config = provider.properties_for(LOCATION) + + refresh_fn.assert_not_called() + assert config == BASE_CREDENTIAL.config + + +def test_properties_for_far_expiry_does_not_refresh() -> None: + refresh_fn = MagicMock() + provider = CredentialsProvider([_far_expiry_credential()], refresh_fn=refresh_fn) + + provider.properties_for(LOCATION) + + refresh_fn.assert_not_called() + + +def test_properties_for_near_expiry_triggers_refresh_once() -> None: + refreshed_credential = StorageCredential( + prefix="s3://warehouse/", + config={ + "s3.access-key-id": "refreshed-key", + "s3.secret-access-key": "refreshed-secret", + "s3.session-token": "refreshed-token", + }, + ) + refresh_fn = MagicMock(return_value=LoadCredentialsResponse(storage_credentials=[refreshed_credential])) + provider = CredentialsProvider([_near_expiry_credential()], refresh_fn=refresh_fn) + + config = provider.properties_for(LOCATION) + + refresh_fn.assert_called_once() + assert config == refreshed_credential.config + + +def test_properties_for_empty_storage_credentials_returns_empty() -> None: + provider = CredentialsProvider([], refresh_fn=MagicMock()) + assert provider.properties_for(LOCATION) == {} + + +def test_properties_for_no_match_returns_empty() -> None: + provider = CredentialsProvider( + [StorageCredential(prefix="s3://other-bucket/", config={"s3.access-key-id": "no-match"})], refresh_fn=MagicMock() + ) + assert provider.properties_for(LOCATION) == {} + + +def test_properties_for_far_expiry_skips_lock_entirely() -> None: + """The outer un-locked check should short-circuit before ever touching the lock.""" + refresh_fn = MagicMock() + provider = CredentialsProvider([_far_expiry_credential()], refresh_fn=refresh_fn) + mock_lock = MagicMock(spec=threading.Lock()) + provider._lock = mock_lock + + provider.properties_for(LOCATION) + + mock_lock.__enter__.assert_not_called() + refresh_fn.assert_not_called() + + +def test_properties_for_concurrent_near_expiry_refreshes_exactly_once() -> None: + """Double-checked locking: concurrent callers must trigger only a single refresh.""" + refreshed_credential = StorageCredential( + prefix="s3://warehouse/", + config={ + "s3.access-key-id": "refreshed-key", + "s3.secret-access-key": "refreshed-secret", + "s3.session-token": "refreshed-token", + }, + ) + call_count = 0 + count_lock = threading.Lock() + + def slow_refresh_fn() -> LoadCredentialsResponse: + nonlocal call_count + with count_lock: + call_count += 1 + # Widen the race window so other threads pile up waiting on the provider's lock + # while this refresh is still in flight. + time.sleep(0.2) + return LoadCredentialsResponse(storage_credentials=[refreshed_credential]) + + provider = CredentialsProvider([_near_expiry_credential()], refresh_fn=slow_refresh_fn) + + thread_count = 10 + barrier = threading.Barrier(thread_count) + results: list[Properties] = [] + results_lock = threading.Lock() + + def worker() -> None: + barrier.wait() + config = provider.properties_for(LOCATION) + with results_lock: + results.append(config) + + threads = [threading.Thread(target=worker) for _ in range(thread_count)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert call_count == 1 + assert len(results) == thread_count + assert all(config == refreshed_credential.config for config in results) + + +def test_properties_for_second_thread_reuses_refresh_done_by_first() -> None: + """A thread that blocks on the lock must see the already-refreshed config, not trigger its own refresh.""" + refreshed_credential = StorageCredential(prefix="s3://warehouse/", config={"s3.access-key-id": "refreshed-key"}) + call_count = 0 + first_thread_holding_lock = threading.Event() + release_first_thread = threading.Event() + + def refresh_fn() -> LoadCredentialsResponse: + nonlocal call_count + call_count += 1 + first_thread_holding_lock.set() + release_first_thread.wait() + return LoadCredentialsResponse(storage_credentials=[refreshed_credential]) + + provider = CredentialsProvider([_near_expiry_credential()], refresh_fn=refresh_fn) + + first_result: list[Properties] = [] + second_result: list[Properties] = [] + + def first_call() -> None: + first_result.append(provider.properties_for(LOCATION)) + + first_thread = threading.Thread(target=first_call) + first_thread.start() + first_thread_holding_lock.wait() + + # Second thread starts once the first is inside the lock refreshing; it should block on + # the lock, then see the refreshed config once it acquires the lock and re-checks. + second_thread = threading.Thread(target=lambda: second_result.append(provider.properties_for(LOCATION))) + second_thread.start() + + release_first_thread.set() + first_thread.join() + second_thread.join() + + assert call_count == 1 + assert first_result[0] == refreshed_credential.config + assert second_result[0] == refreshed_credential.config diff --git a/tests/catalog/test_rest.py b/tests/catalog/test_rest.py index 43790f778a..e3bf0e9269 100644 --- a/tests/catalog/test_rest.py +++ b/tests/catalog/test_rest.py @@ -3085,29 +3085,33 @@ def test_endpoint_parsing_from_string_with_invalid_http_method() -> None: def test_resolve_storage_credentials_longest_prefix_wins() -> None: + from pyiceberg.catalog.rest.credential_provider import resolve_storage_credentials from pyiceberg.catalog.rest.scan_planning import StorageCredential credentials = [ StorageCredential(prefix="s3://warehouse/", config={"s3.access-key-id": "short-prefix-key"}), StorageCredential(prefix="s3://warehouse/database/table", config={"s3.access-key-id": "long-prefix-key"}), ] - result = RestCatalog._resolve_storage_credentials(credentials, "s3://warehouse/database/table/metadata/00001.json") + result = resolve_storage_credentials(credentials, "s3://warehouse/database/table/metadata/00001.json") assert result == {"s3.access-key-id": "long-prefix-key"} def test_resolve_storage_credentials_no_match() -> None: + from pyiceberg.catalog.rest.credential_provider import resolve_storage_credentials from pyiceberg.catalog.rest.scan_planning import StorageCredential credentials = [ StorageCredential(prefix="s3://other-bucket/", config={"s3.access-key-id": "no-match"}), ] - result = RestCatalog._resolve_storage_credentials(credentials, "s3://warehouse/database/table/metadata/00001.json") + result = resolve_storage_credentials(credentials, "s3://warehouse/database/table/metadata/00001.json") assert result == {} def test_resolve_storage_credentials_empty() -> None: - assert RestCatalog._resolve_storage_credentials([], "s3://warehouse/foo") == {} - assert RestCatalog._resolve_storage_credentials([], None) == {} + from pyiceberg.catalog.rest.credential_provider import resolve_storage_credentials + + assert resolve_storage_credentials([], "s3://warehouse/foo") == {} + assert resolve_storage_credentials([], None) == {} def test_load_table_with_storage_credentials(rest_mock: Mocker, example_table_metadata_with_snapshot_v1: dict[str, Any]) -> None: @@ -3144,6 +3148,80 @@ def test_load_table_with_storage_credentials(rest_mock: Mocker, example_table_me assert table.io.properties["s3.session-token"] == "vended-token" +def _mock_table_with_storage_credentials(rest_mock: Mocker, metadata: dict[str, Any]) -> None: + rest_mock.get( + f"{TEST_URI}v1/namespaces/fokko/tables/table", + json={ + "metadata-location": "s3://warehouse/database/table/metadata/00001.metadata.json", + "metadata": metadata, + "config": {}, + "storage-credentials": [ + { + "prefix": "s3://warehouse/database/table", + "config": {"s3.access-key-id": "vended-key"}, + } + ], + }, + status_code=200, + request_headers=TEST_HEADERS, + ) + + +def test_load_table_attaches_credentials_provider_when_enabled( + rest_mock: Mocker, example_table_metadata_with_snapshot_v1: dict[str, Any] +) -> None: + from pyiceberg.catalog.rest.credential_provider import REFRESH_CREDENTIALS_ENABLED, CredentialsProvider + from pyiceberg.io import FileIO + + _mock_table_with_storage_credentials(rest_mock, example_table_metadata_with_snapshot_v1) + catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN, **{REFRESH_CREDENTIALS_ENABLED: "true"}) + with mock.patch.object(FileIO, "set_credentials_provider") as mock_set_credentials_provider: + catalog.load_table(("fokko", "table")) + + mock_set_credentials_provider.assert_called_once() + (provider,), _ = mock_set_credentials_provider.call_args + assert isinstance(provider, CredentialsProvider) + + +def test_load_table_no_provider_when_flag_disabled( + rest_mock: Mocker, example_table_metadata_with_snapshot_v1: dict[str, Any] +) -> None: + from pyiceberg.io import FileIO + + _mock_table_with_storage_credentials(rest_mock, example_table_metadata_with_snapshot_v1) + catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN) + with mock.patch.object(FileIO, "set_credentials_provider") as mock_set_credentials_provider: + catalog.load_table(("fokko", "table")) + + mock_set_credentials_provider.assert_not_called() + + +def test_attached_provider_refresh_fn_returns_full_response( + rest_mock: Mocker, example_table_metadata_with_snapshot_v1: dict[str, Any] +) -> None: + """The refresh callback must return the raw LoadCredentialsResponse (full credential list), + not the already-resolved Properties, so the provider can re-run longest-prefix matching.""" + from pyiceberg.catalog.rest import LoadCredentialsResponse + from pyiceberg.catalog.rest.credential_provider import REFRESH_CREDENTIALS_ENABLED + from pyiceberg.io import FileIO + + _mock_table_with_storage_credentials(rest_mock, example_table_metadata_with_snapshot_v1) + rest_mock.get( + f"{TEST_URI}v1/namespaces/fokko/tables/table/credentials", + json={"storage-credentials": [{"prefix": "s3://warehouse/database/table", "config": {"s3.access-key-id": "refreshed"}}]}, + status_code=200, + request_headers=TEST_HEADERS, + ) + catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN, **{REFRESH_CREDENTIALS_ENABLED: "true"}) + captured: dict[str, LoadCredentialsResponse] = {} + with mock.patch.object(FileIO, "set_credentials_provider", side_effect=lambda p: captured.update(provider=p)): + catalog.load_table(("fokko", "table")) + + response = captured["provider"]._refresh_fn() + assert isinstance(response, LoadCredentialsResponse) + assert response.storage_credentials[0].config == {"s3.access-key-id": "refreshed"} + + def test_load_credentials_with_longest_prefix(rest_mock: Mocker) -> None: rest_mock.get( f"{TEST_URI}v1/namespaces/fokko/tables/table/credentials", diff --git a/tests/io/test_credential_refresh.py b/tests/io/test_credential_refresh.py new file mode 100644 index 0000000000..e3cbee390f --- /dev/null +++ b/tests/io/test_credential_refresh.py @@ -0,0 +1,97 @@ +# 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. +"""End-to-end tests that the FileIO backends rebuild their filesystem when vended creds change.""" + +from typing import Any +from unittest import mock + +from pyiceberg.io.fsspec import FsspecFileIO +from pyiceberg.io.pyarrow import PyArrowFileIO +from pyiceberg.typedef import Properties + +LOCATION = "s3://warehouse/database/table/data/00000.parquet" + + +class _FakeProvider: + """A credentials provider whose returned creds can be swapped between calls.""" + + def __init__(self, creds: Properties) -> None: + self.creds = creds + + def properties_for(self, location: str) -> Properties: + return self.creds + + +def test_pyarrow_rebuilds_fs_when_credentials_change() -> None: + provider = _FakeProvider({"s3.session-token": "token-1"}) + io = PyArrowFileIO() + io.set_credentials_provider(provider) + + seen_tokens: list[Any] = [] + + def fake_s3_fs(netloc: Any, properties: Properties) -> object: + seen_tokens.append(properties.get("s3.session-token")) + return object() + + with mock.patch.object(PyArrowFileIO, "_initialize_s3_fs", side_effect=fake_s3_fs): + fs1 = io.new_input(LOCATION)._filesystem + fs2 = io.new_input(LOCATION)._filesystem # unchanged creds -> cached + provider.creds = {"s3.session-token": "token-2"} + fs3 = io.new_input(LOCATION)._filesystem # rotated creds -> rebuilt + + assert fs1 is fs2 + assert fs3 is not fs1 + assert seen_tokens == ["token-1", "token-2"] + + +def test_pyarrow_without_provider_uses_single_fs() -> None: + io = PyArrowFileIO() + build_count = 0 + + def fake_s3_fs(netloc: Any, properties: Properties) -> object: + nonlocal build_count + build_count += 1 + return object() + + with mock.patch.object(PyArrowFileIO, "_initialize_s3_fs", side_effect=fake_s3_fs): + io.new_input(LOCATION) + io.new_input(LOCATION) + + assert build_count == 1 + + +def test_fsspec_rebuilds_fs_when_credentials_change() -> None: + provider = _FakeProvider({"s3.session-token": "token-1"}) + io = FsspecFileIO(properties={}) + io.set_credentials_provider(provider) + + seen_tokens: list[Any] = [] + + def fake_s3(properties: Properties) -> object: + seen_tokens.append(properties.get("s3.session-token")) + return object() + + io._scheme_to_fs = {"s3": fake_s3} + + fs1 = io.new_input(LOCATION)._fs + fs2 = io.new_input(LOCATION)._fs # unchanged creds -> cached + provider.creds = {"s3.session-token": "token-2"} + fs3 = io.new_input(LOCATION)._fs # rotated creds -> rebuilt + + assert fs1 is fs2 + assert fs3 is not fs1 + assert seen_tokens == ["token-1", "token-2"]