From d92ef4db571c7130b61f892356a427e487e0aed2 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Mon, 17 Aug 2026 06:07:39 -0400 Subject: [PATCH] feat(io): refresh vended storage credentials before they expire --- src/iceberg/arrow/s3/arrow_s3_file_io.cc | 332 ++++++++++++++-- src/iceberg/arrow/s3/s3_properties.h | 3 + src/iceberg/catalog/rest/json_serde.cc | 42 +- .../catalog/rest/json_serde_internal.h | 4 + src/iceberg/catalog/rest/rest_catalog.cc | 61 ++- src/iceberg/catalog/rest/rest_catalog.h | 11 +- src/iceberg/catalog/rest/rest_file_io.cc | 8 +- src/iceberg/catalog/rest/rest_file_io.h | 6 +- src/iceberg/catalog/rest/types.h | 17 + src/iceberg/file_io.h | 12 +- src/iceberg/resolving_file_io.cc | 18 +- src/iceberg/resolving_file_io.h | 20 +- src/iceberg/storage_credential.h | 9 + src/iceberg/test/arrow_s3_file_io_test.cc | 366 +++++++++++++++++- src/iceberg/test/resolving_file_io_test.cc | 54 ++- src/iceberg/test/rest_file_io_test.cc | 36 +- src/iceberg/test/rest_json_serde_test.cc | 24 ++ 17 files changed, 959 insertions(+), 64 deletions(-) diff --git a/src/iceberg/arrow/s3/arrow_s3_file_io.cc b/src/iceberg/arrow/s3/arrow_s3_file_io.cc index e3118453f8..6a97681305 100644 --- a/src/iceberg/arrow/s3/arrow_s3_file_io.cc +++ b/src/iceberg/arrow/s3/arrow_s3_file_io.cc @@ -17,9 +17,14 @@ * under the License. */ +#include +#include +#include #include #include +#include #include +#include #include #include #include @@ -194,11 +199,76 @@ std::string CanonicalizeS3Scheme(std::string_view location) { return std::string(location); } +// Lead time before expiry, matching Java's VendedCredentialsProvider. +constexpr auto kRefreshLeadTime = std::chrono::minutes(5); + +// After a failed refresh, how long to keep the current credentials before +// asking again, so an unreachable catalog is not queried per file operation. +constexpr auto kRefreshRetryBackoff = std::chrono::seconds(30); + +// Floor on a backoff shortened to land on the expiry. +constexpr auto kMinRefreshRetryBackoff = std::chrono::seconds(1); + +// How long an operation with expired credentials waits for a refresh already +// under way. Bounded: the catalog request behind it has no deadline of its own. +constexpr auto kExpiredCredentialWait = std::chrono::seconds(10); + +// When the earliest of `credentials` stops being valid, or nullopt if none of +// them does. No session token means static keys, which never expire; a token +// with no usable expiry is reported as already expired so it gets replaced +// rather than used until it fails, as Java does. +std::optional EarliestExpiry( + const std::vector& credentials) { + std::optional earliest; + const auto note = [&earliest](std::chrono::system_clock::time_point expires_at) { + if (!earliest.has_value() || expires_at < *earliest) { + earliest = expires_at; + } + }; + for (const auto& credential : credentials) { + if (!IsS3FileIOCredentialPrefix(credential.prefix) || + FindProperty(credential.config, S3Properties::kSessionToken) == nullptr) { + continue; + } + const auto* value = + FindProperty(credential.config, S3Properties::kSessionTokenExpiresAtMs); + if (value == nullptr) { + ICEBERG_LOG_WARN("Credential \"{}\" has a session token but no \"{}\"", + credential.prefix, S3Properties::kSessionTokenExpiresAtMs); + note(std::chrono::system_clock::now()); + continue; + } + auto millis = StringUtils::ParseNumber(*value); + if (!millis.has_value()) { + ICEBERG_LOG_WARN( + "Credential \"{}\" has a session token but an unparseable \"{}\" value \"{}\"", + credential.prefix, S3Properties::kSessionTokenExpiresAtMs, *value); + note(std::chrono::system_clock::now()); + continue; + } + // Beyond what the clock can hold, converting would overflow it. + constexpr auto kMaxMillis = std::chrono::duration_cast( + std::chrono::system_clock::duration::max()) + .count(); + constexpr auto kMinMillis = std::chrono::duration_cast( + std::chrono::system_clock::duration::min()) + .count(); + if (*millis > kMaxMillis || *millis < kMinMillis) { + ICEBERG_LOG_WARN("Credential \"{}\" has an out-of-range \"{}\" value \"{}\"", + credential.prefix, S3Properties::kSessionTokenExpiresAtMs, *value); + note(std::chrono::system_clock::now()); + continue; + } + note(std::chrono::system_clock::time_point(std::chrono::milliseconds(*millis))); + } + return earliest; +} + class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { public: ArrowS3FileIO(std::shared_ptr<::arrow::fs::FileSystem> arrow_fs, std::unordered_map default_properties) - : default_file_io_(std::move(arrow_fs)), + : default_file_io_(std::make_shared(std::move(arrow_fs))), default_properties_(std::move(default_properties)) {} Result> NewInputFile(std::string file_location) override; @@ -215,28 +285,96 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { Status SetStorageCredentials( const std::vector& storage_credentials) override; - const std::vector& credentials() const override { + std::vector credentials() const override { + std::shared_lock lock(mutex_); return storage_credentials_; } + void SetCredentialRefresher(StorageCredentialRefresher refresher) override { + std::unique_lock lock(mutex_); + refresher_ = std::move(refresher); + // A refresh in flight was started for the refresher just replaced. + ++credential_generation_; + } + SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; } private: - ArrowFileSystemFileIO& FileIOForPath(std::string_view location); - - ArrowFileSystemFileIO default_file_io_; + /// \brief Delegate serving `location`, kept alive by the caller so a + /// concurrent refresh cannot free it mid-operation. + std::shared_ptr FileIOForPath(std::string_view location); + + using DelegatesByPrefix = + std::vector>>; + + /// \brief Build a delegate for each credential this FileIO can serve. + /// + /// Lock-free on purpose: building an S3 client can reach out to discover a + /// bucket region, which would stall every concurrent operation. Reads no + /// mutable member state. + Result BuildDelegates( + const std::vector& storage_credentials) const; + + /// \brief Swap in credentials and the delegates built from them. + /// + /// Callers must hold `mutex_` exclusively. + void InstallCredentials(const std::vector& storage_credentials, + DelegatesByPrefix delegates); + + /// \brief Whether the installed credentials are close enough to expiring to + /// be replaced, and no backoff is in effect. + /// + /// Callers must hold `mutex_`, at least shared. + bool RefreshDue() const; + + /// \brief Whether the installed credentials have already stopped being valid. + /// + /// Callers must hold `mutex_`, at least shared. + bool Expired() const; + + /// \brief When the next refresh attempt becomes allowed after a failure. + /// + /// Never past the point the credentials stop being valid. + /// + /// Callers must hold `mutex_`, at least shared. + std::chrono::steady_clock::time_point BackoffUntil() const; + + /// \brief Replace the credentials once they are close to expiring. + /// + /// Called before each handle is created; a handle keeps the delegate it was + /// built from, so I/O on an open one is not re-checked. A failure keeps the + /// current credentials rather than failing the read. + void MaybeRefreshCredentials(); + + std::shared_ptr default_file_io_; std::unordered_map default_properties_; + // Guards everything below; shared because reads happen per file operation. + mutable std::shared_mutex mutex_; std::vector storage_credentials_; - std::vector>> - file_io_by_prefix_; + DelegatesByPrefix file_io_by_prefix_; + StorageCredentialRefresher refresher_; + std::optional expires_at_; + std::chrono::steady_clock::time_point retry_refresh_at_; + // Bumped whenever the credentials or the refresher change, so a refresh that + // fetched before one of those happened can tell its result is already stale. + uint64_t credential_generation_ = 0; + // Held across a refresh so concurrent operations skip it. Timed, so waiting + // on it is bounded. + std::timed_mutex refresh_mutex_; }; Status ArrowS3FileIO::SetStorageCredentials( const std::vector& storage_credentials) { - std::vector>> - file_io_by_prefix; - file_io_by_prefix.reserve(storage_credentials.size()); - // TODO(gangwu): Refresh vended credentials via credentials.uri before tokens expire. + ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials)); + std::unique_lock lock(mutex_); + InstallCredentials(storage_credentials, std::move(delegates)); + return {}; +} + +Result ArrowS3FileIO::BuildDelegates( + const std::vector& storage_credentials) const { + DelegatesByPrefix delegates; + delegates.reserve(storage_credentials.size()); for (const auto& credential : storage_credentials) { ICEBERG_RETURN_UNEXPECTED(credential.Validate()); // A server may vend credentials for several storage systems at once; @@ -250,11 +388,10 @@ Status ArrowS3FileIO::SetStorageCredentials( properties[key] = value; } ICEBERG_ASSIGN_OR_RAISE(auto fs, BuildArrowS3FileSystem(properties)); - file_io_by_prefix.emplace_back( - CanonicalizeS3Scheme(credential.prefix), - std::make_unique(std::move(fs))); + delegates.emplace_back(CanonicalizeS3Scheme(credential.prefix), + std::make_shared(std::move(fs))); } - if (file_io_by_prefix.empty() && !storage_credentials.empty()) { + if (delegates.empty() && !storage_credentials.empty()) { // Silent skipping of every vended credential is hard to diagnose: S3 access // would proceed with the default credentials and fail only at IO time. ICEBERG_LOG_WARN( @@ -262,50 +399,191 @@ Status ArrowS3FileIO::SetStorageCredentials( "S3 access will use the default credentials", storage_credentials.size()); } - file_io_by_prefix_ = std::move(file_io_by_prefix); + return delegates; +} + +void ArrowS3FileIO::InstallCredentials( + const std::vector& storage_credentials, + DelegatesByPrefix delegates) { + file_io_by_prefix_ = std::move(delegates); storage_credentials_ = storage_credentials; - return {}; + expires_at_ = EarliestExpiry(storage_credentials); + retry_refresh_at_ = {}; + ++credential_generation_; } -ArrowFileSystemFileIO& ArrowS3FileIO::FileIOForPath(std::string_view location) { +bool ArrowS3FileIO::RefreshDue() const { + return expires_at_.has_value() && + std::chrono::system_clock::now() + kRefreshLeadTime >= *expires_at_ && + std::chrono::steady_clock::now() >= retry_refresh_at_; +} + +bool ArrowS3FileIO::Expired() const { + return expires_at_.has_value() && std::chrono::system_clock::now() >= *expires_at_; +} + +std::chrono::steady_clock::time_point ArrowS3FileIO::BackoffUntil() const { + auto delay = + std::chrono::duration_cast(kRefreshRetryBackoff); + if (expires_at_.has_value()) { + const auto remaining = std::chrono::duration_cast( + *expires_at_ - std::chrono::system_clock::now()); + // Worth retrying before they run out; once they have, faster retries only + // hammer a catalog that is already failing. + if (remaining > std::chrono::milliseconds::zero()) { + delay = std::clamp( + remaining, + std::chrono::duration_cast(kMinRefreshRetryBackoff), + delay); + } + } + return std::chrono::steady_clock::now() + delay; +} + +void ArrowS3FileIO::MaybeRefreshCredentials() { + { + // Cheap pre-check, so the common case costs one shared lock and no more. + std::shared_lock lock(mutex_); + if (!refresher_ || !RefreshDue()) { + return; + } + } + + std::unique_lock refresh_lock(refresh_mutex_, std::defer_lock); + if (!refresh_lock.try_lock()) { + // Another operation is already fetching; normally just use what we have. + { + std::shared_lock lock(mutex_); + if (!Expired()) { + return; + } + } + // Expired credentials leave nothing to proceed with, so wait instead. + if (!refresh_lock.try_lock_for(kExpiredCredentialWait)) { + return; + } + } + // Read together: pairing this refresher with a generation bumped by another + // one installed in between would make its result look current. + StorageCredentialRefresher refresher; + uint64_t generation = 0; + { + // Whoever held the lock may also have just finished, leaving nothing to do. + std::shared_lock lock(mutex_); + if (!refresher_ || !RefreshDue()) { + return; + } + refresher = refresher_; + generation = credential_generation_; + } + + // Outside `mutex_`: both are slow and must not block readers. + Status status; + DelegatesByPrefix delegates; + auto refreshed = refresher(); + if (refreshed.has_value()) { + auto built = BuildDelegates(*refreshed); + if (!built.has_value()) { + status = std::unexpected(built.error()); + } else if (built->empty()) { + // Installing this would drop working credentials for whatever ambient + // identity the AWS chain finds. Java refuses an empty list too. + status = NotFound("Refreshed credentials contain no S3-compatible prefix"); + } else { + delegates = std::move(built).value(); + } + } else { + status = std::unexpected(refreshed.error()); + } + + std::unique_lock lock(mutex_); + // Credentials installed meanwhile supersede this refresh: what it fetched is + // by now the older set. + const bool superseded = credential_generation_ != generation; + if (!status.has_value()) { + // Reported either way, so a failing catalog stays visible. + if (superseded) { + ICEBERG_LOG_WARN( + "Failed to refresh vended storage credentials ({}); they have since been " + "replaced", + status.error().message); + return; + } + retry_refresh_at_ = BackoffUntil(); + ICEBERG_LOG_WARN( + "Failed to refresh vended storage credentials ({}); keeping the current " + "ones and retrying in {}ms", + status.error().message, + std::chrono::duration_cast( + retry_refresh_at_ - std::chrono::steady_clock::now()) + .count()); + return; + } + if (superseded) { + return; + } + + InstallCredentials(*refreshed, std::move(delegates)); + if (RefreshDue()) { + // Tokens shorter-lived than the lead time come back due again at once. + retry_refresh_at_ = BackoffUntil(); + } +} + +std::shared_ptr ArrowS3FileIO::FileIOForPath( + std::string_view location) { + MaybeRefreshCredentials(); + + std::shared_lock lock(mutex_); if (file_io_by_prefix_.empty()) { return default_file_io_; } const std::string canonical = CanonicalizeS3Scheme(location); - ArrowFileSystemFileIO* best = &default_file_io_; + auto best = default_file_io_; size_t best_len = 0; for (const auto& [prefix, file_io] : file_io_by_prefix_) { if (prefix.size() > best_len && canonical.starts_with(prefix)) { - best = file_io.get(); + best = file_io; best_len = prefix.size(); } } - return *best; + return best; } Result> ArrowS3FileIO::NewInputFile( std::string file_location) { - return FileIOForPath(file_location).NewInputFile(std::move(file_location)); + return FileIOForPath(file_location)->NewInputFile(std::move(file_location)); } Result> ArrowS3FileIO::NewInputFile(std::string file_location, size_t length) { - return FileIOForPath(file_location).NewInputFile(std::move(file_location), length); + return FileIOForPath(file_location)->NewInputFile(std::move(file_location), length); } Result> ArrowS3FileIO::NewOutputFile( std::string file_location) { - return FileIOForPath(file_location).NewOutputFile(std::move(file_location)); + return FileIOForPath(file_location)->NewOutputFile(std::move(file_location)); } Status ArrowS3FileIO::DeleteFile(const std::string& file_location) { - return FileIOForPath(file_location).DeleteFile(file_location); + return FileIOForPath(file_location)->DeleteFile(file_location); } Status ArrowS3FileIO::DeleteFiles(const std::vector& file_locations) { - std::unordered_map> locations_by_io; + // Grouped by delegate, of which there are only ever a handful, so a linear + // scan beats hashing. + std::vector, std::vector>> + locations_by_io; for (const auto& file_location : file_locations) { - locations_by_io[&FileIOForPath(file_location)].push_back(file_location); + auto file_io = FileIOForPath(file_location); + auto it = std::ranges::find_if( + locations_by_io, [&](const auto& entry) { return entry.first == file_io; }); + if (it == locations_by_io.end()) { + locations_by_io.emplace_back(std::move(file_io), + std::vector{file_location}); + } else { + it->second.push_back(file_location); + } } for (auto& [file_io, locations] : locations_by_io) { ICEBERG_RETURN_UNEXPECTED(file_io->DeleteFiles(locations)); diff --git a/src/iceberg/arrow/s3/s3_properties.h b/src/iceberg/arrow/s3/s3_properties.h index 03d1492ade..11892e0986 100644 --- a/src/iceberg/arrow/s3/s3_properties.h +++ b/src/iceberg/arrow/s3/s3_properties.h @@ -40,6 +40,9 @@ struct S3Properties { static constexpr std::string_view kSecretAccessKey = "s3.secret-access-key"; /// AWS session token (for temporary credentials) static constexpr std::string_view kSessionToken = "s3.session-token"; + /// Epoch milliseconds at which a vended session token stops being valid + static constexpr std::string_view kSessionTokenExpiresAtMs = + "s3.session-token-expires-at-ms"; /// AWS region, standard Iceberg client property. static constexpr std::string_view kClientRegion = "client.region"; /// Custom endpoint override (for MinIO, LocalStack, etc.) diff --git a/src/iceberg/catalog/rest/json_serde.cc b/src/iceberg/catalog/rest/json_serde.cc index 3ce753f18c..098f2ba088 100644 --- a/src/iceberg/catalog/rest/json_serde.cc +++ b/src/iceberg/catalog/rest/json_serde.cc @@ -152,6 +152,25 @@ Result StorageCredentialFromJson(const nlohmann::json& json) return credential; } +/// \brief Reads the optional `storage-credentials` array shared by the +/// LoadTable and LoadCredentials responses. +Result> StorageCredentialsFromJson( + const nlohmann::json& json) { + std::vector credentials; + auto it = json.find(kStorageCredentials); + if (it == json.end() || it->is_null()) { + return credentials; + } + if (!it->is_array()) { + return JsonParseError("Cannot parse storage credentials from non-array"); + } + for (const auto& entry : *it) { + ICEBERG_ASSIGN_OR_RAISE(auto credential, StorageCredentialFromJson(entry)); + credentials.push_back(std::move(credential)); + } + return credentials; +} + template Result> KeyValueMapFromJson(const nlohmann::json& json, std::string_view key) { @@ -738,19 +757,24 @@ Result LoadTableResultFromJson(const nlohmann::json& json) { ICEBERG_ASSIGN_OR_RAISE(result.metadata, TableMetadataFromJson(metadata_json)); ICEBERG_ASSIGN_OR_RAISE(result.config, GetJsonValueOrDefault(json, kConfig)); - if (auto it = json.find(kStorageCredentials); it != json.end() && !it->is_null()) { - if (!it->is_array()) { - return JsonParseError("Cannot parse storage credentials from non-array"); - } - for (const auto& entry : *it) { - ICEBERG_ASSIGN_OR_RAISE(auto cred, StorageCredentialFromJson(entry)); - result.storage_credentials.push_back(std::move(cred)); - } - } + ICEBERG_ASSIGN_OR_RAISE(result.storage_credentials, StorageCredentialsFromJson(json)); ICEBERG_RETURN_UNEXPECTED(result.Validate()); return result; } +Result LoadCredentialsResponseFromJson( + const nlohmann::json& json) { + // Required here, unlike in LoadTable: reading a malformed response as "no + // credentials" would look like a refresh that succeeded and dropped them. + if (auto it = json.find(kStorageCredentials); it == json.end() || it->is_null()) { + return JsonParseError("Missing '{}'", kStorageCredentials); + } + LoadCredentialsResponse response; + ICEBERG_ASSIGN_OR_RAISE(response.storage_credentials, StorageCredentialsFromJson(json)); + ICEBERG_RETURN_UNEXPECTED(response.Validate()); + return response; +} + nlohmann::json ToJson(const ListNamespacesResponse& response) { nlohmann::json json; SetOptionalStringField(json, kNextPageToken, response.next_page_token); diff --git a/src/iceberg/catalog/rest/json_serde_internal.h b/src/iceberg/catalog/rest/json_serde_internal.h index 6e148e0d31..aedd2cd1a1 100644 --- a/src/iceberg/catalog/rest/json_serde_internal.h +++ b/src/iceberg/catalog/rest/json_serde_internal.h @@ -71,6 +71,10 @@ template <> ICEBERG_REST_EXPORT Result FromJson(const nlohmann::json& json); ICEBERG_REST_EXPORT Result ToJson(const LoadTableResult& model); +// Response-only model: a client never serializes it, so no ToJson. +ICEBERG_REST_EXPORT Result LoadCredentialsResponseFromJson( + const nlohmann::json& json); + ICEBERG_REST_EXPORT Result CreateTableRequestFromJson( const nlohmann::json& json); template <> diff --git a/src/iceberg/catalog/rest/rest_catalog.cc b/src/iceberg/catalog/rest/rest_catalog.cc index 349071f427..646147cfdb 100644 --- a/src/iceberg/catalog/rest/rest_catalog.cc +++ b/src/iceberg/catalog/rest/rest_catalog.cc @@ -42,6 +42,7 @@ #include "iceberg/catalog/rest/rest_util.h" #include "iceberg/catalog/rest/types.h" #include "iceberg/json_serde_internal.h" +#include "iceberg/logging/log_macros.h" #include "iceberg/metrics/metrics_reporters.h" #include "iceberg/partition_spec.h" #include "iceberg/result.h" @@ -508,12 +509,56 @@ Result> RestCatalog::TableAuthSession( std::move(contextual_session)); } +StorageCredentialRefresher RestCatalog::MakeCredentialRefresher( + const TableIdentifier& identifier, + std::shared_ptr table_session) const { + if (!supported_endpoints_.contains(Endpoint::TableCredentials())) { + // Not an error, but it surfaces much later as credentials expiring. + ICEBERG_LOG_DEBUG( + "Catalog does not advertise {}; vended credentials for '{}' will not be " + "refreshed", + Endpoint::TableCredentials().ToString(), ToString(identifier)); + return nullptr; + } + auto path = paths_->Credentials(identifier); + if (!path.has_value()) { + ICEBERG_LOG_WARN( + "Cannot build the credentials path for '{}' ({}); its vended credentials " + "will not be refreshed", + ToString(identifier), path.error().message); + return nullptr; + } + auto client = client_; + auto credentials_path = std::move(path.value()); + auto session = std::move(table_session); + // The catalog's destructor closes the session, and a table's FileIO can + // outlive the table keeping the catalog alive. No cycle: the catalog's own + // FileIO never gets a refresher. + auto catalog = shared_from_this(); + return [catalog, client, credentials_path, + session]() -> Result> { + ICEBERG_ASSIGN_OR_RAISE(const auto response, + client->Get(credentials_path, /*params=*/{}, /*headers=*/{}, + *TableErrorHandler::Instance(), *session)); + ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); + ICEBERG_ASSIGN_OR_RAISE(auto result, LoadCredentialsResponseFromJson(json)); + return std::move(result.storage_credentials); + }; +} + Result> RestCatalog::TableFileIO( - const SessionContext& /*context*/, + const SessionContext& /*context*/, const TableIdentifier& identifier, const std::unordered_map& table_config, - const std::vector& storage_credentials) const { + const std::vector& storage_credentials, + std::shared_ptr table_session) const { if (!table_config.empty() || !storage_credentials.empty()) { - return MakeTableFileIO(config_.configs(), table_config, storage_credentials); + // Only vended credentials expire, so only they need a refresher. + StorageCredentialRefresher refresher; + if (!storage_credentials.empty()) { + refresher = MakeCredentialRefresher(identifier, std::move(table_session)); + } + return MakeTableFileIO(config_.configs(), table_config, storage_credentials, + std::move(refresher)); } return file_io_; @@ -772,11 +817,12 @@ Result> RestCatalog::StageCreateTable( /*stage_create=*/true, *contextual_session)); auto table_config = std::move(result.config); auto storage_credentials = std::move(result.storage_credentials); - ICEBERG_ASSIGN_OR_RAISE(auto table_io, - TableFileIO(context, table_config, storage_credentials)); + // Before the FileIO: refreshing its credentials reuses the table session. ICEBERG_ASSIGN_OR_RAISE( auto table_session, TableAuthSession(identifier, table_config, std::move(contextual_session))); + ICEBERG_ASSIGN_OR_RAISE(auto table_io, TableFileIO(context, identifier, table_config, + storage_credentials, table_session)); ICEBERG_ASSIGN_OR_RAISE(auto reporter, MakeTableReporter(identifier, table_session)); auto table_catalog = std::make_shared( shared_from_this(), context, identifier, table_config, std::move(table_session), @@ -890,11 +936,12 @@ Result> RestCatalog::MakeTableFromLoadResult( std::shared_ptr contextual_session) { auto table_config = std::move(result.config); auto storage_credentials = std::move(result.storage_credentials); - ICEBERG_ASSIGN_OR_RAISE(auto table_io, - TableFileIO(context, table_config, storage_credentials)); + // Before the FileIO: refreshing its credentials reuses the table session. ICEBERG_ASSIGN_OR_RAISE( auto table_session, TableAuthSession(identifier, table_config, std::move(contextual_session))); + ICEBERG_ASSIGN_OR_RAISE(auto table_io, TableFileIO(context, identifier, table_config, + storage_credentials, table_session)); ICEBERG_ASSIGN_OR_RAISE(auto reporter, MakeTableReporter(identifier, table_session)); auto table_catalog = std::make_shared( shared_from_this(), context, identifier, table_config, table_session, table_io); diff --git a/src/iceberg/catalog/rest/rest_catalog.h b/src/iceberg/catalog/rest/rest_catalog.h index 65b0b5eaba..8b80c95a77 100644 --- a/src/iceberg/catalog/rest/rest_catalog.h +++ b/src/iceberg/catalog/rest/rest_catalog.h @@ -85,9 +85,16 @@ class ICEBERG_REST_EXPORT RestCatalog final std::shared_ptr contextual_session); Result> TableFileIO( - const SessionContext& context, + const SessionContext& context, const TableIdentifier& identifier, const std::unordered_map& table_config, - const std::vector& storage_credentials) const; + const std::vector& storage_credentials, + std::shared_ptr table_session) const; + + /// \brief Callback that reloads this table's vended credentials, or nullptr + /// when the catalog does not serve the LoadCredentials endpoint. + StorageCredentialRefresher MakeCredentialRefresher( + const TableIdentifier& identifier, + std::shared_ptr table_session) const; Result> ListNamespaces(const Namespace& ns, auth::AuthSession& session) const; diff --git a/src/iceberg/catalog/rest/rest_file_io.cc b/src/iceberg/catalog/rest/rest_file_io.cc index 4fc5122fe4..1dbdbec004 100644 --- a/src/iceberg/catalog/rest/rest_file_io.cc +++ b/src/iceberg/catalog/rest/rest_file_io.cc @@ -21,6 +21,7 @@ #include #include +#include #include #include "iceberg/catalog/rest/types.h" @@ -61,7 +62,8 @@ Result> MakeCatalogFileIO(const RestCatalogProperties& c Result> MakeTableFileIO( const std::unordered_map& catalog_config, const std::unordered_map& table_config, - const std::vector& storage_credentials) { + const std::vector& storage_credentials, + StorageCredentialRefresher refresher) { const auto default_properties = MergeFileIOProperties(catalog_config, table_config); ICEBERG_ASSIGN_OR_RAISE( auto io, MakeCatalogFileIO(RestCatalogProperties::FromMap(default_properties))); @@ -69,6 +71,10 @@ Result> MakeTableFileIO( if (storage_credentials.empty()) { return io; } else if (auto* credentialed = io->AsSupportsStorageCredentials()) { + // First, so the FileIO never briefly holds credentials it cannot replace. + if (refresher) { + credentialed->SetCredentialRefresher(std::move(refresher)); + } ICEBERG_RETURN_UNEXPECTED(credentialed->SetStorageCredentials(storage_credentials)); } else { return NotSupported("Configured FileIO does not support vended storage credentials"); diff --git a/src/iceberg/catalog/rest/rest_file_io.h b/src/iceberg/catalog/rest/rest_file_io.h index e2316c3e85..301bf62ff2 100644 --- a/src/iceberg/catalog/rest/rest_file_io.h +++ b/src/iceberg/catalog/rest/rest_file_io.h @@ -41,9 +41,13 @@ ICEBERG_REST_EXPORT Result> MakeCatalogFileIO( const RestCatalogProperties& config); /// \brief Build the configured table FileIO and apply storage credentials if present. +/// +/// \param refresher Optional callback used to replace the vended credentials +/// before they expire; ignored when the FileIO cannot tell when they do. ICEBERG_REST_EXPORT Result> MakeTableFileIO( const std::unordered_map& catalog_config, const std::unordered_map& table_config, - const std::vector& storage_credentials); + const std::vector& storage_credentials, + StorageCredentialRefresher refresher = nullptr); } // namespace iceberg::rest diff --git a/src/iceberg/catalog/rest/types.h b/src/iceberg/catalog/rest/types.h index 20a59fa595..9ff03bc06a 100644 --- a/src/iceberg/catalog/rest/types.h +++ b/src/iceberg/catalog/rest/types.h @@ -209,6 +209,23 @@ using CreateTableResponse = LoadTableResult; /// \brief Alias of LoadTableResult used as the body of LoadTableResponse using LoadTableResponse = LoadTableResult; +/// \brief Response body of the LoadCredentials API. +/// +/// Used to replace the credentials vended alongside a table before they expire. +struct ICEBERG_REST_EXPORT LoadCredentialsResponse { + std::vector storage_credentials; + + /// \brief Validates the LoadCredentialsResponse. + Status Validate() const { + for (const auto& credential : storage_credentials) { + ICEBERG_RETURN_UNEXPECTED(credential.Validate()); + } + return {}; + } + + bool operator==(const LoadCredentialsResponse& other) const = default; +}; + /// \brief Response body for listing namespaces. struct ICEBERG_REST_EXPORT ListNamespacesResponse { PageToken next_page_token; diff --git a/src/iceberg/file_io.h b/src/iceberg/file_io.h index 3ea4afa499..14eca99673 100644 --- a/src/iceberg/file_io.h +++ b/src/iceberg/file_io.h @@ -193,8 +193,16 @@ class ICEBERG_EXPORT SupportsStorageCredentials { virtual Status SetStorageCredentials( const std::vector& storage_credentials) = 0; - /// \brief Return currently installed storage credentials. - virtual const std::vector& credentials() const = 0; + /// \brief Return the storage credentials this FileIO holds. + /// + /// By value because a refresh may replace them concurrently. An + /// implementation that delegates may report what was installed on it. + virtual std::vector credentials() const = 0; + + /// \brief Install a callback that re-fetches credentials before they expire. + /// + /// Ignored by implementations that cannot tell when theirs expire. + virtual void SetCredentialRefresher(StorageCredentialRefresher /*refresher*/) {} }; } // namespace iceberg diff --git a/src/iceberg/resolving_file_io.cc b/src/iceberg/resolving_file_io.cc index ce91a52e8a..00848d4d73 100644 --- a/src/iceberg/resolving_file_io.cc +++ b/src/iceberg/resolving_file_io.cc @@ -61,8 +61,12 @@ Result ResolvingFileIO::FileIOForPath(std::string_view location) { FileIORegistry::Load(std::string(name), properties_)); // Forward all credentials; each implementation applies the prefixes it // understands. - if (!storage_credentials_.empty()) { - if (auto* credentialed = io->AsSupportsStorageCredentials()) { + if (auto* credentialed = io->AsSupportsStorageCredentials()) { + // Before the credentials, so the delegate can always replace them. + if (refresher_) { + credentialed->SetCredentialRefresher(refresher_); + } + if (!storage_credentials_.empty()) { ICEBERG_RETURN_UNEXPECTED( credentialed->SetStorageCredentials(storage_credentials_)); } @@ -117,8 +121,16 @@ Status ResolvingFileIO::SetStorageCredentials( return {}; } -const std::vector& ResolvingFileIO::credentials() const { +std::vector ResolvingFileIO::credentials() const { + std::lock_guard lock(mutex_); return storage_credentials_; } +void ResolvingFileIO::SetCredentialRefresher(StorageCredentialRefresher refresher) { + // Drop the cached delegates so they are rebuilt with the refresher. + std::lock_guard lock(mutex_); + refresher_ = std::move(refresher); + io_by_name_.clear(); +} + } // namespace iceberg diff --git a/src/iceberg/resolving_file_io.h b/src/iceberg/resolving_file_io.h index d96b0ca6e8..342d1c6ae6 100644 --- a/src/iceberg/resolving_file_io.h +++ b/src/iceberg/resolving_file_io.h @@ -47,9 +47,10 @@ namespace iceberg { /// them; each applies the prefixes it understands and ignores the rest. /// /// Lazy resolution is internally synchronized, so file operations may run -/// concurrently. Credentials are not: install them before sharing the instance, -/// since credentials() hands out a reference that SetStorageCredentials -/// replaces. +/// concurrently -- provided credentials and any refresher were installed before +/// the instance was shared, as the REST catalog does when it binds a table's +/// FileIO. Installing either drops the resolved implementations so they are +/// rebuilt, which would free one that a concurrent operation is still using. class ICEBERG_EXPORT ResolvingFileIO final : public FileIO, public SupportsStorageCredentials { public: @@ -70,7 +71,13 @@ class ICEBERG_EXPORT ResolvingFileIO final : public FileIO, Status SetStorageCredentials( const std::vector& storage_credentials) override; - const std::vector& credentials() const override; + /// \brief Return the credentials installed on this resolver. + /// + /// Not necessarily the ones in use: a resolved implementation refreshes its + /// own without reporting back. + std::vector credentials() const override; + + void SetCredentialRefresher(StorageCredentialRefresher refresher) override; SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; } @@ -79,9 +86,10 @@ class ICEBERG_EXPORT ResolvingFileIO final : public FileIO, Result FileIOForPath(std::string_view location); std::unordered_map properties_; - // Guards lazy resolution; set credentials before sharing across threads. - std::mutex mutex_; + // Guards lazy resolution and credential state. + mutable std::mutex mutex_; std::vector storage_credentials_; + StorageCredentialRefresher refresher_; std::unordered_map, StringHash, StringEqual> io_by_name_; }; diff --git a/src/iceberg/storage_credential.h b/src/iceberg/storage_credential.h index 746bad086c..fb1708efaf 100644 --- a/src/iceberg/storage_credential.h +++ b/src/iceberg/storage_credential.h @@ -22,8 +22,10 @@ /// \file iceberg/storage_credential.h /// \brief Define storage credential metadata. +#include #include #include +#include #include "iceberg/iceberg_export.h" #include "iceberg/result.h" @@ -48,4 +50,11 @@ struct ICEBERG_EXPORT StorageCredential { bool operator==(const StorageCredential& other) const = default; }; +/// \brief Re-fetches the storage credentials that are currently valid. +/// +/// Lets a FileIO replace expiring credentials without knowing how they are +/// delivered. Returns the whole vended list. +using StorageCredentialRefresher = + std::function>()>; + } // namespace iceberg diff --git a/src/iceberg/test/arrow_s3_file_io_test.cc b/src/iceberg/test/arrow_s3_file_io_test.cc index aa949ab129..0959ac22ec 100644 --- a/src/iceberg/test/arrow_s3_file_io_test.cc +++ b/src/iceberg/test/arrow_s3_file_io_test.cc @@ -18,12 +18,17 @@ */ #include +#include +#include +#include #include #include #include +#include #include #include #include +#include #include #include #include @@ -104,6 +109,9 @@ class ArrowS3FileIOTest : public ::testing::Test { protected: #if ICEBERG_S3_ENABLED static void SetUpTestSuite() { + // Off EC2 every S3 client build waits for this to time out. Not overwritten, + // so a run that does want those credentials can still ask. + ::setenv("AWS_EC2_METADATA_DISABLED", "true", /*overwrite=*/0); auto io = MakeS3FileIO({}); ASSERT_THAT(io, IsOk()); } @@ -130,10 +138,26 @@ class ArrowS3FileIOTest : public ::testing::Test { std::optional base_uri_; }; -bool HasWarning(const CapturingLogger& logger) { +bool HasWarning(const CapturingLogger& logger, std::string_view substring = {}) { const auto records = logger.records(); - return std::ranges::any_of( - records, [](const LogMessage& record) { return record.level == LogLevel::kWarn; }); + return std::ranges::any_of(records, [substring](const LogMessage& record) { + return record.level == LogLevel::kWarn && + record.message.find(substring) != std::string::npos; + }); +} + +constexpr auto kOutlastsAShortenedBackoff = std::chrono::milliseconds(1200); + +std::vector ExpiringCredentials(std::chrono::milliseconds valid_for, + std::string_view access_key) { + const auto expires_at = std::chrono::duration_cast( + (std::chrono::system_clock::now() + valid_for).time_since_epoch()); + return {{.prefix = "s3", + .config = {{std::string(S3Properties::kAccessKeyId), std::string(access_key)}, + {std::string(S3Properties::kSecretAccessKey), "secret"}, + {std::string(S3Properties::kSessionToken), "token"}, + {std::string(S3Properties::kSessionTokenExpiresAtMs), + std::to_string(expires_at.count())}}}}; } Status CheckReadWrite(FileIO& io, const std::string& object_uri, @@ -226,6 +250,342 @@ TEST_F(ArrowS3FileIOTest, WarnsWhenNoCredentialApplies) { EXPECT_TRUE(HasWarning(*logger)); } +TEST_F(ArrowS3FileIOTest, RefreshesCredentialsCloseToExpiry) { + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + const auto refreshed = ExpiringCredentials(std::chrono::hours(1), "refreshed-key"); + int refresh_calls = 0; + credentialed->SetCredentialRefresher([&]() -> Result> { + ++refresh_calls; + return refreshed; + }); + ASSERT_THAT(credentialed->SetStorageCredentials( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key")), + IsOk()); + + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(refresh_calls, 1); + EXPECT_EQ(credentialed->credentials(), refreshed); + + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(refresh_calls, 1); +} + +TEST_F(ArrowS3FileIOTest, DoesNotRefreshCredentialsThatAreNotCloseToExpiry) { + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + int refresh_calls = 0; + credentialed->SetCredentialRefresher([&]() -> Result> { + ++refresh_calls; + return std::vector{}; + }); + + ASSERT_THAT(credentialed->SetStorageCredentials( + ExpiringCredentials(std::chrono::hours(1), "access-key")), + IsOk()); + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(refresh_calls, 0); + + const std::vector static_credentials = { + {.prefix = "s3", + .config = {{std::string(S3Properties::kAccessKeyId), "access-key"}, + {std::string(S3Properties::kSecretAccessKey), "secret"}}}}; + ASSERT_THAT(credentialed->SetStorageCredentials(static_credentials), IsOk()); + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(refresh_calls, 0); +} + +TEST_F(ArrowS3FileIOTest, RefreshesOnceWhenOperationsRaceForIt) { + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + std::atomic refresh_calls = 0; + credentialed->SetCredentialRefresher([&]() -> Result> { + ++refresh_calls; + return ExpiringCredentials(std::chrono::hours(1), "refreshed-key"); + }); + ASSERT_THAT(credentialed->SetStorageCredentials( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key")), + IsOk()); + + constexpr int kThreads = 8; + std::atomic failures = 0; + std::vector threads; + threads.reserve(kThreads); + for (int i = 0; i < kThreads; ++i) { + threads.emplace_back([&] { + for (int op = 0; op < 4; ++op) { + if (!result.value()->NewInputFile("s3://bucket/key").has_value()) { + ++failures; + } + } + }); + } + for (auto& thread : threads) { + thread.join(); + } + + EXPECT_EQ(failures, 0); + EXPECT_EQ(refresh_calls, 1); +} + +TEST_F(ArrowS3FileIOTest, RefreshesOnceWhenCredentialsHaveExpired) { + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + std::mutex mutex; + std::condition_variable cv; + bool refresh_started = false; + bool release_refresh = false; + std::atomic refresh_calls = 0; + + credentialed->SetCredentialRefresher([&]() -> Result> { + ++refresh_calls; + std::unique_lock lock(mutex); + refresh_started = true; + cv.notify_all(); + cv.wait(lock, [&] { return release_refresh; }); + return ExpiringCredentials(std::chrono::hours(1), "refreshed-key"); + }); + ASSERT_THAT(credentialed->SetStorageCredentials( + ExpiringCredentials(-std::chrono::minutes(1), "expired-key")), + IsOk()); + + std::thread winner( + [&] { EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); }); + { + std::unique_lock lock(mutex); + cv.wait(lock, [&] { return refresh_started; }); + } + + bool loser_ready = false; + std::thread loser([&] { + { + std::lock_guard lock(mutex); + loser_ready = true; + } + cv.notify_all(); + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + }); + { + std::unique_lock lock(mutex); + cv.wait(lock, [&] { return loser_ready; }); + release_refresh = true; + } + cv.notify_all(); + winner.join(); + loser.join(); + + EXPECT_EQ(refresh_calls, 1); + EXPECT_THAT(credentialed->credentials(), ::testing::Not(::testing::IsEmpty())); +} + +TEST_F(ArrowS3FileIOTest, RefreshDoesNotUndoCredentialsInstalledWhileItRan) { + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + std::mutex mutex; + std::condition_variable cv; + bool refresh_started = false; + bool release_refresh = false; + + credentialed->SetCredentialRefresher([&]() -> Result> { + std::unique_lock lock(mutex); + refresh_started = true; + cv.notify_all(); + cv.wait(lock, [&] { return release_refresh; }); + return ExpiringCredentials(std::chrono::hours(1), "fetched-by-refresh"); + }); + ASSERT_THAT(credentialed->SetStorageCredentials( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key")), + IsOk()); + + std::thread operation( + [&] { EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); }); + + const auto installed = + ExpiringCredentials(std::chrono::hours(2), "installed-meanwhile"); + { + std::unique_lock lock(mutex); + cv.wait(lock, [&] { return refresh_started; }); + } + ASSERT_THAT(credentialed->SetStorageCredentials(installed), IsOk()); + { + std::lock_guard lock(mutex); + release_refresh = true; + } + cv.notify_all(); + operation.join(); + + EXPECT_EQ(credentialed->credentials(), installed); +} + +TEST_F(ArrowS3FileIOTest, RefreshesSessionCredentialsWithoutAUsableExpiry) { + for (std::string_view expiry : {"", "not-a-number"}) { + SCOPED_TRACE(expiry); + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + int refresh_calls = 0; + credentialed->SetCredentialRefresher([&]() -> Result> { + ++refresh_calls; + return ExpiringCredentials(std::chrono::hours(1), "refreshed-key"); + }); + + auto logger = std::make_shared(); + ScopedDefaultLogger scoped(logger); + std::unordered_map config = { + {std::string(S3Properties::kAccessKeyId), "access-key"}, + {std::string(S3Properties::kSecretAccessKey), "secret"}, + {std::string(S3Properties::kSessionToken), "token"}}; + if (!expiry.empty()) { + config[std::string(S3Properties::kSessionTokenExpiresAtMs)] = std::string(expiry); + } + ASSERT_THAT(credentialed->SetStorageCredentials( + {{.prefix = "s3", .config = std::move(config)}}), + IsOk()); + EXPECT_TRUE(HasWarning(*logger, "session token")); + + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(refresh_calls, 1); + } +} + +TEST_F(ArrowS3FileIOTest, BacksOffWhenReplacementsAlsoLackAnExpiry) { + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + const std::vector undatable = { + {.prefix = "s3", + .config = {{std::string(S3Properties::kAccessKeyId), "access-key"}, + {std::string(S3Properties::kSecretAccessKey), "secret"}, + {std::string(S3Properties::kSessionToken), "token"}}}}; + int refresh_calls = 0; + credentialed->SetCredentialRefresher([&]() -> Result> { + ++refresh_calls; + return undatable; + }); + ASSERT_THAT(credentialed->SetStorageCredentials(undatable), IsOk()); + + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + ASSERT_EQ(refresh_calls, 1); + + std::this_thread::sleep_for(kOutlastsAShortenedBackoff); + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(refresh_calls, 1); +} + +TEST_F(ArrowS3FileIOTest, IgnoresUnparseableExpiry) { + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + int refresh_calls = 0; + credentialed->SetCredentialRefresher([&]() -> Result> { + ++refresh_calls; + return std::vector{}; + }); + + auto logger = std::make_shared(); + ScopedDefaultLogger scoped(logger); + std::unordered_map config = { + {std::string(S3Properties::kAccessKeyId), "access-key"}, + {std::string(S3Properties::kSecretAccessKey), "secret"}, + {std::string(S3Properties::kSessionTokenExpiresAtMs), "not-a-number"}}; + const std::vector credentials = { + {.prefix = "s3", .config = std::move(config)}}; + ASSERT_THAT(credentialed->SetStorageCredentials(credentials), IsOk()); + EXPECT_EQ(credentialed->credentials(), credentials); + + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(refresh_calls, 0); +} + +TEST_F(ArrowS3FileIOTest, BacksOffWhenTheReplacementIsAlsoCloseToExpiry) { + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + int refresh_calls = 0; + credentialed->SetCredentialRefresher([&]() -> Result> { + ++refresh_calls; + return ExpiringCredentials(std::chrono::minutes(1), "short-lived-key"); + }); + ASSERT_THAT(credentialed->SetStorageCredentials( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key")), + IsOk()); + + for (int i = 0; i < 3; ++i) { + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + } + EXPECT_EQ(refresh_calls, 1); +} + +TEST_F(ArrowS3FileIOTest, KeepsCredentialsWhenRefreshFails) { + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + int refresh_calls = 0; + credentialed->SetCredentialRefresher([&]() -> Result> { + ++refresh_calls; + return NotFound("catalog unreachable"); + }); + const auto expiring = ExpiringCredentials(std::chrono::minutes(1), "expiring-key"); + ASSERT_THAT(credentialed->SetStorageCredentials(expiring), IsOk()); + + auto logger = std::make_shared(); + ScopedDefaultLogger scoped(logger); + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(credentialed->credentials(), expiring); + EXPECT_TRUE(HasWarning(*logger, "Failed to refresh")); + + EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(refresh_calls, 1); +} + +TEST_F(ArrowS3FileIOTest, DeleteFilesDispatchesAcrossCredentialPrefixes) { + auto result = MakeS3FileIO({}); + ASSERT_THAT(result, IsOk()); + auto* credentialed = result.value()->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + auto credential = [](std::string_view prefix, std::string_view access_key) { + return StorageCredential{ + .prefix = std::string(prefix), + .config = {{std::string(S3Properties::kAccessKeyId), std::string(access_key)}, + {std::string(S3Properties::kSecretAccessKey), "secret"}}}; + }; + ASSERT_THAT(credentialed->SetStorageCredentials({credential("s3://bucket-a", "key-a"), + credential("s3://bucket-b", "key-b")}), + IsOk()); + + auto status = result.value()->DeleteFiles({"s3://bucket-a/%ZZ.parquet", + "s3://bucket-a/second.parquet", + "s3://bucket-b/other.parquet"}); + EXPECT_FALSE(status.has_value()); + EXPECT_THAT(status, HasErrorMessage("Cannot parse URI")); +} + TEST_F(ArrowS3FileIOTest, RejectsIncompleteStaticCredentials) { auto result = MakeS3FileIO({{std::string(S3Properties::kAccessKeyId), "access-key-only"}}); diff --git a/src/iceberg/test/resolving_file_io_test.cc b/src/iceberg/test/resolving_file_io_test.cc index 97c1d78731..addad42643 100644 --- a/src/iceberg/test/resolving_file_io_test.cc +++ b/src/iceberg/test/resolving_file_io_test.cc @@ -55,14 +55,27 @@ class RecordingCredentialedFileIO : public RecordingFileIO, return {}; } - const std::vector& credentials() const override { - return credentials_; + std::vector credentials() const override { return credentials_; } + + void SetCredentialRefresher(StorageCredentialRefresher refresher) override { + refresher_ = std::move(refresher); } SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; } + Status Refresh() { + if (!refresher_) { + return NotFound("no refresher installed"); + } + ICEBERG_ASSIGN_OR_RAISE(auto refreshed, refresher_()); + return SetStorageCredentials(refreshed); + } + + bool has_refresher() const { return static_cast(refresher_); } + private: std::vector credentials_; + StorageCredentialRefresher refresher_; }; // File-scope recording state: registry factories are process-global, so they @@ -180,4 +193,41 @@ TEST(ResolvingFileIOTest, ForwardsAllCredentialsToResolvedImplementations) { ASSERT_NE(last_local_io, nullptr); } +TEST(ResolvingFileIOTest, ForwardsCredentialRefresherToResolvedImplementations) { + RegisterRecordingFileIOs(); + ResolvingFileIO io({}); + + std::vector refreshed = {{.prefix = "s3", .config = {{"k2", "v2"}}}}; + io.SetCredentialRefresher( + [&]() -> Result> { return refreshed; }); + EXPECT_THAT(io.SetStorageCredentials({{.prefix = "s3", .config = {{"k1", "v1"}}}}), + IsOk()); + + (void)io.NewInputFile("s3://bucket/db/table/data/file.parquet"); + ASSERT_NE(last_s3_io, nullptr); + ASSERT_TRUE(last_s3_io->has_refresher()); + EXPECT_THAT(last_s3_io->Refresh(), IsOk()); + EXPECT_EQ(last_s3_io->credentials(), refreshed); +} + +TEST(ResolvingFileIOTest, RebuildsResolvedImplementationsForALaterRefresher) { + RegisterRecordingFileIOs(); + ResolvingFileIO io({}); + + EXPECT_THAT(io.SetStorageCredentials({{.prefix = "s3", .config = {{"k1", "v1"}}}}), + IsOk()); + (void)io.NewInputFile("s3://bucket/db/table/data/file.parquet"); + ASSERT_NE(last_s3_io, nullptr); + EXPECT_FALSE(last_s3_io->has_refresher()); + EXPECT_EQ(s3_factory_calls, 1); + + io.SetCredentialRefresher([]() -> Result> { + return std::vector{}; + }); + (void)io.NewInputFile("s3://bucket/db/table/data/other.parquet"); + EXPECT_EQ(s3_factory_calls, 2); + ASSERT_NE(last_s3_io, nullptr); + EXPECT_TRUE(last_s3_io->has_refresher()); +} + } // namespace iceberg diff --git a/src/iceberg/test/rest_file_io_test.cc b/src/iceberg/test/rest_file_io_test.cc index 6e584ea220..ecb3264dd0 100644 --- a/src/iceberg/test/rest_file_io_test.cc +++ b/src/iceberg/test/rest_file_io_test.cc @@ -51,6 +51,7 @@ class MockFileIO : public FileIO { std::vector captured_storage_credentials; std::unordered_map captured_file_io_properties; +StorageCredentialRefresher captured_refresher; class MockCredentialedFileIO : public MockFileIO, public SupportsStorageCredentials { public: @@ -60,10 +61,14 @@ class MockCredentialedFileIO : public MockFileIO, public SupportsStorageCredenti return {}; } - const std::vector& credentials() const override { + std::vector credentials() const override { return captured_storage_credentials; } + void SetCredentialRefresher(StorageCredentialRefresher refresher) override { + captured_refresher = std::move(refresher); + } + SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; } }; @@ -166,6 +171,35 @@ TEST(RestFileIOTest, TableFileIOMergesConfigAndCredentials) { EXPECT_EQ(credentialed->credentials(), captured_storage_credentials); } +TEST(RestFileIOTest, InstallsCredentialRefresherOnlyForVendedCredentials) { + const std::string custom_impl = "rest-file-io-test-refresher"; + FileIORegistry::Register( + custom_impl, + [](const std::unordered_map& /*properties*/) + -> Result> { + return std::make_unique(); + }); + const std::unordered_map table_config{ + {"io-impl", custom_impl}}; + + captured_refresher = nullptr; + std::vector refreshed = {{.prefix = "s3", .config = {{"k", "v2"}}}}; + auto result = MakeTableFileIO( + {}, table_config, {{.prefix = "s3", .config = {{"k", "v1"}}}}, + [&]() -> Result> { return refreshed; }); + ASSERT_THAT(result, IsOk()); + ASSERT_TRUE(captured_refresher); + EXPECT_THAT(captured_refresher(), IsOk()); + EXPECT_EQ(captured_refresher().value(), refreshed); + + captured_refresher = nullptr; + result = MakeTableFileIO( + {}, table_config, /*storage_credentials=*/{}, + [&]() -> Result> { return refreshed; }); + ASSERT_THAT(result, IsOk()); + EXPECT_FALSE(captured_refresher); +} + TEST(RestFileIOTest, TableImplOverridesWarehouseScheme) { captured_file_io_properties.clear(); FileIORegistry::Register( diff --git a/src/iceberg/test/rest_json_serde_test.cc b/src/iceberg/test/rest_json_serde_test.cc index ec41e4a668..1dcb0238d7 100644 --- a/src/iceberg/test/rest_json_serde_test.cc +++ b/src/iceberg/test/rest_json_serde_test.cc @@ -1229,6 +1229,30 @@ INSTANTIATE_TEST_SUITE_P( return info.param.test_name; }); +TEST(LoadCredentialsResponseTest, ParsesVendedCredentials) { + auto json = nlohmann::json::parse( + R"({"storage-credentials":[{"prefix":"s3","config":{"s3.access-key-id":"AKIAtest"}}]})"); + auto response = LoadCredentialsResponseFromJson(json); + ASSERT_THAT(response, IsOk()); + EXPECT_EQ(response->storage_credentials, + (std::vector{ + {.prefix = "s3", .config = {{"s3.access-key-id", "AKIAtest"}}}})); +} + +TEST(LoadCredentialsResponseTest, RejectsResponseWithoutCredentials) { + auto response = LoadCredentialsResponseFromJson(nlohmann::json::parse("{}")); + EXPECT_THAT(response, IsError(ErrorKind::kJsonParseError)); + EXPECT_THAT(response, HasErrorMessage("Missing 'storage-credentials'")); +} + +TEST(LoadCredentialsResponseTest, RejectsNonArrayCredentials) { + auto response = LoadCredentialsResponseFromJson( + nlohmann::json::parse(R"({"storage-credentials":"oops"})")); + EXPECT_THAT(response, IsError(ErrorKind::kJsonParseError)); + EXPECT_THAT(response, + HasErrorMessage("Cannot parse storage credentials from non-array")); +} + DECLARE_ROUNDTRIP_TEST(CommitTableRequest) INSTANTIATE_TEST_SUITE_P(