From 098ba85287f04f1b6171a34fa19d05e813250228 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Tue, 13 Oct 2026 00:42:20 -0700 Subject: [PATCH 1/7] feat(io): refresh vended storage credentials before they expire Vended storage credentials expire, so new file operations on a reused table could eventually fail. ArrowS3FileIO now refreshes them through the StorageCredentialProvider installed by InitializeStorageCredentials, starting five minutes before the earliest applicable s3.session-token-expires-at-ms. Later SetStorageCredentials calls keep the provider. Each vended S3 client signs through a VendedCredentialsProvider, which the AWS SDK asks before every request: it refreshes once due and returns the keys the FileIO currently vends for the client's prefix, as Java's VendedCredentialsProvider does. Handles and streams opened before a refresh therefore use the new keys, including those on clients a later refresh replaced. A refresh that changes only keys keeps the existing clients. A request cannot build clients, since Arrow holds its client registry for the request and building one locks it exclusively. A refresh it starts that needs new clients is left to the next file operation, and it never waits on another refresh. A refresh is installed only if every S3 credential is usable: a session token must carry a parseable expiry that has not passed. Otherwise, or if loading fails, the current set is kept and the refresh backs off for up to 30 seconds. Concurrent operations share one refresh, and operations with expired credentials wait up to 10 seconds for one in flight. A concurrent credential install supersedes any older refresh result. Once a delegate's credentials have expired and could not be refreshed, creating a handle or deleting under its prefix fails with TokenExpired instead of handing out credentials known to be unusable. The LoadCredentials request had no deadline, so a stalled catalog could block a file operation indefinitely. HttpClient::Get takes an optional timeout, after which the request fails with an IOError, and RestStorageCredentialProvider uses 10 seconds. The bundle now links aws-cpp-sdk-core when S3 is enabled. --- mkdocs/docs/file-io.md | 1 + src/iceberg/CMakeLists.txt | 10 + src/iceberg/arrow/s3/arrow_s3_file_io.cc | 572 +++++++++++++++++- src/iceberg/arrow/s3/s3_properties.h | 3 + src/iceberg/catalog/rest/http_client.cc | 13 +- src/iceberg/catalog/rest/http_client.h | 8 +- src/iceberg/catalog/rest/rest_catalog.cc | 31 +- src/iceberg/catalog/rest/rest_table_scan.cc | 3 +- src/iceberg/test/CMakeLists.txt | 1 + src/iceberg/test/arrow_s3_file_io_test.cc | 552 ++++++++++++++++- src/iceberg/test/auth_manager_test.cc | 4 +- src/iceberg/test/http_client_test.cc | 76 +++ src/iceberg/test/rest_arrow_file_io_test.cc | 58 ++ .../test/rest_catalog_integration_test.cc | 4 +- src/iceberg/test/rest_catalog_unit_test.cc | 43 +- src/iceberg/test/rest_table_scan_test.cc | 24 +- 16 files changed, 1324 insertions(+), 79 deletions(-) create mode 100644 src/iceberg/test/http_client_test.cc diff --git a/mkdocs/docs/file-io.md b/mkdocs/docs/file-io.md index 63ce34de9a..62f998b1fc 100644 --- a/mkdocs/docs/file-io.md +++ b/mkdocs/docs/file-io.md @@ -63,6 +63,7 @@ each file location's scheme. | `s3.access-key-id` | `admin` | Static access key ID; must be set together with the secret key | | `s3.secret-access-key` | `password` | Static secret access key | | `s3.session-token` | `AQoDYXdzEJr...` | Session token, for temporary credentials. Ignored unless both static keys are set | +| `s3.session-token-expires-at-ms` | `1767225600000` | When the session token expires, in epoch milliseconds. Vended credentials are refreshed from the catalog five minutes before, including for files already open; once they have expired and cannot be refreshed, new operations under their prefix fail | | `client.region` | `us-east-1` | Region to sign requests for | | `s3.endpoint` | `https://127.0.0.1:9000` | Endpoint to use instead of the AWS one. When absent, the `AWS_ENDPOINT_URL_S3` / `AWS_ENDPOINT_URL` environment variables are consulted | | `s3.path-style-access` | `true` | Address buckets as a path (`endpoint/bucket`) instead of a virtual host (`bucket.endpoint`). Only takes effect together with a custom endpoint | diff --git a/src/iceberg/CMakeLists.txt b/src/iceberg/CMakeLists.txt index 717f1da4b2..ebd9a67237 100644 --- a/src/iceberg/CMakeLists.txt +++ b/src/iceberg/CMakeLists.txt @@ -371,6 +371,16 @@ if(ICEBERG_BUILD_BUNDLE) endforeach() endif() + # The S3 FileIO signs with its own AWS credentials provider. + if(ICEBERG_S3) + list(APPEND ICEBERG_BUNDLE_STATIC_BUILD_INTERFACE_LIBS aws-cpp-sdk-core) + list(APPEND ICEBERG_BUNDLE_SHARED_BUILD_INTERFACE_LIBS aws-cpp-sdk-core) + if(NOT ICEBERG_AWSSDK_BUNDLED) + list(APPEND ICEBERG_BUNDLE_STATIC_INSTALL_INTERFACE_LIBS aws-cpp-sdk-core) + list(APPEND ICEBERG_BUNDLE_SHARED_INSTALL_INTERFACE_LIBS aws-cpp-sdk-core) + endif() + endif() + add_iceberg_lib(iceberg_bundle SOURCES ${ICEBERG_BUNDLE_SOURCES} diff --git a/src/iceberg/arrow/s3/arrow_s3_file_io.cc b/src/iceberg/arrow/s3/arrow_s3_file_io.cc index 24b89bdd61..3c46cb9e41 100644 --- a/src/iceberg/arrow/s3/arrow_s3_file_io.cc +++ b/src/iceberg/arrow/s3/arrow_s3_file_io.cc @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -40,6 +41,8 @@ #include #if ICEBERG_S3_ENABLED # include +# include +# include #endif #include "iceberg/arrow/arrow_io_internal.h" @@ -295,10 +298,17 @@ Result<::arrow::fs::S3Options> ConfigureS3Options( namespace { +// `credentials_provider`, if set, signs requests instead of the keys in +// `properties`. Result> BuildArrowS3FileSystem( - const std::unordered_map& properties) { + const std::unordered_map& properties, + std::shared_ptr credentials_provider = nullptr) { ICEBERG_RETURN_UNEXPECTED(EnsureS3Initialized()); ICEBERG_ASSIGN_OR_RAISE(auto options, ConfigureS3Options(properties)); + if (credentials_provider != nullptr) { + options.credentials_provider = std::move(credentials_provider); + options.credentials_kind = ::arrow::fs::S3CredentialsKind::Explicit; + } ICEBERG_ARROW_ASSIGN_OR_RETURN(auto fs, ::arrow::fs::S3FileSystem::Make(options)); return std::shared_ptr<::arrow::fs::FileSystem>(std::move(fs)); } @@ -314,6 +324,206 @@ 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: a provider may not limit how long its request takes. +constexpr auto kExpiredCredentialWait = std::chrono::seconds(10); + +// When `credential`'s session token stops being valid: nullopt for static +// keys, which never expire, or an error if its expiry is missing or unusable. +Result> ExpiryOf( + const StorageCredential& credential) { + if (FindProperty(credential.config, S3Properties::kSessionToken) == nullptr) { + return std::nullopt; + } + const auto* value = + FindProperty(credential.config, S3Properties::kSessionTokenExpiresAtMs); + if (value == nullptr) { + return InvalidArgument("Credential \"{}\" has a session token but no \"{}\"", + credential.prefix, S3Properties::kSessionTokenExpiresAtMs); + } + auto millis = StringUtils::ParseNumber(*value); + if (!millis.has_value()) { + return InvalidArgument( + "Credential \"{}\" has a session token but an unparseable \"{}\" value \"{}\"", + credential.prefix, S3Properties::kSessionTokenExpiresAtMs, *value); + } + // 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) { + return InvalidArgument("Credential \"{}\" has an out-of-range \"{}\" value \"{}\"", + credential.prefix, S3Properties::kSessionTokenExpiresAtMs, + *value); + } + return std::chrono::system_clock::time_point(std::chrono::milliseconds(*millis)); +} + +// When installed credentials need replacing. +struct CredentialsExpiry { + // The earliest expiry known, if any. + std::optional expires_at; + // A session token has no usable expiry. It is replaced at once rather than + // used until it fails, as Java does, but not reported as expired. + bool unknown = false; +}; + +CredentialsExpiry ExpiryOf(const std::vector& credentials) { + CredentialsExpiry expiry; + for (const auto& credential : credentials) { + if (!IsS3CredentialPrefix(credential.prefix)) { + continue; + } + auto expires_at = ExpiryOf(credential); + if (!expires_at.has_value()) { + ICEBERG_LOG_WARN("{}", expires_at.error().message); + expiry.unknown = true; + } else if (expires_at->has_value() && + (!expiry.expires_at.has_value() || **expires_at < *expiry.expires_at)) { + expiry.expires_at = **expires_at; + } + } + return expiry; +} + +std::vector S3CredentialsOf( + const std::vector& credentials) { + std::vector s3_credentials; + for (const auto& credential : credentials) { + if (IsS3CredentialPrefix(credential.prefix)) { + s3_credentials.push_back(&credential); + } + } + return s3_credentials; +} + +// A refresh must not replace working credentials with ones that cannot be used +// or whose expiry is unknown, which would also defer the next refresh. +Status ValidateRefreshed(const std::vector& credentials) { + const auto now = std::chrono::system_clock::now(); + for (const auto& credential : credentials) { + ICEBERG_RETURN_UNEXPECTED(credential.Validate()); + } + const auto s3_credentials = S3CredentialsOf(credentials); + if (s3_credentials.empty()) { + // Installing this would drop working credentials for whatever ambient + // identity the AWS chain finds. Java refuses an empty list too. + return NotFound("Refreshed credentials contain no S3-compatible prefix"); + } + for (const auto* credential : s3_credentials) { + ICEBERG_ASSIGN_OR_RAISE(auto expires_at, ExpiryOf(*credential)); + if (expires_at.has_value() && *expires_at <= now) { + return InvalidArgument("Credential \"{}\" has already expired", credential->prefix); + } + } + return {}; +} + +// Whether `config` carries vended keys, which reach the S3 client through a +// VendedCredentialsProvider. +bool HasKeys(const std::unordered_map& config) { + return FindProperty(config, S3Properties::kAccessKeyId) != nullptr && + FindProperty(config, S3Properties::kSecretAccessKey) != nullptr; +} + +// Requires HasKeys(config). +Aws::Auth::AWSCredentials AwsCredentialsOf( + const std::unordered_map& config) { + Aws::Auth::AWSCredentials credentials; + credentials.SetAWSAccessKeyId( + FindProperty(config, S3Properties::kAccessKeyId)->c_str()); + credentials.SetAWSSecretKey( + FindProperty(config, S3Properties::kSecretAccessKey)->c_str()); + if (const auto* session_token = FindProperty(config, S3Properties::kSessionToken); + session_token != nullptr) { + credentials.SetSessionToken(session_token->c_str()); + } + return credentials; +} + +// Whether `refreshed` serves the same prefix with the same settings as +// `current`, so that only its keys need to reach the existing S3 client. +bool DiffersOnlyInKeys(const StorageCredential& current, + const StorageCredential& refreshed) { + if (current.prefix != refreshed.prefix || !HasKeys(refreshed.config)) { + return false; + } + const auto settings = [](std::unordered_map config) { + for (const auto key : + {S3Properties::kAccessKeyId, S3Properties::kSecretAccessKey, + S3Properties::kSessionToken, S3Properties::kSessionTokenExpiresAtMs}) { + config.erase(std::string(key)); + } + return config; + }; + return settings(current.config) == settings(refreshed.config); +} + +class ArrowS3FileIO; + +// Where the S3 clients of an ArrowS3FileIO get their keys. Open handles keep +// clients alive past the FileIO, which therefore detaches on destruction. +class KeySource { + public: + explicit KeySource(ArrowS3FileIO* io) : io_(io) {} + + void Detach() { + std::unique_lock lock(mutex_); + io_ = nullptr; + } + + /// \brief The keys the FileIO vends for the canonical `prefix`, after the + /// refresh a request may run; nullopt once it is gone or vends none for + /// `prefix`. + std::optional KeysFor(std::string_view prefix); + + private: + std::shared_mutex mutex_; + ArrowS3FileIO* io_; +}; + +// Signs an S3 client's requests with the keys its FileIO currently vends for +// the client's prefix. The AWS SDK asks before each request, so a refresh +// reaches handles and streams opened before it, even those on a client it +// replaced, as each client's VendedCredentialsProvider does in Java. Without +// current keys for the prefix, the last ones stay. +class VendedCredentialsProvider final : public Aws::Auth::AWSCredentialsProvider { + public: + VendedCredentialsProvider(std::shared_ptr source, std::string prefix, + Aws::Auth::AWSCredentials credentials) + : source_(std::move(source)), + prefix_(std::move(prefix)), + credentials_(std::move(credentials)) {} + + Aws::Auth::AWSCredentials GetAWSCredentials() override { + auto current = source_->KeysFor(prefix_); + std::lock_guard lock(mutex_); + if (current.has_value()) { + credentials_ = std::move(current).value(); + } + return credentials_; + } + + private: + const std::shared_ptr source_; + const std::string prefix_; + std::mutex mutex_; + Aws::Auth::AWSCredentials credentials_; +}; + class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { public: ArrowS3FileIO(std::shared_ptr<::arrow::fs::FileSystem> arrow_fs, @@ -323,6 +533,8 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { default_properties_(std::move(default_properties)), delete_parallelism_(delete_parallelism) {} + ~ArrowS3FileIO() override { key_source_->Detach(); } + Result> NewInputFile(std::string file_location) override; Result> NewInputFile(std::string file_location, @@ -342,18 +554,38 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { return storage_credentials_; } + Status InitializeStorageCredentials( + const std::vector& storage_credentials, + std::shared_ptr provider) override; + SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; } private: + friend class KeySource; + + /// \brief The S3 client serving the locations under `prefix`. + struct Delegate { + std::string prefix; + std::shared_ptr file_io; + // Whether the client signs with the keys vended for `prefix`, which it + // reads through a VendedCredentialsProvider. + bool vended_keys = false; + // When those keys stop being valid, if known. + std::optional expires_at; + }; + + using DelegatesByPrefix = std::vector; + /// \brief Delegate serving `location`, pinned by the caller against a /// concurrent credential install. - std::shared_ptr FileIOForPath(std::string_view location); - - using DelegatesByPrefix = - std::vector>>; + Result> FileIOForPath(std::string_view location); - /// \brief Longest-prefix match against one consistent view of the delegates. - static std::shared_ptr MatchDelegate( + /// \brief Longest-prefix match against one consistent view of the delegates, + /// or `fallback` if none matches. + /// + /// Fails once the matching delegate's keys have expired: a handle built on + /// them would only fail at S3. + static Result> MatchDelegate( const std::shared_ptr& fallback, const DelegatesByPrefix& by_prefix, std::string_view location); @@ -373,15 +605,81 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { void InstallCredentials(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, and by the S3 clients before each + /// request, so open handles see the new keys too. A failure is logged and + /// keeps the current credentials, which MatchDelegate refuses once expired. + /// + /// \param from_request Whether an S3 client calls while signing a request. + /// Arrow then holds its client registry, which building a client + /// locks exclusively, so a refresh needing new clients is left to the + /// next file operation. Nor does it wait for another refresh, which + /// may be building clients. + void MaybeRefreshCredentials(bool from_request = false); + + /// \brief The current delegates, with `refreshed`'s expiries, if keys are all + /// it changes: their S3 clients stay and read the new keys through their + /// providers. Nullopt if it needs new clients. + static std::optional WithRefreshedKeys( + const std::vector& current, const DelegatesByPrefix& delegates, + const std::vector& refreshed); + + /// \brief The keys installed for the canonical `prefix`, if any. + std::optional KeysFor(std::string_view prefix) const; + std::shared_ptr default_file_io_; std::unordered_map default_properties_; const size_t delete_parallelism_; + const std::shared_ptr key_source_ = std::make_shared(this); // Guards everything below; shared because reads happen per file operation. mutable std::shared_mutex mutex_; std::vector storage_credentials_; DelegatesByPrefix file_io_by_prefix_; + std::shared_ptr provider_; + CredentialsExpiry expiry_; + std::chrono::steady_clock::time_point retry_refresh_at_; + // Until when requests leave the refresh to file operations, after finding + // that it needs new clients. + std::chrono::steady_clock::time_point request_refresh_at_; + // Bumped on each install so a refresh cannot overwrite newer credentials. + 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::InitializeStorageCredentials( + const std::vector& storage_credentials, + std::shared_ptr provider) { + ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials)); + auto credentials = storage_credentials; + { + std::unique_lock lock(mutex_); + provider_.swap(provider); + InstallCredentials(credentials, delegates); + } + return {}; +} + Status ArrowS3FileIO::SetStorageCredentials( const std::vector& storage_credentials) { ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials)); @@ -399,7 +697,6 @@ Result ArrowS3FileIO::BuildDelegates( const std::vector& storage_credentials) const { DelegatesByPrefix delegates; delegates.reserve(storage_credentials.size()); - // TODO(gangwu): Refresh vended credentials via credentials.uri before tokens expire. for (const auto& credential : storage_credentials) { ICEBERG_RETURN_UNEXPECTED(credential.Validate()); // A server may vend credentials for several storage systems at once; @@ -412,9 +709,20 @@ Result ArrowS3FileIO::BuildDelegates( for (const auto& [key, value] : credential.config) { properties[key] = value; } - ICEBERG_ASSIGN_OR_RAISE(auto fs, BuildArrowS3FileSystem(properties)); - delegates.emplace_back(CanonicalizeS3Scheme(credential.prefix), - std::make_shared(std::move(fs))); + auto prefix = CanonicalizeS3Scheme(credential.prefix); + std::shared_ptr keys; + if (HasKeys(credential.config)) { + keys = std::make_shared( + key_source_, prefix, AwsCredentialsOf(credential.config)); + } + const bool vended_keys = keys != nullptr; + ICEBERG_ASSIGN_OR_RAISE(auto fs, BuildArrowS3FileSystem(properties, std::move(keys))); + delegates.push_back( + {.prefix = std::move(prefix), + .file_io = std::make_shared(std::move(fs)), + .vended_keys = vended_keys, + // Unusable expiries are reported by ExpiryOf when installed. + .expires_at = ExpiryOf(credential).value_or(std::nullopt)}); } if (delegates.empty() && !storage_credentials.empty()) { // Silent skipping of every vended credential is hard to diagnose: S3 access @@ -430,53 +738,250 @@ Result ArrowS3FileIO::BuildDelegates( void ArrowS3FileIO::InstallCredentials( std::vector& storage_credentials, DelegatesByPrefix& delegates) { file_io_by_prefix_.swap(delegates); + expiry_ = ExpiryOf(storage_credentials); + retry_refresh_at_ = {}; + request_refresh_at_ = {}; + ++credential_generation_; storage_credentials_.swap(storage_credentials); } -std::shared_ptr ArrowS3FileIO::MatchDelegate( +bool ArrowS3FileIO::RefreshDue() const { + return (expiry_.unknown || + (expiry_.expires_at.has_value() && + std::chrono::system_clock::now() + kRefreshLeadTime >= *expiry_.expires_at)) && + std::chrono::steady_clock::now() >= retry_refresh_at_; +} + +bool ArrowS3FileIO::Expired() const { + return expiry_.expires_at.has_value() && + std::chrono::system_clock::now() >= *expiry_.expires_at; +} + +std::chrono::steady_clock::time_point ArrowS3FileIO::BackoffUntil() const { + auto delay = + std::chrono::duration_cast(kRefreshRetryBackoff); + if (expiry_.expires_at.has_value()) { + const auto remaining = std::chrono::duration_cast( + *expiry_.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(bool from_request) { + const auto due = [&] { + return provider_ && RefreshDue() && + (!from_request || std::chrono::steady_clock::now() >= request_refresh_at_); + }; + { + // Cheap pre-check, so the common case costs one shared lock and no more. + std::shared_lock lock(mutex_); + if (!due()) { + 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. + if (from_request) { + return; + } + { + 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; + } + } + // Snapshot the provider, credentials and their generation together. + std::shared_ptr provider; + uint64_t generation = 0; + std::vector current; + DelegatesByPrefix current_delegates; + { + // Whoever held the lock may also have just finished, leaving nothing to do. + std::shared_lock lock(mutex_); + if (!due()) { + return; + } + provider = provider_; + generation = credential_generation_; + current = storage_credentials_; + current_delegates = file_io_by_prefix_; + } + + // Outside `mutex_`: both are slow and must not block readers. + Status status; + DelegatesByPrefix delegates; + bool needs_clients = false; + auto refreshed = provider->Load(); + if (refreshed.has_value()) { + status = ValidateRefreshed(*refreshed); + } else { + status = ::iceberg::unexpected(refreshed.error()); + } + if (status.has_value()) { + if (auto same_clients = WithRefreshedKeys(current, current_delegates, *refreshed); + same_clients.has_value()) { + delegates = std::move(same_clients).value(); + } else if (from_request) { + needs_clients = true; + } else if (auto built = BuildDelegates(*refreshed); built.has_value()) { + delegates = std::move(built).value(); + } else { + status = ::iceberg::unexpected(built.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; + } + if (needs_clients) { + // Fetched again by the next file operation, so it is not stale by then. + request_refresh_at_ = std::chrono::steady_clock::now() + kRefreshRetryBackoff; + ICEBERG_LOG_DEBUG( + "Refreshed vended storage credentials need new S3 clients; leaving them to " + "the next file operation"); + return; + } + + // `delegates`/`*refreshed` take the retired generation; both are declared + // before the lock, so it destructs only after the lock releases. + InstallCredentials(*refreshed, delegates); + if (RefreshDue()) { + // Tokens shorter-lived than the lead time come back due again at once. + retry_refresh_at_ = BackoffUntil(); + } +} + +std::optional ArrowS3FileIO::WithRefreshedKeys( + const std::vector& current, const DelegatesByPrefix& delegates, + const std::vector& refreshed) { + // BuildDelegates made one delegate per S3 credential, in order. + const auto current_s3 = S3CredentialsOf(current); + const auto refreshed_s3 = S3CredentialsOf(refreshed); + if (current_s3.size() != delegates.size() || refreshed_s3.size() != delegates.size()) { + return std::nullopt; + } + auto result = delegates; + for (size_t i = 0; i < result.size(); ++i) { + if (!result[i].vended_keys || !DiffersOnlyInKeys(*current_s3[i], *refreshed_s3[i])) { + return std::nullopt; + } + // ValidateRefreshed has checked the expiry. + result[i].expires_at = ExpiryOf(*refreshed_s3[i]).value_or(std::nullopt); + } + return result; +} + +std::optional ArrowS3FileIO::KeysFor( + std::string_view prefix) const { + std::shared_lock lock(mutex_); + for (const auto* credential : S3CredentialsOf(storage_credentials_)) { + if (HasKeys(credential->config) && + CanonicalizeS3Scheme(credential->prefix) == prefix) { + return AwsCredentialsOf(credential->config); + } + } + return std::nullopt; +} + +Result> ArrowS3FileIO::MatchDelegate( const std::shared_ptr& fallback, const DelegatesByPrefix& by_prefix, std::string_view location) { if (by_prefix.empty()) { return fallback; } const std::string canonical = CanonicalizeS3Scheme(location); - auto best = fallback; + const Delegate* best = nullptr; size_t best_len = 0; - for (const auto& [prefix, file_io] : by_prefix) { - if (prefix.size() > best_len && canonical.starts_with(prefix)) { - best = file_io; - best_len = prefix.size(); + for (const auto& delegate : by_prefix) { + if (delegate.prefix.size() > best_len && canonical.starts_with(delegate.prefix)) { + best = &delegate; + best_len = delegate.prefix.size(); } } - return best; + if (best == nullptr) { + return fallback; + } + if (best->expires_at.has_value() && + std::chrono::system_clock::now() >= *best->expires_at) { + return TokenExpired( + "Vended credentials for \"{}\" have expired and could not be refreshed", + best->prefix); + } + return best->file_io; } -std::shared_ptr ArrowS3FileIO::FileIOForPath( +Result> ArrowS3FileIO::FileIOForPath( std::string_view location) { + MaybeRefreshCredentials(); + std::shared_lock lock(mutex_); return MatchDelegate(default_file_io_, file_io_by_prefix_, location); } Result> ArrowS3FileIO::NewInputFile( std::string file_location) { - return FileIOForPath(file_location)->NewInputFile(std::move(file_location)); + ICEBERG_ASSIGN_OR_RAISE(auto file_io, FileIOForPath(file_location)); + return file_io->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); + ICEBERG_ASSIGN_OR_RAISE(auto file_io, FileIOForPath(file_location)); + return file_io->NewInputFile(std::move(file_location), length); } Result> ArrowS3FileIO::NewOutputFile( std::string file_location) { - return FileIOForPath(file_location)->NewOutputFile(std::move(file_location)); + ICEBERG_ASSIGN_OR_RAISE(auto file_io, FileIOForPath(file_location)); + return file_io->NewOutputFile(std::move(file_location)); } Status ArrowS3FileIO::DeleteFile(const std::string& file_location) { - return FileIOForPath(file_location)->DeleteFile(file_location); + ICEBERG_ASSIGN_OR_RAISE(auto file_io, FileIOForPath(file_location)); + return file_io->DeleteFile(file_location); } Status ArrowS3FileIO::DeleteFiles(const std::vector& file_locations) { + MaybeRefreshCredentials(); + // Keep one delegate snapshot for the whole batch. std::shared_ptr fallback; DelegatesByPrefix by_prefix; @@ -485,11 +990,22 @@ Status ArrowS3FileIO::DeleteFiles(const std::vector& file_locations fallback = default_file_io_; by_prefix = file_io_by_prefix_; } - return BulkDeleteFiles( - file_locations, delete_executor.Get(), delete_parallelism_, - [&](const std::string& location) { - return MatchDelegate(fallback, by_prefix, location)->DeleteFile(location); - }); + return BulkDeleteFiles(file_locations, delete_executor.Get(), delete_parallelism_, + [&](const std::string& location) -> Status { + ICEBERG_ASSIGN_OR_RAISE( + auto file_io, + MatchDelegate(fallback, by_prefix, location)); + return file_io->DeleteFile(location); + }); +} + +std::optional KeySource::KeysFor(std::string_view prefix) { + std::shared_lock lock(mutex_); + if (io_ == nullptr) { + return std::nullopt; + } + io_->MaybeRefreshCredentials(/*from_request=*/true); + return io_->KeysFor(prefix); } } // namespace diff --git a/src/iceberg/arrow/s3/s3_properties.h b/src/iceberg/arrow/s3/s3_properties.h index fe0a9aa3c9..14cd2b4fdc 100644 --- a/src/iceberg/arrow/s3/s3_properties.h +++ b/src/iceberg/arrow/s3/s3_properties.h @@ -42,6 +42,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 S3-compatible object stores) diff --git a/src/iceberg/catalog/rest/http_client.cc b/src/iceberg/catalog/rest/http_client.cc index d48471cde3..1bfd27109f 100644 --- a/src/iceberg/catalog/rest/http_client.cc +++ b/src/iceberg/catalog/rest/http_client.cc @@ -152,6 +152,10 @@ Status HandleFailureResponse(const cpr::Response& response, if (IsSuccessful(response.status_code)) { return {}; } + // No response to hand the error handler. + if (response.error.code == cpr::ErrorCode::OPERATION_TIMEDOUT) { + return IOError("HTTP request timed out: {}", response.error.message); + } auto parse_result = error_handler.ParseResponse( static_cast(response.status_code), response.text); const ErrorResponse final_error = @@ -176,13 +180,16 @@ HttpClient::~HttpClient() = default; Result HttpClient::Get( const std::string& path, const std::unordered_map& params, const std::unordered_map& headers, - const ErrorHandler& error_handler, auth::AuthSession& session) { + const ErrorHandler& error_handler, auth::AuthSession& session, + std::optional timeout) { ICEBERG_ASSIGN_OR_RAISE(auto url, AppendQueryString(path, params)); ICEBERG_ASSIGN_OR_RAISE(auto authenticated, AuthenticateRequest(session, HttpMethod::kGet, std::move(url), MergeHeaders(default_headers_, headers))); - cpr::Response response = cpr::Get(cpr::Url{authenticated.url}, - ToCprHeader(authenticated), *connection_pool_); + // Zero leaves the request unbounded. + cpr::Response response = + cpr::Get(cpr::Url{authenticated.url}, ToCprHeader(authenticated), *connection_pool_, + cpr::Timeout{timeout.value_or(std::chrono::milliseconds::zero())}); ICEBERG_RETURN_UNEXPECTED(HandleFailureResponse(response, error_handler)); HttpResponse http_response; diff --git a/src/iceberg/catalog/rest/http_client.h b/src/iceberg/catalog/rest/http_client.h index 9400202dc8..1f2d02eac2 100644 --- a/src/iceberg/catalog/rest/http_client.h +++ b/src/iceberg/catalog/rest/http_client.h @@ -19,8 +19,10 @@ #pragma once +#include #include #include +#include #include #include @@ -80,10 +82,14 @@ class ICEBERG_REST_EXPORT HttpClient { HttpClient& operator=(HttpClient&&) = delete; /// \brief Sends a GET request. + /// + /// \param timeout Limit on the whole request, after which it fails with an + /// IOError; nullopt for none. virtual Result Get( const std::string& path, const std::unordered_map& params, const std::unordered_map& headers, - const ErrorHandler& error_handler, auth::AuthSession& session); + const ErrorHandler& error_handler, auth::AuthSession& session, + std::optional timeout); /// \brief Sends a POST request. virtual Result Post( diff --git a/src/iceberg/catalog/rest/rest_catalog.cc b/src/iceberg/catalog/rest/rest_catalog.cc index c450673105..58ee2db11d 100644 --- a/src/iceberg/catalog/rest/rest_catalog.cc +++ b/src/iceberg/catalog/rest/rest_catalog.cc @@ -19,6 +19,7 @@ #include "iceberg/catalog/rest/rest_catalog.h" +#include #include #include #include @@ -69,6 +70,10 @@ namespace iceberg::rest { namespace { +// Bounds a credential refresh, which runs on the path of a file operation. The +// FileIO keeps its current credentials, or fails once they have expired. +constexpr auto kLoadCredentialsTimeout = std::chrono::seconds(10); + class RestStorageCredentialProvider final : public StorageCredentialProvider { public: RestStorageCredentialProvider(std::shared_ptr client, @@ -83,9 +88,10 @@ class RestStorageCredentialProvider final : public StorageCredentialProvider { Result> Load() override { // The provider can be shared by several FileIOs; keep fetches serialized. std::lock_guard lock(mutex_); - ICEBERG_ASSIGN_OR_RAISE(const auto response, - client_->Get(path_, params_, /*headers=*/{}, - *TableErrorHandler::Instance(), *session_)); + ICEBERG_ASSIGN_OR_RAISE( + const auto response, + client_->Get(path_, params_, /*headers=*/{}, *TableErrorHandler::Instance(), + *session_, kLoadCredentialsTimeout)); // Parse errors can contain credential data; return a fixed message. auto json = FromJsonString(response.body()); if (!json.has_value()) { @@ -147,7 +153,8 @@ Result FetchServerConfig(HttpClient& client, const ResourcePaths& ICEBERG_ASSIGN_OR_RAISE(const auto response, client.Get(config_path, params, /*headers=*/{}, - *ConfigErrorHandler::Instance(), session)); + *ConfigErrorHandler::Instance(), session, + /*timeout=*/std::nullopt)); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); return CatalogConfigFromJson(json); } @@ -650,7 +657,8 @@ Result> RestCatalog::ListNamespaces( } ICEBERG_ASSIGN_OR_RAISE(const auto response, client_->Get(path, params, /*headers=*/{}, - *NamespaceErrorHandler::Instance(), session)); + *NamespaceErrorHandler::Instance(), session, + /*timeout=*/std::nullopt)); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); ICEBERG_ASSIGN_OR_RAISE(auto list_response, ListNamespacesResponseFromJson(json)); result.insert(result.end(), list_response.namespaces.begin(), @@ -684,7 +692,8 @@ Result> RestCatalog::GetNamespacePr ICEBERG_ASSIGN_OR_RAISE(auto path, paths_->Namespace_(ns)); ICEBERG_ASSIGN_OR_RAISE(const auto response, client_->Get(path, /*params=*/{}, /*headers=*/{}, - *NamespaceErrorHandler::Instance(), session)); + *NamespaceErrorHandler::Instance(), session, + /*timeout=*/std::nullopt)); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); ICEBERG_ASSIGN_OR_RAISE(auto get_response, GetNamespaceResponseFromJson(json)); return get_response.properties; @@ -743,7 +752,8 @@ Result> RestCatalog::ListTables( } ICEBERG_ASSIGN_OR_RAISE(const auto response, client_->Get(path, params, /*headers=*/{}, - *TableErrorHandler::Instance(), session)); + *TableErrorHandler::Instance(), session, + /*timeout=*/std::nullopt)); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); ICEBERG_ASSIGN_OR_RAISE(auto list_response, ListTablesResponseFromJson(json)); result.insert(result.end(), list_response.identifiers.begin(), @@ -938,9 +948,10 @@ Result RestCatalog::LoadTableInternal(const TableIdentifier& id params["snapshots"] = "all"; } - ICEBERG_ASSIGN_OR_RAISE(const auto response, - client_->Get(path, params, /*headers=*/{}, - *TableErrorHandler::Instance(), session)); + ICEBERG_ASSIGN_OR_RAISE( + const auto response, + client_->Get(path, params, /*headers=*/{}, *TableErrorHandler::Instance(), session, + /*timeout=*/std::nullopt)); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); return LoadTableResultFromJson(json); } diff --git a/src/iceberg/catalog/rest/rest_table_scan.cc b/src/iceberg/catalog/rest/rest_table_scan.cc index e7dae0f492..6317b87efe 100644 --- a/src/iceberg/catalog/rest/rest_table_scan.cc +++ b/src/iceberg/catalog/rest/rest_table_scan.cc @@ -205,7 +205,8 @@ Result FetchPlanningResult( for (int32_t retry = 0; retry <= poll_config.max_retries; ++retry) { ICEBERG_ASSIGN_OR_RAISE(auto response, ctx.client->Get(path, /*params=*/{}, /*headers=*/{}, - *PlanErrorHandler::Instance(), *ctx.session)); + *PlanErrorHandler::Instance(), *ctx.session, + /*timeout=*/std::nullopt)); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); ICEBERG_ASSIGN_OR_RAISE(auto result, FetchPlanningResultResponseFromJson(json, specs, schema)); diff --git a/src/iceberg/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt index 292132bb6d..48191ef37d 100644 --- a/src/iceberg/test/CMakeLists.txt +++ b/src/iceberg/test/CMakeLists.txt @@ -328,6 +328,7 @@ if(ICEBERG_BUILD_REST) catalog_properties_test.cc error_handlers_test.cc endpoint_test.cc + http_client_test.cc rest_catalog_unit_test.cc rest_file_io_test.cc rest_json_serde_test.cc diff --git a/src/iceberg/test/arrow_s3_file_io_test.cc b/src/iceberg/test/arrow_s3_file_io_test.cc index 6bb3f4178e..0499bc30a7 100644 --- a/src/iceberg/test/arrow_s3_file_io_test.cc +++ b/src/iceberg/test/arrow_s3_file_io_test.cc @@ -22,10 +22,12 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -157,10 +159,37 @@ 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; + }); +} + +class MockStorageCredentialProvider : public StorageCredentialProvider { + public: + MOCK_METHOD((Result>), Load, (), (override)); +}; + +constexpr auto kOutlastsAShortenedBackoff = std::chrono::milliseconds(1200); + +std::string ExpiresAtMs(std::chrono::milliseconds valid_for) { + return std::to_string( + std::chrono::duration_cast( + (std::chrono::system_clock::now() + valid_for).time_since_epoch()) + .count()); +} + +std::vector ExpiringCredentials(std::chrono::milliseconds valid_for, + std::string_view access_key, + std::string prefix = "s3") { + return {{.prefix = std::move(prefix), + .config = {{std::string(S3Properties::kAccessKeyId), std::string(access_key)}, + {std::string(S3Properties::kSecretAccessKey), "secret"}, + {std::string(S3Properties::kSessionToken), "token"}, + {std::string(S3Properties::kSessionTokenExpiresAtMs), + ExpiresAtMs(valid_for)}}}}; } Status CheckReadWrite(FileIO& io, const std::string& object_uri, @@ -510,6 +539,449 @@ 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; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> Result> { + ++refresh_calls; + return refreshed; + }); + ASSERT_THAT(credentialed->InitializeStorageCredentials( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), provider), + IsOk()); + EXPECT_EQ(refresh_calls, 0); + + 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; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> Result> { + ++refresh_calls; + return std::vector{}; + }); + + ASSERT_THAT(credentialed->InitializeStorageCredentials( + ExpiringCredentials(std::chrono::hours(1), "access-key"), provider), + 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, KeepsProviderWhenCredentialsAreUpdated) { + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + const auto refreshed = ExpiringCredentials(std::chrono::hours(2), "refreshed-key"); + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).WillOnce(::testing::Return(refreshed)); + ASSERT_THAT(credentialed->InitializeStorageCredentials( + ExpiringCredentials(std::chrono::hours(1), "initial-key"), provider), + IsOk()); + ASSERT_THAT(credentialed->SetStorageCredentials( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key")), + IsOk()); + + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(credentialed->credentials(), refreshed); +} + +TEST_F(ArrowS3FileIOTest, FailedInitializationDoesNotInstallProvider) { + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + const auto initial = ExpiringCredentials(std::chrono::minutes(1), "initial-key"); + ASSERT_THAT(credentialed->InitializeStorageCredentials(initial, nullptr), IsOk()); + auto invalid = initial; + invalid.front().prefix.clear(); + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).Times(0); + EXPECT_THAT(credentialed->InitializeStorageCredentials(invalid, provider), + IsError(ErrorKind::kValidationFailed)); + + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(credentialed->credentials(), initial); +} + +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; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> Result> { + ++refresh_calls; + return ExpiringCredentials(std::chrono::hours(1), "refreshed-key"); + }); + ASSERT_THAT(credentialed->InitializeStorageCredentials( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), provider), + 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; + + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> 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->InitializeStorageCredentials( + ExpiringCredentials(-std::chrono::minutes(1), "expired-key"), provider), + 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; + + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> 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->InitializeStorageCredentials( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), provider), + 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; }); + } + auto install_status = credentialed->SetStorageCredentials(installed); + { + std::lock_guard lock(mutex); + release_refresh = true; + } + cv.notify_all(); + operation.join(); + + ASSERT_THAT(install_status, IsOk()); + 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; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> 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->InitializeStorageCredentials( + {{.prefix = "s3", .config = std::move(config)}}, provider), + 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; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> Result> { + ++refresh_calls; + return undatable; + }); + ASSERT_THAT(credentialed->InitializeStorageCredentials(undatable, provider), 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; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> 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->InitializeStorageCredentials(credentials, provider), 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; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> Result> { + ++refresh_calls; + return ExpiringCredentials(std::chrono::minutes(1), "short-lived-key"); + }); + ASSERT_THAT(credentialed->InitializeStorageCredentials( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), provider), + 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; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> Result> { + ++refresh_calls; + return NotFound("catalog unreachable"); + }); + const auto expiring = ExpiringCredentials(std::chrono::minutes(1), "expiring-key"); + ASSERT_THAT(credentialed->InitializeStorageCredentials(expiring, provider), 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, KeepsCredentialsWhenTheReplacementIsUnusable) { + struct Case { + std::string_view name; + std::optional expires_at_ms; + }; + const std::vector cases = { + {.name = "expired", .expires_at_ms = ExpiresAtMs(-std::chrono::minutes(1))}, + {.name = "no expiry", .expires_at_ms = std::nullopt}, + {.name = "unparseable expiry", .expires_at_ms = "not-a-number"}}; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.name); + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + std::unordered_map config = { + {std::string(S3Properties::kAccessKeyId), "unusable-key"}, + {std::string(S3Properties::kSecretAccessKey), "secret"}, + {std::string(S3Properties::kSessionToken), "token"}}; + if (test_case.expires_at_ms.has_value()) { + config[std::string(S3Properties::kSessionTokenExpiresAtMs)] = + *test_case.expires_at_ms; + } + const std::vector unusable = { + {.prefix = "s3", .config = std::move(config)}}; + int refresh_calls = 0; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillRepeatedly([&]() -> Result> { + ++refresh_calls; + return unusable; + }); + const auto expiring = ExpiringCredentials(std::chrono::minutes(1), "expiring-key"); + ASSERT_THAT(credentialed->InitializeStorageCredentials(expiring, provider), IsOk()); + + auto logger = std::make_shared(); + ScopedDefaultLogger scoped(logger); + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(credentialed->credentials(), expiring); + EXPECT_TRUE(HasWarning(*logger, "Failed to refresh")); + + // Rejected like any failed refresh, so it is not asked for again at once. + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(refresh_calls, 1); + } +} + +TEST_F(ArrowS3FileIOTest, FailsOnceCredentialsHaveExpiredAndCannotBeRefreshed) { + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).WillOnce([]() -> Result> { + return NotFound("catalog unreachable"); + }); + ASSERT_THAT( + credentialed->InitializeStorageCredentials( + ExpiringCredentials(-std::chrono::minutes(1), "expired-key", "s3://bucket/"), + provider), + IsOk()); + + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsError(ErrorKind::kTokenExpired)); + EXPECT_THAT(io->NewOutputFile("s3://bucket/key"), IsError(ErrorKind::kTokenExpired)); + EXPECT_THAT(io->DeleteFile("s3://bucket/key"), IsError(ErrorKind::kTokenExpired)); + EXPECT_THAT(io->DeleteFiles({"s3://bucket/key"}), + HasErrorMessage("Failed to delete 1 of 1 files")); + // Locations outside the vended prefix keep the default credentials. + EXPECT_THAT(io->NewInputFile("s3://other-bucket/key"), IsOk()); +} + TEST_F(ArrowS3FileIOTest, OperationsSurviveConcurrentCredentialInstalls) { auto result = MakeS3FileIO({}); ASSERT_THAT(result, IsOk()); @@ -642,6 +1114,80 @@ TEST_F(ArrowS3FileIOTest, LongestCredentialPrefix) { IsOk()); } +TEST_F(ArrowS3FileIOTest, RefreshReachesHandlesOpenedBeforeIt) { + if (!HasIntegrationEnv()) { + GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test"; + } + + auto properties = PropertiesFromEnv(); + const auto access_key = properties.find(std::string(S3Properties::kAccessKeyId)); + const auto secret_key = properties.find(std::string(S3Properties::kSecretAccessKey)); + if (access_key == properties.end() || secret_key == properties.end()) { + GTEST_SKIP() << "Set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY to enable " + "credential refresh test"; + } + + // A refresh that changes only keys keeps the S3 client, and the handle's own + // request runs it. One that changes a setting replaces the client, which a + // request cannot do, so a file operation runs it. Either way the handle + // opened before it, still on the first client, must sign with the new keys. + struct Case { + std::string_view name; + std::unordered_map refreshed_settings; + bool file_operation_refreshes; + }; + const std::vector cases = { + {.name = "keys only", .refreshed_settings = {}, .file_operation_refreshes = false}, + {.name = "new client", + .refreshed_settings = {{std::string(S3Properties::kConnectTimeoutMs), "5000"}}, + .file_operation_refreshes = true}}; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.name); + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO(properties)); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + const auto prefix = ObjectUri("refresh_in_place/"); + const auto uri = prefix + "object"; + constexpr std::string_view kContent = "read with keys refreshed after it was opened"; + ASSERT_THAT(io->WriteFile(uri, kContent), IsOk()); + + // The vended keys are bad, and fall due for refresh just after the handle + // is opened. Only the refreshed keys can read the object. + constexpr auto kDueIn = std::chrono::seconds(2); + auto refreshed_config = test_case.refreshed_settings; + refreshed_config.insert(*access_key); + refreshed_config.insert(*secret_key); + const std::vector refreshed = { + {.prefix = prefix, .config = std::move(refreshed_config)}}; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).WillOnce(::testing::Return(refreshed)); + ASSERT_THAT( + credentialed->InitializeStorageCredentials( + {{.prefix = prefix, + .config = {{std::string(S3Properties::kAccessKeyId), "bad-access-key"}, + {std::string(S3Properties::kSecretAccessKey), "bad-secret"}, + {std::string(S3Properties::kSessionToken), "bad-token"}, + {std::string(S3Properties::kSessionTokenExpiresAtMs), + ExpiresAtMs(std::chrono::minutes(5) + kDueIn)}}}}, + provider), + IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto file, io->NewInputFile(uri, kContent.size())); + + std::this_thread::sleep_for(kDueIn + std::chrono::seconds(1)); + if (test_case.file_operation_refreshes) { + ASSERT_THAT(io->NewInputFile(uri), IsOk()); + } + ICEBERG_UNWRAP_OR_FAIL(auto stream, file->Open()); + std::string read(kContent.size(), '\0'); + EXPECT_THAT(stream->ReadFully(0, std::as_writable_bytes(std::span(read))), IsOk()); + EXPECT_EQ(read, kContent); + EXPECT_THAT(stream->Close(), IsOk()); + EXPECT_EQ(credentialed->credentials(), refreshed); + EXPECT_THAT(io->DeleteFile(uri), IsOk()); + } +} + TEST_F(ArrowS3FileIOTest, DeleteFilesAttemptsEveryFile) { if (!HasIntegrationEnv()) { GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test"; diff --git a/src/iceberg/test/auth_manager_test.cc b/src/iceberg/test/auth_manager_test.cc index fd9605fc31..848356e30a 100644 --- a/src/iceberg/test/auth_manager_test.cc +++ b/src/iceberg/test/auth_manager_test.cc @@ -24,6 +24,7 @@ #include #include #include +#include #include #include #include @@ -253,7 +254,8 @@ TEST_F(AuthManagerTest, HttpClientRejectsParamsWhenUrlAlreadyHasQuery) { auto session = AuthSession::MakeDefault({}); auto result = client_->Get("http://127.0.0.1/v1/config?existing=true", {{"warehouse", "prod"}}, - /*headers=*/{}, *rest::DefaultErrorHandler::Instance(), *session); + /*headers=*/{}, *rest::DefaultErrorHandler::Instance(), *session, + /*timeout=*/std::nullopt); EXPECT_THAT(result, IsError(ErrorKind::kInvalidArgument)); EXPECT_THAT(result, HasErrorMessage("must not contain a query string")); diff --git a/src/iceberg/test/http_client_test.cc b/src/iceberg/test/http_client_test.cc new file mode 100644 index 0000000000..98c587247f --- /dev/null +++ b/src/iceberg/test/http_client_test.cc @@ -0,0 +1,76 @@ +/* + * 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. + */ + +#include "iceberg/catalog/rest/http_client.h" + +#include +#include +#include + +#include +#include + +#ifndef _WIN32 +# include + +# include +# include +# include +#endif + +#include "iceberg/catalog/rest/auth/auth_session.h" +#include "iceberg/catalog/rest/error_handlers.h" +#include "iceberg/test/matchers.h" + +namespace iceberg::rest { + +#ifndef _WIN32 + +TEST(HttpClientTest, GetFailsOnceItsTimeoutPasses) { + // A listening socket that never accepts: the kernel completes the connection, + // but no response ever comes. + const int server = ::socket(AF_INET, SOCK_STREAM, 0); + ASSERT_GE(server, 0); + struct Closer { + int fd; + ~Closer() { ::close(fd); } + } closer{server}; + sockaddr_in address{}; + address.sin_family = AF_INET; + address.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + ASSERT_EQ(::bind(server, reinterpret_cast(&address), sizeof(address)), 0); + ASSERT_EQ(::listen(server, 1), 0); + socklen_t length = sizeof(address); + ASSERT_EQ(::getsockname(server, reinterpret_cast(&address), &length), 0); + const auto url = std::format("http://127.0.0.1:{}/v1/config", ntohs(address.sin_port)); + + HttpClient client; + auto session = auth::AuthSession::MakeDefault({}); + const auto start = std::chrono::steady_clock::now(); + auto response = + client.Get(url, /*params=*/{}, /*headers=*/{}, *DefaultErrorHandler::Instance(), + *session, std::chrono::milliseconds(200)); + EXPECT_THAT(response, IsError(ErrorKind::kIOError)); + EXPECT_THAT(response, HasErrorMessage("timed out")); + EXPECT_LT(std::chrono::steady_clock::now() - start, std::chrono::seconds(10)); +} + +#endif + +} // namespace iceberg::rest diff --git a/src/iceberg/test/rest_arrow_file_io_test.cc b/src/iceberg/test/rest_arrow_file_io_test.cc index 83966dc96a..30066f946e 100644 --- a/src/iceberg/test/rest_arrow_file_io_test.cc +++ b/src/iceberg/test/rest_arrow_file_io_test.cc @@ -21,6 +21,7 @@ /// \brief Covers REST -> ResolvingFileIO -> registry -> Arrow FileIO against the /// real registered implementations, which mock delegates cannot exercise. +#include #include #include #include @@ -67,6 +68,11 @@ TEST_F(RestArrowFileIOTest, ReadsBackWhatItWroteThroughRealLocalFileIO) { #if ICEBERG_S3_ENABLED +class MockStorageCredentialProvider : public StorageCredentialProvider { + public: + MOCK_METHOD((Result>), Load, (), (override)); +}; + std::optional GetEnvIfSet(const char* key) { const char* value = std::getenv(key); if (value == nullptr || std::string_view(value).empty()) { @@ -195,6 +201,58 @@ TEST_F(RestArrowFileIOTest, ReadsBackWhatItWroteThroughAnOssLocation) { EXPECT_THAT(io.value()->DeleteFile(object_uri), IsOk()); } +TEST_F(RestArrowFileIOTest, RefreshesCredentialsThroughTheRealS3FileIO) { + const auto base_uri = GetEnvIfSet("ICEBERG_TEST_S3_URI"); + const auto access_key = GetEnvIfSet("AWS_ACCESS_KEY_ID"); + const auto secret_key = GetEnvIfSet("AWS_SECRET_ACCESS_KEY"); + if (!base_uri || !access_key || !secret_key) { + GTEST_SKIP() + << "Set ICEBERG_TEST_S3_URI, AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY"; + } + + std::unordered_map config = { + {"s3.access-key-id", *access_key}, {"s3.secret-access-key", *secret_key}}; + if (const auto token = GetEnvIfSet("AWS_SESSION_TOKEN")) { + config["s3.session-token"] = *token; + // A refreshed session token must say when it expires. + config["s3.session-token-expires-at-ms"] = std::to_string( + std::chrono::duration_cast( + (std::chrono::system_clock::now() + std::chrono::hours(1)).time_since_epoch()) + .count()); + } + if (const auto endpoint = GetEnvIfSet("ICEBERG_TEST_S3_ENDPOINT")) { + config["s3.endpoint"] = *endpoint; + } + if (const auto region = GetEnvIfSet("AWS_REGION")) { + config["client.region"] = *region; + } + const std::vector refreshed = {{.prefix = "s3", .config = config}}; + config["s3.access-key-id"] = "bad-access-key"; + config["s3.secret-access-key"] = "bad-secret-key"; + config["s3.session-token"] = "expired-token"; + config["s3.session-token-expires-at-ms"] = "0"; + + // Only the provider can supply credentials that authenticate this round trip. + ScopedScrubbedAwsCredentialEnv scrubbed; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).WillOnce(::testing::Return(refreshed)); + auto io = MakeTableFileIO({{"warehouse", "logical_warehouse_name"}}, + /*table_config=*/{}, + {{.prefix = "s3", .config = std::move(config)}}, provider); + ASSERT_THAT(io, IsOk()); + + auto object_uri = *base_uri; + if (!object_uri.ends_with('/')) { + object_uri += '/'; + } + object_uri += "iceberg_refresh_provider.txt"; + constexpr std::string_view kContent = "written after refreshing vended credentials"; + ASSERT_THAT(io.value()->WriteFile(object_uri, kContent), IsOk()); + EXPECT_THAT(io.value()->ReadFile(object_uri, std::nullopt), + HasValue(::testing::Eq(std::string(kContent)))); + EXPECT_THAT(io.value()->DeleteFile(object_uri), IsOk()); +} + #endif // ICEBERG_S3_ENABLED } // namespace diff --git a/src/iceberg/test/rest_catalog_integration_test.cc b/src/iceberg/test/rest_catalog_integration_test.cc index c03c9c2716..1fb174a314 100644 --- a/src/iceberg/test/rest_catalog_integration_test.cc +++ b/src/iceberg/test/rest_catalog_integration_test.cc @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -404,7 +405,8 @@ TEST_F(RestCatalogIntegrationTest, FetchServerConfigDirect) { ICEBERG_UNWRAP_OR_FAIL(const auto response, client.Get(config_url, {}, /*headers=*/{}, - *DefaultErrorHandler::Instance(), *noop_session)); + *DefaultErrorHandler::Instance(), *noop_session, + /*timeout=*/std::nullopt)); ICEBERG_UNWRAP_OR_FAIL(auto json, FromJsonString(response.body())); EXPECT_TRUE(json.contains("defaults")); diff --git a/src/iceberg/test/rest_catalog_unit_test.cc b/src/iceberg/test/rest_catalog_unit_test.cc index 97c9c36504..67a2f04988 100644 --- a/src/iceberg/test/rest_catalog_unit_test.cc +++ b/src/iceberg/test/rest_catalog_unit_test.cc @@ -17,7 +17,9 @@ * under the License. */ +#include #include +#include #include #include #include @@ -66,7 +68,8 @@ class MockHttpClient : public HttpClient { (const std::string& path, (const std::unordered_map&)params, (const std::unordered_map&)headers, - const ErrorHandler& error_handler, auth::AuthSession& session), + const ErrorHandler& error_handler, auth::AuthSession& session, + std::optional timeout), (override)); MOCK_METHOD(Result, Post, @@ -203,7 +206,7 @@ class RestCatalogLoadTableTest : public ::testing::Test { props["scan-planning-mode"] = client_scan_mode; } CatalogConfig server_config{.endpoints = {endpoints.begin(), endpoints.end()}}; - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/config"), _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/config"), _, _, _, _, _)) .WillOnce(Return(HttpResponse(200, ToJson(server_config).dump()))); auto config = RestCatalogProperties::FromMap(std::move(props)); return RestCatalog::Make(config, /*metrics_executor=*/nullptr, @@ -247,7 +250,7 @@ class RestCatalogLabelsTest : public RestCatalogLoadTableTest, }; TEST_P(RestCatalogLabelsTest, LoadExposesLabels) { - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _, _)) .WillOnce(Return(HttpResponse(200, LabelsResponse()))); ICEBERG_UNWRAP_OR_FAIL(auto catalog, MakeLabelsCatalog()); ICEBERG_UNWRAP_OR_FAIL(auto table, catalog->LoadTable(identifier_)); @@ -292,9 +295,9 @@ TEST_F(RestCatalogLoadTableTest, InitializationPropagatesClientFactoryErrors) { for (int fail_on_call : {1, 2}) { SCOPED_TRACE(fail_on_call); if (fail_on_call == 1) { - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)).Times(0); } else { - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/config"), _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/config"), _, _, _, _, _)) .WillOnce(Return(HttpResponse(200, ToJson(CatalogConfig{}).dump()))); } int calls = 0; @@ -321,14 +324,14 @@ TEST_F(RestCatalogLoadTableTest, InitializationUsesMergedServerConfig) { {"scan-planning-mode", "server"}, {"header.X-Shared", "server"}}, .endpoints = {all_plan_endpoints_.begin(), all_plan_endpoints_.end()}}; - EXPECT_CALL( - *mock_client_, - Get(::testing::EndsWith("/v1/config"), - ::testing::Contains(::testing::Pair("warehouse", "client-warehouse")), _, _, _)) + EXPECT_CALL(*mock_client_, + Get(::testing::EndsWith("/v1/config"), + ::testing::Contains(::testing::Pair("warehouse", "client-warehouse")), + _, _, _, _)) .WillOnce(Return(HttpResponse(200, ToJson(server_config).dump()))); EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/catalog/namespaces/default/tables/my_table"), - ::testing::Contains(::testing::Pair("snapshots", "refs")), _, _, _)) + ::testing::Contains(::testing::Pair("snapshots", "refs")), _, _, _, _)) .WillOnce(Return(HttpResponse(200, kLoadTableDefaultResponse))); EXPECT_CALL( *mock_client_, @@ -368,7 +371,7 @@ TEST_F(RestCatalogLoadTableTest, InitializationUsesMergedServerConfig) { // LoadTable returns NotSupported. // -------------------------------------------------------------------------- TEST_F(RestCatalogLoadTableTest, ServerScanWithMissingEndpointReturnsNotSupported) { - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)) .WillOnce(Return(HttpResponse(200, kLoadTableServerScanResponse))); ICEBERG_UNWRAP_OR_FAIL(auto catalog, MakeCatalog(no_plan_endpoint_set_)); @@ -384,14 +387,14 @@ TEST_F(RestCatalogLoadTableTest, ServerScanWithMissingEndpointReturnsNotSupporte TEST_F(RestCatalogLoadTableTest, ScanCredentialProviderWithoutPlanIdOmitsPlanId) { RegisterCredentialedFileIO(); EXPECT_CALL(*mock_client_, - Get(::testing::Not(::testing::HasSubstr("/credentials")), _, _, _, _)) + Get(::testing::Not(::testing::HasSubstr("/credentials")), _, _, _, _, _)) .WillOnce(Return(HttpResponse(200, kLoadTableServerScanResponse))); EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) .WillOnce(Return(HttpResponse( 200, R"({"status":"completed","storage-credentials":[{"prefix":"s3://bucket","config":{"k":"v"}}]})"))); - EXPECT_CALL(*mock_client_, - Get(::testing::HasSubstr("/credentials"), ::testing::IsEmpty(), _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::HasSubstr("/credentials"), + ::testing::IsEmpty(), _, _, _, ::testing::Optional(_))) .WillOnce(Return(HttpResponse( 200, R"({"storage-credentials":[{"prefix":"s3://bucket","config":{"k":"v2"}}]})"))); @@ -418,7 +421,7 @@ TEST_F(RestCatalogLoadTableTest, ScanCredentialProviderWithoutPlanIdOmitsPlanId) TEST_F(RestCatalogLoadTableTest, ScanCredentialsWithoutCredentialsEndpointHaveNoProvider) { RegisterCredentialedFileIO(); - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)) .WillOnce(Return(HttpResponse(200, kLoadTableServerScanResponse))); EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) .WillOnce(Return(HttpResponse(200, std::string(kPlanWithCredentialsResponse)))); @@ -453,7 +456,7 @@ TEST_P(RestCatalogPlanningModeTest, LoadAndCommitKeepTheEffectivePlanningMode) { const auto& config = GetParam(); auto response = nlohmann::json::parse(kLoadTableDefaultResponse); if (!config.table.empty()) response["config"]["scan-planning-mode"] = config.table; - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _, _)) .WillOnce(Return(HttpResponse(200, response.dump()))); EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/tables/my_table"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, kCommitTableResponse))); @@ -506,13 +509,13 @@ TEST_P(RestCatalogScanLifetimeTest, ScanKeepsCredentialRefreshAlive) { {"timestamp-ms", 2}, {"manifest-list", "s3://bucket/manifests-2.avro"}, {"schema-id", 1}}}); - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _, _)) .WillOnce(Return(HttpResponse(200, response.dump()))); EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, std::string(kPlanWithCredentialsResponse)))); - EXPECT_CALL(*mock_client_, - Get(::testing::EndsWith("/credentials"), - ::testing::Contains(::testing::Pair("planId", "plan-1")), _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/credentials"), + ::testing::Contains(::testing::Pair("planId", "plan-1")), + _, _, _, ::testing::Optional(_))) .WillOnce(Return(HttpResponse( 200, R"({"storage-credentials":[{"prefix":"s3://bucket","config":{"k":"v2"}}]})"))); diff --git a/src/iceberg/test/rest_table_scan_test.cc b/src/iceberg/test/rest_table_scan_test.cc index a1b8e6e867..b9f9112f6d 100644 --- a/src/iceberg/test/rest_table_scan_test.cc +++ b/src/iceberg/test/rest_table_scan_test.cc @@ -20,6 +20,7 @@ #include "iceberg/catalog/rest/rest_table_scan.h" #include +#include #include #include #include @@ -114,7 +115,8 @@ class MockHttpClient : public HttpClient { (const std::string& path, (const std::unordered_map&)params, (const std::unordered_map&)headers, - const ErrorHandler& error_handler, auth::AuthSession& session), + const ErrorHandler& error_handler, auth::AuthSession& session, + std::optional timeout), (override)); MOCK_METHOD(Result, Post, @@ -604,7 +606,7 @@ TEST_F(RestTableScanTest, PlanFilesStreamPollNumRetriesZero) { R"({"status":"submitted","plan-id":"plan-no-retry"})"; EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) .WillOnce(Return(HttpResponse(200, std::string(kSubmittedBody)))); - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)) .WillOnce(Return(HttpResponse(200, std::string(kSubmittedBody)))); EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)) .WillOnce(Return(HttpResponse(200, "{}"))); @@ -649,7 +651,7 @@ TEST_F(RestTableScanTest, CancelCalledOnFetchPlanningResultGetError) { EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) .WillOnce(Return(HttpResponse(200, std::string(kSubmittedBody)))); - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)) .WillOnce(Return(IOError("network failure"))); EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)) .WillOnce(Return(IOError("cancel failure"))); @@ -968,7 +970,7 @@ TEST_P(RestScanResidualTest, PreservesFilteringAndDeleteFiles) { if (submitted) { completed.erase("plan-id"); EXPECT_CALL(*mock_client_, - Get(::testing::EndsWith("/plan/residual-plan"), _, _, _, _)) + Get(::testing::EndsWith("/plan/residual-plan"), _, _, _, _, _)) .WillOnce(Return(HttpResponse(200, completed.dump()))); } if (fetch_token) { @@ -1084,7 +1086,7 @@ TEST_P(RestScanRegressionTest, MissingPollingEndpointCancelsSubmittedPlan) { EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) .WillOnce(Return( HttpResponse(200, R"({"status":"submitted","plan-id":"missing-poll"})"))); - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)).Times(0); EXPECT_CALL(*mock_client_, Delete(::testing::EndsWith("/plan/missing-poll"), _, _, _, _)) .WillOnce(Return(HttpResponse(204, ""))); @@ -1164,10 +1166,10 @@ class RestScanStreamTest : public RestTableScanTest, .WillOnce(Return(HttpResponse(200, initial.dump()))); if (GetParam()) { completed.erase("plan-id"); - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/plan/stream"), _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/plan/stream"), _, _, _, _, _)) .WillOnce(Return(HttpResponse(200, completed.dump()))); } else { - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)).Times(0); } } }; @@ -1237,10 +1239,10 @@ TEST_P(RestScanRegressionTest, EmptyCompletedPlan) { EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, initial))); if (submitted) { - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/plan/empty"), _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/plan/empty"), _, _, _, _, _)) .WillOnce(Return(HttpResponse(200, R"({"status":"completed"})"))); } else { - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)).Times(0); } EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); ICEBERG_UNWRAP_OR_FAIL(auto tasks, PlanScan(GetParam(), MakeContext(), [](auto&) {})); @@ -1254,7 +1256,7 @@ TEST_P(RestScanRegressionTest, FailedPlanReturnsError) { .WillOnce(Return(HttpResponse( 200, R"({"status":"failed","error":{"message":"server error","type":"ServerError","code":500}})"))); - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)).Times(0); EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); EXPECT_THAT( PlanScan(GetParam(), MakeContext(), [](auto&) {}), @@ -1303,7 +1305,7 @@ TEST_P(RestScanRegressionTest, InvalidPlanResponseCancelsKnownId) { SCOPED_TRACE(test_case.body); EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, test_case.body))); - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)).Times(0); if (test_case.cancel) { EXPECT_CALL(*mock_client_, Delete(::testing::EndsWith("/plan/invalid"), _, _, _, _)) .WillOnce(Return(HttpResponse(204, ""))); From 439397baef3d427667d119ff528ce460ba72e191 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Wed, 14 Oct 2026 03:51:08 -0700 Subject: [PATCH 2/7] fix(io): make credential refresh safe for in-flight S3 requests A refresh may now only rotate keys: refreshed credentials must keep the S3 prefixes, matched by canonical prefix in any order, their settings and which of them sign with keys. It therefore never builds S3 clients, and a request whose keys have expired waits up to 10 seconds for a refresh in flight, as file operations do, instead of signing with the expired keys. A response that changes that layout is rejected like any unusable refresh, which drops the deferred client rebuild. Only keys vended with the settings a client was built with reach it, so keys meant for another endpoint are never sent to an old one. Keys, settings and token expiries come from the vended config merged over the FileIO's properties, and only a token that signs requests is tracked. For a repeated prefix, the first credential is used throughout. Expiries are tracked per prefix from the installed credentials, and a batch delete resolves each file as a single delete does, so a long batch follows refreshes. HttpClient::Get checks the transfer before the status: a timeout, or a transport error after a successful status, is an IOError. Writes keep their status, since one may have taken effect: failing a commit outright would clean up its files. An installed Arrow does not export the AWS SDK it bundles, so with one the system SDK is used for the S3 FileIO's credentials provider. --- .../IcebergThirdpartyToolchain.cmake | 8 + src/iceberg/arrow/s3/arrow_s3_file_io.cc | 472 ++++++++---------- src/iceberg/catalog/rest/http_client.cc | 20 +- src/iceberg/catalog/rest/rest_catalog.cc | 3 +- src/iceberg/test/arrow_s3_file_io_test.cc | 454 +++++++++-------- src/iceberg/test/http_client_test.cc | 150 +++++- 6 files changed, 603 insertions(+), 504 deletions(-) diff --git a/cmake_modules/IcebergThirdpartyToolchain.cmake b/cmake_modules/IcebergThirdpartyToolchain.cmake index fef63f6513..8e9cc8b0cc 100644 --- a/cmake_modules/IcebergThirdpartyToolchain.cmake +++ b/cmake_modules/IcebergThirdpartyToolchain.cmake @@ -937,6 +937,14 @@ endif() if(ICEBERG_BUILD_BUNDLE) resolve_arrow_dependency() + # An installed Arrow does not export its bundled AWS SDK, which the S3 FileIO + # uses directly. Use the system SDK, which must match Arrow's. + if(ICEBERG_AWSSDK_BUNDLED AND NOT ARROW_VENDORED) + message(STATUS "Using the system AWS SDK, since Arrow is provided") + set(ICEBERG_AWSSDK_BUNDLED FALSE) + set(ICEBERG_AWSSDK_COMPONENTS core) + resolve_aws_sdk_dependency() + endif() resolve_avro_dependency() resolve_zstd_dependency() endif() diff --git a/src/iceberg/arrow/s3/arrow_s3_file_io.cc b/src/iceberg/arrow/s3/arrow_s3_file_io.cc index 3c46cb9e41..7653b9dd56 100644 --- a/src/iceberg/arrow/s3/arrow_s3_file_io.cc +++ b/src/iceberg/arrow/s3/arrow_s3_file_io.cc @@ -35,6 +35,7 @@ #include #include #include +#include #include #include @@ -327,35 +328,63 @@ std::string CanonicalizeS3Scheme(std::string_view 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. +// How long to wait after a failed refresh before asking the catalog again. 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: a provider may not limit how long its request takes. +// How long a caller with expired credentials waits for a refresh under way. constexpr auto kExpiredCredentialWait = std::chrono::seconds(10); -// When `credential`'s session token stops being valid: nullopt for static -// keys, which never expire, or an error if its expiry is missing or unusable. -Result> ExpiryOf( +// Whether `config` carries an access key pair. +bool HasKeys(const std::unordered_map& config) { + return FindProperty(config, S3Properties::kAccessKeyId) != nullptr && + FindProperty(config, S3Properties::kSecretAccessKey) != nullptr; +} + +// The properties an S3 client for `credential` is built from. +std::unordered_map PropertiesFor( + const std::unordered_map& default_properties, const StorageCredential& credential) { - if (FindProperty(credential.config, S3Properties::kSessionToken) == nullptr) { + auto properties = default_properties; + for (const auto& [key, value] : credential.config) { + properties[key] = value; + } + return properties; +} + +// The client settings in `properties`: everything but the keys. +std::unordered_map SettingsOf( + std::unordered_map properties) { + for (const auto key : + {S3Properties::kAccessKeyId, S3Properties::kSecretAccessKey, + S3Properties::kSessionToken, S3Properties::kSessionTokenExpiresAtMs}) { + properties.erase(std::string(key)); + } + return properties; +} + +// When the session token that `properties` sign with expires: nullopt if they +// sign without one, or an error if the expiry is missing or invalid. +Result> ExpiryOf( + const std::unordered_map& properties, + std::string_view prefix) { + // Only a nonempty token alongside an access key pair signs requests. + const auto* token = FindProperty(properties, S3Properties::kSessionToken); + if (!HasKeys(properties) || token == nullptr || token->empty()) { return std::nullopt; } - const auto* value = - FindProperty(credential.config, S3Properties::kSessionTokenExpiresAtMs); + const auto* value = FindProperty(properties, S3Properties::kSessionTokenExpiresAtMs); if (value == nullptr) { - return InvalidArgument("Credential \"{}\" has a session token but no \"{}\"", - credential.prefix, S3Properties::kSessionTokenExpiresAtMs); + return InvalidArgument("Credential \"{}\" has a session token but no \"{}\"", prefix, + S3Properties::kSessionTokenExpiresAtMs); } auto millis = StringUtils::ParseNumber(*value); if (!millis.has_value()) { return InvalidArgument( "Credential \"{}\" has a session token but an unparseable \"{}\" value \"{}\"", - credential.prefix, S3Properties::kSessionTokenExpiresAtMs, *value); + prefix, S3Properties::kSessionTokenExpiresAtMs, *value); } // Beyond what the clock can hold, converting would overflow it. constexpr auto kMaxMillis = std::chrono::duration_cast( @@ -366,28 +395,42 @@ Result> ExpiryOf( .count(); if (*millis > kMaxMillis || *millis < kMinMillis) { return InvalidArgument("Credential \"{}\" has an out-of-range \"{}\" value \"{}\"", - credential.prefix, S3Properties::kSessionTokenExpiresAtMs, - *value); + prefix, S3Properties::kSessionTokenExpiresAtMs, *value); } return std::chrono::system_clock::time_point(std::chrono::milliseconds(*millis)); } +// The S3 credentials in effect: the first for each canonical prefix, which is +// the one its locations match. +std::vector S3CredentialsOf( + const std::vector& credentials) { + std::vector s3_credentials; + std::unordered_set prefixes; + for (const auto& credential : credentials) { + if (IsS3CredentialPrefix(credential.prefix) && + prefixes.insert(CanonicalizeS3Scheme(credential.prefix)).second) { + s3_credentials.push_back(&credential); + } + } + return s3_credentials; +} + // When installed credentials need replacing. struct CredentialsExpiry { // The earliest expiry known, if any. std::optional expires_at; - // A session token has no usable expiry. It is replaced at once rather than - // used until it fails, as Java does, but not reported as expired. + // A session token lacks a usable expiry: refresh at once, as Java does, but + // do not treat it as expired. bool unknown = false; }; -CredentialsExpiry ExpiryOf(const std::vector& credentials) { +CredentialsExpiry ExpiryOf( + const std::unordered_map& default_properties, + const std::vector& credentials) { CredentialsExpiry expiry; - for (const auto& credential : credentials) { - if (!IsS3CredentialPrefix(credential.prefix)) { - continue; - } - auto expires_at = ExpiryOf(credential); + for (const auto* credential : S3CredentialsOf(credentials)) { + auto expires_at = + ExpiryOf(PropertiesFor(default_properties, *credential), credential->prefix); if (!expires_at.has_value()) { ICEBERG_LOG_WARN("{}", expires_at.error().message); expiry.unknown = true; @@ -399,32 +442,59 @@ CredentialsExpiry ExpiryOf(const std::vector& credentials) { return expiry; } -std::vector S3CredentialsOf( +// Known expiries by canonical prefix. +std::unordered_map ExpiriesOf( + const std::unordered_map& default_properties, const std::vector& credentials) { - std::vector s3_credentials; - for (const auto& credential : credentials) { - if (IsS3CredentialPrefix(credential.prefix)) { - s3_credentials.push_back(&credential); + std::unordered_map expiries; + for (const auto* credential : S3CredentialsOf(credentials)) { + // Invalid expiries are logged on install. + if (auto expires_at = + ExpiryOf(PropertiesFor(default_properties, *credential), credential->prefix); + expires_at.has_value() && expires_at->has_value()) { + expiries.emplace(CanonicalizeS3Scheme(credential->prefix), **expires_at); } } - return s3_credentials; + return expiries; } -// A refresh must not replace working credentials with ones that cannot be used -// or whose expiry is unknown, which would also defer the next refresh. -Status ValidateRefreshed(const std::vector& credentials) { - const auto now = std::chrono::system_clock::now(); +// The S3 prefixes `credentials` serve, each with its client settings and +// whether it signs with keys. +using Layout = + std::unordered_map, bool>>; + +Layout LayoutOf(const std::unordered_map& default_properties, + const std::vector& credentials) { + Layout layout; + for (const auto* credential : S3CredentialsOf(credentials)) { + auto properties = PropertiesFor(default_properties, *credential); + const bool keys = HasKeys(properties); + layout.try_emplace(CanonicalizeS3Scheme(credential->prefix), + SettingsOf(std::move(properties)), keys); + } + return layout; +} + +// A refresh may only rotate keys, so the current clients serve it: it must keep +// the layout of `current`, and its tokens must be usable. +Status ValidateRefreshed( + const std::unordered_map& default_properties, + const std::vector& current, + const std::vector& credentials) { for (const auto& credential : credentials) { ICEBERG_RETURN_UNEXPECTED(credential.Validate()); } - const auto s3_credentials = S3CredentialsOf(credentials); - if (s3_credentials.empty()) { - // Installing this would drop working credentials for whatever ambient - // identity the AWS chain finds. Java refuses an empty list too. - return NotFound("Refreshed credentials contain no S3-compatible prefix"); + if (LayoutOf(default_properties, credentials) != + LayoutOf(default_properties, current)) { + return InvalidArgument( + "Refreshed credentials change the S3 prefixes, their settings or keys"); } - for (const auto* credential : s3_credentials) { - ICEBERG_ASSIGN_OR_RAISE(auto expires_at, ExpiryOf(*credential)); + const auto now = std::chrono::system_clock::now(); + for (const auto* credential : S3CredentialsOf(credentials)) { + ICEBERG_ASSIGN_OR_RAISE( + auto expires_at, + ExpiryOf(PropertiesFor(default_properties, *credential), credential->prefix)); if (expires_at.has_value() && *expires_at <= now) { return InvalidArgument("Credential \"{}\" has already expired", credential->prefix); } @@ -432,13 +502,6 @@ Status ValidateRefreshed(const std::vector& credentials) { return {}; } -// Whether `config` carries vended keys, which reach the S3 client through a -// VendedCredentialsProvider. -bool HasKeys(const std::unordered_map& config) { - return FindProperty(config, S3Properties::kAccessKeyId) != nullptr && - FindProperty(config, S3Properties::kSecretAccessKey) != nullptr; -} - // Requires HasKeys(config). Aws::Auth::AWSCredentials AwsCredentialsOf( const std::unordered_map& config) { @@ -454,28 +517,10 @@ Aws::Auth::AWSCredentials AwsCredentialsOf( return credentials; } -// Whether `refreshed` serves the same prefix with the same settings as -// `current`, so that only its keys need to reach the existing S3 client. -bool DiffersOnlyInKeys(const StorageCredential& current, - const StorageCredential& refreshed) { - if (current.prefix != refreshed.prefix || !HasKeys(refreshed.config)) { - return false; - } - const auto settings = [](std::unordered_map config) { - for (const auto key : - {S3Properties::kAccessKeyId, S3Properties::kSecretAccessKey, - S3Properties::kSessionToken, S3Properties::kSessionTokenExpiresAtMs}) { - config.erase(std::string(key)); - } - return config; - }; - return settings(current.config) == settings(refreshed.config); -} - class ArrowS3FileIO; -// Where the S3 clients of an ArrowS3FileIO get their keys. Open handles keep -// clients alive past the FileIO, which therefore detaches on destruction. +// Gives S3 clients their FileIO's current keys. Clients may outlive the +// FileIO, which detaches on destruction. class KeySource { public: explicit KeySource(ArrowS3FileIO* io) : io_(io) {} @@ -485,31 +530,32 @@ class KeySource { io_ = nullptr; } - /// \brief The keys the FileIO vends for the canonical `prefix`, after the - /// refresh a request may run; nullopt once it is gone or vends none for - /// `prefix`. - std::optional KeysFor(std::string_view prefix); + /// \brief Refreshes if due, then returns the FileIO's keys for `prefix` and + /// `settings`; nullopt if it has none or is gone. + std::optional KeysFor( + std::string_view prefix, + const std::unordered_map& settings); private: std::shared_mutex mutex_; ArrowS3FileIO* io_; }; -// Signs an S3 client's requests with the keys its FileIO currently vends for -// the client's prefix. The AWS SDK asks before each request, so a refresh -// reaches handles and streams opened before it, even those on a client it -// replaced, as each client's VendedCredentialsProvider does in Java. Without -// current keys for the prefix, the last ones stay. +// Supplies an S3 client's keys. The AWS SDK asks before each request, so open +// handles pick up refreshed keys, as in Java. Keys vended with other settings +// (e.g. another endpoint) are never used; the last keys stay instead. class VendedCredentialsProvider final : public Aws::Auth::AWSCredentialsProvider { public: VendedCredentialsProvider(std::shared_ptr source, std::string prefix, + std::unordered_map settings, Aws::Auth::AWSCredentials credentials) : source_(std::move(source)), prefix_(std::move(prefix)), + settings_(std::move(settings)), credentials_(std::move(credentials)) {} Aws::Auth::AWSCredentials GetAWSCredentials() override { - auto current = source_->KeysFor(prefix_); + auto current = source_->KeysFor(prefix_, settings_); std::lock_guard lock(mutex_); if (current.has_value()) { credentials_ = std::move(current).value(); @@ -520,6 +566,7 @@ class VendedCredentialsProvider final : public Aws::Auth::AWSCredentialsProvider private: const std::shared_ptr source_; const std::string prefix_; + const std::unordered_map settings_; std::mutex mutex_; Aws::Auth::AWSCredentials credentials_; }; @@ -563,31 +610,21 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { private: friend class KeySource; - /// \brief The S3 client serving the locations under `prefix`. - struct Delegate { - std::string prefix; - std::shared_ptr file_io; - // Whether the client signs with the keys vended for `prefix`, which it - // reads through a VendedCredentialsProvider. - bool vended_keys = false; - // When those keys stop being valid, if known. - std::optional expires_at; - }; - - using DelegatesByPrefix = std::vector; + using DelegatesByPrefix = + std::vector>>; + using ExpiriesByPrefix = + std::unordered_map; /// \brief Delegate serving `location`, pinned by the caller against a /// concurrent credential install. Result> FileIOForPath(std::string_view location); - /// \brief Longest-prefix match against one consistent view of the delegates, - /// or `fallback` if none matches. - /// - /// Fails once the matching delegate's keys have expired: a handle built on - /// them would only fail at S3. + /// \brief Longest-prefix match, or `fallback` if none matches. Fails with + /// TokenExpired once the matched prefix's keys have expired. static Result> MatchDelegate( const std::shared_ptr& fallback, - const DelegatesByPrefix& by_prefix, std::string_view location); + const DelegatesByPrefix& by_prefix, const ExpiriesByPrefix& expiries, + std::string_view location); /// \brief Build a delegate for each credential this FileIO can serve. /// @@ -597,54 +634,37 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { Result BuildDelegates( const std::vector& storage_credentials) const; - /// \brief Swap in credentials and delegates, handing back the retired ones. + /// \brief Swap in credentials and, unless null, delegates, handing back the + /// retired ones. A refresh keeps the layout, so the delegates stay. /// /// Callers must hold `mutex_` exclusively and let the returned generation /// destruct only after releasing it: tearing down an S3 client can block on /// in-flight requests, which would stall every operation. void InstallCredentials(std::vector& storage_credentials, - DelegatesByPrefix& delegates); + 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. + /// \brief Whether a refresh is due and not backed off. Requires `mutex_`. bool RefreshDue() const; - /// \brief Whether the installed credentials have already stopped being valid. - /// - /// Callers must hold `mutex_`, at least shared. + /// \brief Whether the installed credentials have expired. Requires `mutex_`. 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. + /// \brief When to retry after a failed refresh; never past the expiry. + /// Requires `mutex_`. std::chrono::steady_clock::time_point BackoffUntil() const; - /// \brief Replace the credentials once they are close to expiring. - /// - /// Called before each handle is created, and by the S3 clients before each - /// request, so open handles see the new keys too. A failure is logged and - /// keeps the current credentials, which MatchDelegate refuses once expired. + /// \brief Refresh the credentials if they are close to expiring. /// - /// \param from_request Whether an S3 client calls while signing a request. - /// Arrow then holds its client registry, which building a client - /// locks exclusively, so a refresh needing new clients is left to the - /// next file operation. Nor does it wait for another refresh, which - /// may be building clients. - void MaybeRefreshCredentials(bool from_request = false); - - /// \brief The current delegates, with `refreshed`'s expiries, if keys are all - /// it changes: their S3 clients stay and read the new keys through their - /// providers. Nullopt if it needs new clients. - static std::optional WithRefreshedKeys( - const std::vector& current, const DelegatesByPrefix& delegates, - const std::vector& refreshed); - - /// \brief The keys installed for the canonical `prefix`, if any. - std::optional KeysFor(std::string_view prefix) const; + /// Called before file operations and by S3 clients before each request. A + /// refresh only rotates keys and builds no clients, so a request can run it + /// or, with expired credentials, wait for one under way. A failure keeps the + /// current credentials. + void MaybeRefreshCredentials(); + + /// \brief The installed keys for `prefix`, if its credential has `settings`. + std::optional KeysFor( + std::string_view prefix, + const std::unordered_map& settings) const; std::shared_ptr default_file_io_; std::unordered_map default_properties_; @@ -654,16 +674,13 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { mutable std::shared_mutex mutex_; std::vector storage_credentials_; DelegatesByPrefix file_io_by_prefix_; + ExpiriesByPrefix expiries_by_prefix_; std::shared_ptr provider_; CredentialsExpiry expiry_; std::chrono::steady_clock::time_point retry_refresh_at_; - // Until when requests leave the refresh to file operations, after finding - // that it needs new clients. - std::chrono::steady_clock::time_point request_refresh_at_; // Bumped on each install so a refresh cannot overwrite newer credentials. uint64_t credential_generation_ = 0; - // Held across a refresh so concurrent operations skip it. Timed, so waiting - // on it is bounded. + // Serializes refreshes; timed so waits are bounded. std::timed_mutex refresh_mutex_; }; @@ -675,7 +692,7 @@ Status ArrowS3FileIO::InitializeStorageCredentials( { std::unique_lock lock(mutex_); provider_.swap(provider); - InstallCredentials(credentials, delegates); + InstallCredentials(credentials, &delegates); } return {}; } @@ -686,7 +703,7 @@ Status ArrowS3FileIO::SetStorageCredentials( auto credentials = storage_credentials; { std::unique_lock lock(mutex_); - InstallCredentials(credentials, delegates); + InstallCredentials(credentials, &delegates); } // `credentials` and `delegates` now hold the retired generation and destruct // here, outside the lock. @@ -705,24 +722,17 @@ Result ArrowS3FileIO::BuildDelegates( if (!IsS3CredentialPrefix(credential.prefix)) { continue; } - auto properties = default_properties_; - for (const auto& [key, value] : credential.config) { - properties[key] = value; - } auto prefix = CanonicalizeS3Scheme(credential.prefix); + const auto properties = PropertiesFor(default_properties_, credential); std::shared_ptr keys; - if (HasKeys(credential.config)) { + // Signs with what ConfigureS3Options would, defaults included. + if (HasKeys(properties)) { keys = std::make_shared( - key_source_, prefix, AwsCredentialsOf(credential.config)); + key_source_, prefix, SettingsOf(properties), AwsCredentialsOf(properties)); } - const bool vended_keys = keys != nullptr; ICEBERG_ASSIGN_OR_RAISE(auto fs, BuildArrowS3FileSystem(properties, std::move(keys))); - delegates.push_back( - {.prefix = std::move(prefix), - .file_io = std::make_shared(std::move(fs)), - .vended_keys = vended_keys, - // Unusable expiries are reported by ExpiryOf when installed. - .expires_at = ExpiryOf(credential).value_or(std::nullopt)}); + delegates.emplace_back(std::move(prefix), + std::make_shared(std::move(fs))); } if (delegates.empty() && !storage_credentials.empty()) { // Silent skipping of every vended credential is hard to diagnose: S3 access @@ -736,11 +746,13 @@ Result ArrowS3FileIO::BuildDelegates( } void ArrowS3FileIO::InstallCredentials( - std::vector& storage_credentials, DelegatesByPrefix& delegates) { - file_io_by_prefix_.swap(delegates); - expiry_ = ExpiryOf(storage_credentials); + std::vector& storage_credentials, DelegatesByPrefix* delegates) { + if (delegates != nullptr) { + file_io_by_prefix_.swap(*delegates); + } + expiry_ = ExpiryOf(default_properties_, storage_credentials); + expiries_by_prefix_ = ExpiriesOf(default_properties_, storage_credentials); retry_refresh_at_ = {}; - request_refresh_at_ = {}; ++credential_generation_; storage_credentials_.swap(storage_credentials); } @@ -763,8 +775,7 @@ std::chrono::steady_clock::time_point ArrowS3FileIO::BackoffUntil() const { if (expiry_.expires_at.has_value()) { const auto remaining = std::chrono::duration_cast( *expiry_.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. + // Retry sooner while still valid; once expired, use the full backoff. if (remaining > std::chrono::milliseconds::zero()) { delay = std::clamp( remaining, @@ -775,82 +786,56 @@ std::chrono::steady_clock::time_point ArrowS3FileIO::BackoffUntil() const { return std::chrono::steady_clock::now() + delay; } -void ArrowS3FileIO::MaybeRefreshCredentials(bool from_request) { - const auto due = [&] { - return provider_ && RefreshDue() && - (!from_request || std::chrono::steady_clock::now() >= request_refresh_at_); - }; +void ArrowS3FileIO::MaybeRefreshCredentials() { { - // Cheap pre-check, so the common case costs one shared lock and no more. + // Fast path. std::shared_lock lock(mutex_); - if (!due()) { + if (!provider_ || !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. - if (from_request) { - return; - } + // Another caller is refreshing; keep the current keys unless expired. { 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; } } - // Snapshot the provider, credentials and their generation together. std::shared_ptr provider; uint64_t generation = 0; std::vector current; - DelegatesByPrefix current_delegates; { - // Whoever held the lock may also have just finished, leaving nothing to do. + // The previous holder may have just refreshed. std::shared_lock lock(mutex_); - if (!due()) { + if (!provider_ || !RefreshDue()) { return; } provider = provider_; generation = credential_generation_; current = storage_credentials_; - current_delegates = file_io_by_prefix_; } - // Outside `mutex_`: both are slow and must not block readers. - Status status; - DelegatesByPrefix delegates; - bool needs_clients = false; + // Slow, so done without holding `mutex_`. auto refreshed = provider->Load(); + Status status; if (refreshed.has_value()) { - status = ValidateRefreshed(*refreshed); + status = ValidateRefreshed(default_properties_, current, *refreshed); } else { status = ::iceberg::unexpected(refreshed.error()); } - if (status.has_value()) { - if (auto same_clients = WithRefreshedKeys(current, current_delegates, *refreshed); - same_clients.has_value()) { - delegates = std::move(same_clients).value(); - } else if (from_request) { - needs_clients = true; - } else if (auto built = BuildDelegates(*refreshed); built.has_value()) { - delegates = std::move(built).value(); - } else { - status = ::iceberg::unexpected(built.error()); - } - } std::unique_lock lock(mutex_); - // Credentials installed meanwhile supersede this refresh: what it fetched is - // by now the older set. + // Credentials installed meanwhile are newer than what this fetched. const bool superseded = credential_generation_ != generation; if (!status.has_value()) { - // Reported either way, so a failing catalog stays visible. + // Logged even if superseded, so catalog failures stay visible. if (superseded) { ICEBERG_LOG_WARN( "Failed to refresh vended storage credentials ({}); they have since been " @@ -871,89 +856,67 @@ void ArrowS3FileIO::MaybeRefreshCredentials(bool from_request) { if (superseded) { return; } - if (needs_clients) { - // Fetched again by the next file operation, so it is not stale by then. - request_refresh_at_ = std::chrono::steady_clock::now() + kRefreshRetryBackoff; - ICEBERG_LOG_DEBUG( - "Refreshed vended storage credentials need new S3 clients; leaving them to " - "the next file operation"); - return; - } - // `delegates`/`*refreshed` take the retired generation; both are declared - // before the lock, so it destructs only after the lock releases. - InstallCredentials(*refreshed, delegates); + InstallCredentials(*refreshed, /*delegates=*/nullptr); if (RefreshDue()) { - // Tokens shorter-lived than the lead time come back due again at once. + // Tokens shorter-lived than the lead time would be due again at once. retry_refresh_at_ = BackoffUntil(); } } -std::optional ArrowS3FileIO::WithRefreshedKeys( - const std::vector& current, const DelegatesByPrefix& delegates, - const std::vector& refreshed) { - // BuildDelegates made one delegate per S3 credential, in order. - const auto current_s3 = S3CredentialsOf(current); - const auto refreshed_s3 = S3CredentialsOf(refreshed); - if (current_s3.size() != delegates.size() || refreshed_s3.size() != delegates.size()) { - return std::nullopt; - } - auto result = delegates; - for (size_t i = 0; i < result.size(); ++i) { - if (!result[i].vended_keys || !DiffersOnlyInKeys(*current_s3[i], *refreshed_s3[i])) { - return std::nullopt; - } - // ValidateRefreshed has checked the expiry. - result[i].expires_at = ExpiryOf(*refreshed_s3[i]).value_or(std::nullopt); - } - return result; -} - std::optional ArrowS3FileIO::KeysFor( - std::string_view prefix) const { + std::string_view prefix, + const std::unordered_map& settings) const { std::shared_lock lock(mutex_); for (const auto* credential : S3CredentialsOf(storage_credentials_)) { - if (HasKeys(credential->config) && - CanonicalizeS3Scheme(credential->prefix) == prefix) { - return AwsCredentialsOf(credential->config); + if (CanonicalizeS3Scheme(credential->prefix) != prefix) { + continue; + } + const auto properties = PropertiesFor(default_properties_, *credential); + if (!HasKeys(properties) || SettingsOf(properties) != settings) { + return std::nullopt; } + return AwsCredentialsOf(properties); } return std::nullopt; } Result> ArrowS3FileIO::MatchDelegate( const std::shared_ptr& fallback, - const DelegatesByPrefix& by_prefix, std::string_view location) { + const DelegatesByPrefix& by_prefix, const ExpiriesByPrefix& expiries, + std::string_view location) { if (by_prefix.empty()) { return fallback; } const std::string canonical = CanonicalizeS3Scheme(location); - const Delegate* best = nullptr; + auto best = fallback; size_t best_len = 0; - for (const auto& delegate : by_prefix) { - if (delegate.prefix.size() > best_len && canonical.starts_with(delegate.prefix)) { - best = &delegate; - best_len = delegate.prefix.size(); + const std::string* best_prefix = nullptr; + for (const auto& [prefix, file_io] : by_prefix) { + if (prefix.size() > best_len && canonical.starts_with(prefix)) { + best = file_io; + best_len = prefix.size(); + best_prefix = &prefix; } } - if (best == nullptr) { + if (best_prefix == nullptr) { return fallback; } - if (best->expires_at.has_value() && - std::chrono::system_clock::now() >= *best->expires_at) { + if (const auto expiry = expiries.find(*best_prefix); + expiry != expiries.end() && std::chrono::system_clock::now() >= expiry->second) { return TokenExpired( "Vended credentials for \"{}\" have expired and could not be refreshed", - best->prefix); + *best_prefix); } - return best->file_io; + return best; } Result> ArrowS3FileIO::FileIOForPath( std::string_view location) { MaybeRefreshCredentials(); - std::shared_lock lock(mutex_); - return MatchDelegate(default_file_io_, file_io_by_prefix_, location); + return MatchDelegate(default_file_io_, file_io_by_prefix_, expiries_by_prefix_, + location); } Result> ArrowS3FileIO::NewInputFile( @@ -980,32 +943,23 @@ Status ArrowS3FileIO::DeleteFile(const std::string& file_location) { } Status ArrowS3FileIO::DeleteFiles(const std::vector& file_locations) { - MaybeRefreshCredentials(); - - // Keep one delegate snapshot for the whole batch. - std::shared_ptr fallback; - DelegatesByPrefix by_prefix; - { - std::shared_lock lock(mutex_); - fallback = default_file_io_; - by_prefix = file_io_by_prefix_; - } + // Each file is resolved afresh, as a long batch may outlive its credentials. return BulkDeleteFiles(file_locations, delete_executor.Get(), delete_parallelism_, - [&](const std::string& location) -> Status { - ICEBERG_ASSIGN_OR_RAISE( - auto file_io, - MatchDelegate(fallback, by_prefix, location)); + [this](const std::string& location) -> Status { + ICEBERG_ASSIGN_OR_RAISE(auto file_io, FileIOForPath(location)); return file_io->DeleteFile(location); }); } -std::optional KeySource::KeysFor(std::string_view prefix) { +std::optional KeySource::KeysFor( + std::string_view prefix, + const std::unordered_map& settings) { std::shared_lock lock(mutex_); if (io_ == nullptr) { return std::nullopt; } - io_->MaybeRefreshCredentials(/*from_request=*/true); - return io_->KeysFor(prefix); + io_->MaybeRefreshCredentials(); + return io_->KeysFor(prefix, settings); } } // namespace diff --git a/src/iceberg/catalog/rest/http_client.cc b/src/iceberg/catalog/rest/http_client.cc index 1bfd27109f..42dac26d8c 100644 --- a/src/iceberg/catalog/rest/http_client.cc +++ b/src/iceberg/catalog/rest/http_client.cc @@ -146,16 +146,27 @@ ErrorResponse BuildDefaultErrorResponse(const cpr::Response& response) { }; } +/// \brief Fails a GET whose transfer failed, even after a success status. +/// +/// Not for writes: one may have taken effect, and a failed commit's files are +/// cleaned up. +Status CheckGetTransfer(const cpr::Response& response) { + if (response.error.code == cpr::ErrorCode::OPERATION_TIMEDOUT) { + return IOError("HTTP request timed out: {}", response.error.message); + } + if (response.error && IsSuccessful(response.status_code)) { + return IOError("HTTP request failed after status {}: {}", response.status_code, + response.error.message); + } + return {}; +} + /// \brief Handles failure responses by invoking the provided error handler. Status HandleFailureResponse(const cpr::Response& response, const ErrorHandler& error_handler) { if (IsSuccessful(response.status_code)) { return {}; } - // No response to hand the error handler. - if (response.error.code == cpr::ErrorCode::OPERATION_TIMEDOUT) { - return IOError("HTTP request timed out: {}", response.error.message); - } auto parse_result = error_handler.ParseResponse( static_cast(response.status_code), response.text); const ErrorResponse final_error = @@ -191,6 +202,7 @@ Result HttpClient::Get( cpr::Get(cpr::Url{authenticated.url}, ToCprHeader(authenticated), *connection_pool_, cpr::Timeout{timeout.value_or(std::chrono::milliseconds::zero())}); + ICEBERG_RETURN_UNEXPECTED(CheckGetTransfer(response)); ICEBERG_RETURN_UNEXPECTED(HandleFailureResponse(response, error_handler)); HttpResponse http_response; http_response.impl_ = std::make_unique(std::move(response)); diff --git a/src/iceberg/catalog/rest/rest_catalog.cc b/src/iceberg/catalog/rest/rest_catalog.cc index 58ee2db11d..38728deac7 100644 --- a/src/iceberg/catalog/rest/rest_catalog.cc +++ b/src/iceberg/catalog/rest/rest_catalog.cc @@ -70,8 +70,7 @@ namespace iceberg::rest { namespace { -// Bounds a credential refresh, which runs on the path of a file operation. The -// FileIO keeps its current credentials, or fails once they have expired. +// A refresh runs inline in file operations and S3 requests, so bound it. constexpr auto kLoadCredentialsTimeout = std::chrono::seconds(10); class RestStorageCredentialProvider final : public StorageCredentialProvider { diff --git a/src/iceberg/test/arrow_s3_file_io_test.cc b/src/iceberg/test/arrow_s3_file_io_test.cc index 0499bc30a7..ef73622e36 100644 --- a/src/iceberg/test/arrow_s3_file_io_test.cc +++ b/src/iceberg/test/arrow_s3_file_io_test.cc @@ -172,8 +172,6 @@ class MockStorageCredentialProvider : public StorageCredentialProvider { MOCK_METHOD((Result>), Load, (), (override)); }; -constexpr auto kOutlastsAShortenedBackoff = std::chrono::milliseconds(1200); - std::string ExpiresAtMs(std::chrono::milliseconds valid_for) { return std::to_string( std::chrono::duration_cast( @@ -200,6 +198,16 @@ Status CheckReadWrite(FileIO& io, const std::string& object_uri, return io.DeleteFile(object_uri); } +// Reads `length` bytes of `file` through a fresh stream. +Result ReadAll(InputFile& file, size_t length) { + ICEBERG_ASSIGN_OR_RAISE(auto stream, file.Open()); + std::string content(length, '\0'); + ICEBERG_RETURN_UNEXPECTED( + stream->ReadFully(0, std::as_writable_bytes(std::span(content)))); + ICEBERG_RETURN_UNEXPECTED(stream->Close()); + return content; +} + class InlineDeleteExecutor : public Executor { public: Status Submit(ExecutorTask task) override { @@ -539,60 +547,41 @@ 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; - auto provider = std::make_shared(); - EXPECT_CALL(*provider, Load()) - .WillRepeatedly([&]() -> Result> { - ++refresh_calls; - return refreshed; - }); - ASSERT_THAT(credentialed->InitializeStorageCredentials( - ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), provider), - IsOk()); - EXPECT_EQ(refresh_calls, 0); - - 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; - auto provider = std::make_shared(); - EXPECT_CALL(*provider, Load()) - .WillRepeatedly([&]() -> Result> { - ++refresh_calls; - return std::vector{}; - }); +TEST_F(ArrowS3FileIOTest, RefreshesOnlyCredentialsCloseToExpiry) { + struct Case { + std::string_view name; + std::vector initial; + bool refreshes; + }; + const std::vector cases = { + {.name = "close to expiry", + .initial = ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), + .refreshes = true}, + {.name = "not close to expiry", + .initial = ExpiringCredentials(std::chrono::hours(1), "valid-key"), + .refreshes = false}, + {.name = "static keys", + .initial = {{.prefix = "s3", .config = BadS3Credentials()}}, + .refreshes = false}}; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.name); + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); - ASSERT_THAT(credentialed->InitializeStorageCredentials( - ExpiringCredentials(std::chrono::hours(1), "access-key"), provider), - IsOk()); - EXPECT_THAT(result.value()->NewInputFile("s3://bucket/key"), IsOk()); - EXPECT_EQ(refresh_calls, 0); + const auto refreshed = ExpiringCredentials(std::chrono::hours(1), "refreshed-key"); + auto provider = std::make_shared(); + ON_CALL(*provider, Load()).WillByDefault(::testing::Return(refreshed)); + EXPECT_CALL(*provider, Load()).Times(test_case.refreshes ? 1 : 0); + ASSERT_THAT(credentialed->InitializeStorageCredentials(test_case.initial, provider), + IsOk()); - 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); + // The replacement is not close to expiry, so the second does not refresh. + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(credentialed->credentials(), + test_case.refreshes ? refreshed : test_case.initial); + } } TEST_F(ArrowS3FileIOTest, KeepsProviderWhenCredentialsAreUpdated) { @@ -632,44 +621,6 @@ TEST_F(ArrowS3FileIOTest, FailedInitializationDoesNotInstallProvider) { EXPECT_EQ(credentialed->credentials(), initial); } -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; - auto provider = std::make_shared(); - EXPECT_CALL(*provider, Load()) - .WillRepeatedly([&]() -> Result> { - ++refresh_calls; - return ExpiringCredentials(std::chrono::hours(1), "refreshed-key"); - }); - ASSERT_THAT(credentialed->InitializeStorageCredentials( - ExpiringCredentials(std::chrono::minutes(1), "expiring-key"), provider), - 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()); @@ -805,63 +756,6 @@ TEST_F(ArrowS3FileIOTest, RefreshesSessionCredentialsWithoutAUsableExpiry) { } } -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; - auto provider = std::make_shared(); - EXPECT_CALL(*provider, Load()) - .WillRepeatedly([&]() -> Result> { - ++refresh_calls; - return undatable; - }); - ASSERT_THAT(credentialed->InitializeStorageCredentials(undatable, provider), 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; - auto provider = std::make_shared(); - EXPECT_CALL(*provider, Load()) - .WillRepeatedly([&]() -> 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->InitializeStorageCredentials(credentials, provider), 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()); @@ -885,62 +779,52 @@ TEST_F(ArrowS3FileIOTest, BacksOffWhenTheReplacementIsAlsoCloseToExpiry) { 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; - auto provider = std::make_shared(); - EXPECT_CALL(*provider, Load()) - .WillRepeatedly([&]() -> Result> { - ++refresh_calls; - return NotFound("catalog unreachable"); - }); - const auto expiring = ExpiringCredentials(std::chrono::minutes(1), "expiring-key"); - ASSERT_THAT(credentialed->InitializeStorageCredentials(expiring, provider), 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, KeepsCredentialsWhenTheReplacementIsUnusable) { +TEST_F(ArrowS3FileIOTest, KeepsCredentialsWhenTheRefreshFails) { struct Case { std::string_view name; std::optional expires_at_ms; + std::string prefix = "s3"; + std::unordered_map settings; + bool keys = true; + bool load_fails = false; }; + const auto valid = ExpiresAtMs(std::chrono::hours(1)); const std::vector cases = { + {.name = "catalog error", .expires_at_ms = valid, .load_fails = true}, {.name = "expired", .expires_at_ms = ExpiresAtMs(-std::chrono::minutes(1))}, {.name = "no expiry", .expires_at_ms = std::nullopt}, - {.name = "unparseable expiry", .expires_at_ms = "not-a-number"}}; + {.name = "unparseable expiry", .expires_at_ms = "not-a-number"}, + // A refresh may only rotate keys. + {.name = "prefix changed", .expires_at_ms = valid, .prefix = "s3://other/"}, + {.name = "setting changed", + .expires_at_ms = valid, + .settings = {{std::string(S3Properties::kConnectTimeoutMs), "5000"}}}, + {.name = "no keys", .expires_at_ms = valid, .keys = false}}; for (const auto& test_case : cases) { SCOPED_TRACE(test_case.name); ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); auto* credentialed = io->AsSupportsStorageCredentials(); ASSERT_NE(credentialed, nullptr); - std::unordered_map config = { - {std::string(S3Properties::kAccessKeyId), "unusable-key"}, - {std::string(S3Properties::kSecretAccessKey), "secret"}, - {std::string(S3Properties::kSessionToken), "token"}}; + auto config = test_case.settings; + if (test_case.keys) { + config.insert({{std::string(S3Properties::kAccessKeyId), "unusable-key"}, + {std::string(S3Properties::kSecretAccessKey), "secret"}}); + } + config.emplace(std::string(S3Properties::kSessionToken), "token"); if (test_case.expires_at_ms.has_value()) { config[std::string(S3Properties::kSessionTokenExpiresAtMs)] = *test_case.expires_at_ms; } const std::vector unusable = { - {.prefix = "s3", .config = std::move(config)}}; - int refresh_calls = 0; + {.prefix = test_case.prefix, .config = std::move(config)}}; auto provider = std::make_shared(); + // Once: a failed refresh backs off. EXPECT_CALL(*provider, Load()) - .WillRepeatedly([&]() -> Result> { - ++refresh_calls; + .WillOnce([&]() -> Result> { + if (test_case.load_fails) { + return NotFound("catalog unreachable"); + } return unusable; }); const auto expiring = ExpiringCredentials(std::chrono::minutes(1), "expiring-key"); @@ -951,10 +835,78 @@ TEST_F(ArrowS3FileIOTest, KeepsCredentialsWhenTheReplacementIsUnusable) { EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); EXPECT_EQ(credentialed->credentials(), expiring); EXPECT_TRUE(HasWarning(*logger, "Failed to refresh")); - - // Rejected like any failed refresh, so it is not asked for again at once. EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); - EXPECT_EQ(refresh_calls, 1); + } +} + +TEST_F(ArrowS3FileIOTest, AcceptsRefreshedPrefixesInAnotherOrder) { + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + auto initial = + ExpiringCredentials(std::chrono::minutes(1), "expiring-key", "s3://first/"); + initial.push_back( + ExpiringCredentials(std::chrono::minutes(1), "expiring-key", "s3://second/") + .front()); + auto refreshed = + ExpiringCredentials(std::chrono::hours(1), "refreshed-key", "s3://second/"); + refreshed.push_back( + ExpiringCredentials(std::chrono::hours(1), "refreshed-key", "s3://first/").front()); + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).WillOnce(::testing::Return(refreshed)); + ASSERT_THAT(credentialed->InitializeStorageCredentials(initial, provider), IsOk()); + + EXPECT_THAT(io->NewInputFile("s3://first/key"), IsOk()); + EXPECT_EQ(credentialed->credentials(), refreshed); +} + +TEST_F(ArrowS3FileIOTest, TracksTheExpiryOfTheTokenItSignsWith) { + const std::unordered_map expired_token = { + {std::string(S3Properties::kSessionToken), "token"}, + {std::string(S3Properties::kSessionTokenExpiresAtMs), "0"}}; + auto keys_and = [](std::unordered_map config) { + config.merge(BadS3Credentials()); + return config; + }; + struct Case { + std::string_view name; + std::unordered_map properties; + std::vector vended; + bool expired; + }; + const std::vector cases = { + // Without an access key pair the AWS default chain signs instead. + {.name = "token without keys", + .vended = {{.prefix = "s3://bucket/", .config = expired_token}}, + .expired = false}, + {.name = "expiry without a token", + .vended = {{.prefix = "s3://bucket/", + .config = keys_and( + {{std::string(S3Properties::kSessionTokenExpiresAtMs), "0"}})}}, + .expired = false}, + // Inherited from the FileIO's properties, it signs with the vended keys. + {.name = "inherited token", + .properties = expired_token, + .vended = {{.prefix = "s3://bucket/", .config = BadS3Credentials()}}, + .expired = true}, + // Only the first credential for a prefix is used. + {.name = "repeated prefix", + .vended = {{.prefix = "s3://bucket/", .config = BadS3Credentials()}, + {.prefix = "s3://bucket/", .config = keys_and(expired_token)}}, + .expired = false}}; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.name); + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO(test_case.properties)); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + ASSERT_THAT(credentialed->SetStorageCredentials(test_case.vended), IsOk()); + + if (test_case.expired) { + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsError(ErrorKind::kTokenExpired)); + } else { + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + } } } @@ -1127,20 +1079,20 @@ TEST_F(ArrowS3FileIOTest, RefreshReachesHandlesOpenedBeforeIt) { "credential refresh test"; } - // A refresh that changes only keys keeps the S3 client, and the handle's own - // request runs it. One that changes a setting replaces the client, which a - // request cannot do, so a file operation runs it. Either way the handle - // opened before it, still on the first client, must sign with the new keys. + // New keys reach a handle opened before them, from a refresh its own request + // runs, but not keys installed with other client settings. struct Case { std::string_view name; - std::unordered_map refreshed_settings; - bool file_operation_refreshes; + bool refresh; + std::unordered_map settings; + bool handle_reads; }; const std::vector cases = { - {.name = "keys only", .refreshed_settings = {}, .file_operation_refreshes = false}, - {.name = "new client", - .refreshed_settings = {{std::string(S3Properties::kConnectTimeoutMs), "5000"}}, - .file_operation_refreshes = true}}; + {.name = "refresh", .refresh = true, .settings = {}, .handle_reads = true}, + {.name = "install with another setting", + .refresh = false, + .settings = {{std::string(S3Properties::kConnectTimeoutMs), "5000"}}, + .handle_reads = false}}; for (const auto& test_case : cases) { SCOPED_TRACE(test_case.name); ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO(properties)); @@ -1152,16 +1104,19 @@ TEST_F(ArrowS3FileIOTest, RefreshReachesHandlesOpenedBeforeIt) { constexpr std::string_view kContent = "read with keys refreshed after it was opened"; ASSERT_THAT(io->WriteFile(uri, kContent), IsOk()); - // The vended keys are bad, and fall due for refresh just after the handle - // is opened. Only the refreshed keys can read the object. + // The vended keys are bad and come due right after the handle opens; only + // the new ones can read. constexpr auto kDueIn = std::chrono::seconds(2); - auto refreshed_config = test_case.refreshed_settings; - refreshed_config.insert(*access_key); - refreshed_config.insert(*secret_key); - const std::vector refreshed = { - {.prefix = prefix, .config = std::move(refreshed_config)}}; - auto provider = std::make_shared(); - EXPECT_CALL(*provider, Load()).WillOnce(::testing::Return(refreshed)); + auto new_config = test_case.settings; + new_config.insert(*access_key); + new_config.insert(*secret_key); + const std::vector new_credentials = { + {.prefix = prefix, .config = std::move(new_config)}}; + std::shared_ptr provider; + if (test_case.refresh) { + provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).WillOnce(::testing::Return(new_credentials)); + } ASSERT_THAT( credentialed->InitializeStorageCredentials( {{.prefix = prefix, @@ -1174,20 +1129,87 @@ TEST_F(ArrowS3FileIOTest, RefreshReachesHandlesOpenedBeforeIt) { IsOk()); ICEBERG_UNWRAP_OR_FAIL(auto file, io->NewInputFile(uri, kContent.size())); - std::this_thread::sleep_for(kDueIn + std::chrono::seconds(1)); - if (test_case.file_operation_refreshes) { - ASSERT_THAT(io->NewInputFile(uri), IsOk()); + if (test_case.refresh) { + std::this_thread::sleep_for(kDueIn + std::chrono::seconds(1)); + } else { + ASSERT_THAT(credentialed->SetStorageCredentials(new_credentials), IsOk()); } - ICEBERG_UNWRAP_OR_FAIL(auto stream, file->Open()); - std::string read(kContent.size(), '\0'); - EXPECT_THAT(stream->ReadFully(0, std::as_writable_bytes(std::span(read))), IsOk()); - EXPECT_EQ(read, kContent); - EXPECT_THAT(stream->Close(), IsOk()); - EXPECT_EQ(credentialed->credentials(), refreshed); + auto read = ReadAll(*file, kContent.size()); + if (test_case.handle_reads) { + EXPECT_THAT(read, HasValue(::testing::Eq(std::string(kContent)))); + } else { + // Still signs with the bad keys it was opened with. + EXPECT_THAT(read, ::testing::Not(IsOk())); + } + EXPECT_EQ(credentialed->credentials(), new_credentials); EXPECT_THAT(io->DeleteFile(uri), IsOk()); } } +TEST_F(ArrowS3FileIOTest, RequestWaitsForARefreshUnderWayOnceExpired) { + if (!HasIntegrationEnv()) { + GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test"; + } + + auto properties = PropertiesFromEnv(); + const auto access_key = properties.find(std::string(S3Properties::kAccessKeyId)); + const auto secret_key = properties.find(std::string(S3Properties::kSecretAccessKey)); + if (access_key == properties.end() || secret_key == properties.end()) { + GTEST_SKIP() << "Set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY to enable " + "credential refresh test"; + } + + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO(properties)); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + const auto prefix = ObjectUri("refresh_while_expired/"); + const auto uri = prefix + "object"; + constexpr std::string_view kContent = "read with keys a refresh under way fetched"; + ASSERT_THAT(io->WriteFile(uri, kContent), IsOk()); + + // Bad keys that expire right after the handle opens. The first refresh + // fails; the second, after expiry, is slow and returns working keys. + constexpr auto kExpiresIn = std::chrono::seconds(2); + std::promise refresh_started; + auto started = refresh_started.get_future(); + const std::vector refreshed = { + {.prefix = prefix, .config = {*access_key, *secret_key}}}; + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()) + .WillOnce([]() -> Result> { + return ServiceUnavailable("catalog unavailable"); + }) + .WillOnce([&]() -> Result> { + refresh_started.set_value(); + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + return refreshed; + }); + ASSERT_THAT( + credentialed->InitializeStorageCredentials( + {{.prefix = prefix, + .config = {{std::string(S3Properties::kAccessKeyId), "bad-access-key"}, + {std::string(S3Properties::kSecretAccessKey), "bad-secret"}, + {std::string(S3Properties::kSessionToken), "bad-token"}, + {std::string(S3Properties::kSessionTokenExpiresAtMs), + ExpiresAtMs(kExpiresIn)}}}}, + provider), + IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto file, io->NewInputFile(uri, kContent.size())); + + // Past the expiry, and the backoff after the failed refresh. + std::this_thread::sleep_for(kExpiresIn + std::chrono::milliseconds(500)); + std::thread operation([&] { EXPECT_THAT(io->NewInputFile(uri), IsOk()); }); + // Read while that refresh runs. No fatal assertions before the join. + EXPECT_EQ(started.wait_for(std::chrono::seconds(10)), std::future_status::ready); + auto read = ReadAll(*file, kContent.size()); + operation.join(); + + EXPECT_THAT(read, HasValue(::testing::Eq(std::string(kContent)))); + EXPECT_EQ(credentialed->credentials(), refreshed); + EXPECT_THAT(io->DeleteFile(uri), IsOk()); +} + TEST_F(ArrowS3FileIOTest, DeleteFilesAttemptsEveryFile) { if (!HasIntegrationEnv()) { GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test"; diff --git a/src/iceberg/test/http_client_test.cc b/src/iceberg/test/http_client_test.cc index 98c587247f..e553bfadfe 100644 --- a/src/iceberg/test/http_client_test.cc +++ b/src/iceberg/test/http_client_test.cc @@ -19,9 +19,13 @@ #include "iceberg/catalog/rest/http_client.h" +#include #include +#include #include #include +#include +#include #include #include @@ -32,6 +36,7 @@ # include # include # include +# include #endif #include "iceberg/catalog/rest/auth/auth_session.h" @@ -42,33 +47,132 @@ namespace iceberg::rest { #ifndef _WIN32 -TEST(HttpClientTest, GetFailsOnceItsTimeoutPasses) { - // A listening socket that never accepts: the kernel completes the connection, - // but no response ever comes. - const int server = ::socket(AF_INET, SOCK_STREAM, 0); - ASSERT_GE(server, 0); - struct Closer { - int fd; - ~Closer() { ::close(fd); } - } closer{server}; +namespace { + +// A loopback socket listening on an ephemeral port. +struct Listener { + int fd = -1; + uint16_t port = 0; + + ~Listener() { + if (fd >= 0) { + ::close(fd); + } + } +}; + +void Listen(Listener& listener) { + listener.fd = ::socket(AF_INET, SOCK_STREAM, 0); + ASSERT_GE(listener.fd, 0); sockaddr_in address{}; address.sin_family = AF_INET; address.sin_addr.s_addr = htonl(INADDR_LOOPBACK); - ASSERT_EQ(::bind(server, reinterpret_cast(&address), sizeof(address)), 0); - ASSERT_EQ(::listen(server, 1), 0); + ASSERT_EQ(::bind(listener.fd, reinterpret_cast(&address), sizeof(address)), + 0); + ASSERT_EQ(::listen(listener.fd, 1), 0); socklen_t length = sizeof(address); - ASSERT_EQ(::getsockname(server, reinterpret_cast(&address), &length), 0); - const auto url = std::format("http://127.0.0.1:{}/v1/config", ntohs(address.sin_port)); - - HttpClient client; - auto session = auth::AuthSession::MakeDefault({}); - const auto start = std::chrono::steady_clock::now(); - auto response = - client.Get(url, /*params=*/{}, /*headers=*/{}, *DefaultErrorHandler::Instance(), - *session, std::chrono::milliseconds(200)); - EXPECT_THAT(response, IsError(ErrorKind::kIOError)); - EXPECT_THAT(response, HasErrorMessage("timed out")); - EXPECT_LT(std::chrono::steady_clock::now() - start, std::chrono::seconds(10)); + ASSERT_EQ(::getsockname(listener.fd, reinterpret_cast(&address), &length), + 0); + listener.port = ntohs(address.sin_port); +} + +// Accepts one connection, reads the request head and sends `response`. Then +// ends the connection if `end`, or else waits for the client to close it. +void ServeOnce(int listener, std::string_view response, bool end) { + // Bound every wait, so the test cannot hang. + timeval wait{}; + wait.tv_sec = 10; + ::setsockopt(listener, SOL_SOCKET, SO_RCVTIMEO, &wait, sizeof(wait)); + const int connection = ::accept(listener, nullptr, nullptr); + if (connection < 0) { + return; + } + ::setsockopt(connection, SOL_SOCKET, SO_RCVTIMEO, &wait, sizeof(wait)); + std::string request; + std::array buffer; + while (request.find("\r\n\r\n") == std::string::npos) { + const auto received = ::recv(connection, buffer.data(), buffer.size(), 0); + if (received <= 0) { + break; + } + request.append(buffer.data(), static_cast(received)); + } + ::send(connection, response.data(), response.size(), 0); + if (end) { + ::shutdown(connection, SHUT_WR); + } + // Drain until the client closes, so ours closes cleanly. + while (::recv(connection, buffer.data(), buffer.size(), 0) > 0) { + } + ::close(connection); +} + +} // namespace + +TEST(HttpClientTest, GetFailsWhenItsTransferFails) { + // A 200 arriving first does not make a truncated response a success. + constexpr std::string_view kTruncated = + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n" + "Content-Length: 100\r\n\r\n{"; + struct Case { + std::string_view name; + std::string_view response; + bool end; + std::string_view error; + }; + for (const auto& test_case : + {Case{.name = "no response", .response = "", .end = false, .error = "timed out"}, + Case{.name = "body stalls", + .response = kTruncated, + .end = false, + .error = "timed out"}, + Case{.name = "body ends early", + .response = kTruncated, + .end = true, + .error = "failed after status 200"}}) { + SCOPED_TRACE(test_case.name); + Listener listener; + ASSERT_NO_FATAL_FAILURE(Listen(listener)); + const auto url = std::format("http://127.0.0.1:{}/v1/config", listener.port); + + std::thread server( + [&] { ServeOnce(listener.fd, test_case.response, test_case.end); }); + // Destroy the client before the join, so no pooled connection lingers. + auto response = [&] { + HttpClient client; + auto session = auth::AuthSession::MakeDefault({}); + return client.Get(url, /*params=*/{}, /*headers=*/{}, + *DefaultErrorHandler::Instance(), *session, + std::chrono::milliseconds(500)); + }(); + server.join(); + EXPECT_THAT(response, IsError(ErrorKind::kIOError)); + EXPECT_THAT(response, HasErrorMessage(test_case.error)); + } +} + +TEST(HttpClientTest, PostKeepsASuccessfulStatusWhenTheTransferBreaks) { + // A write that got its status may have taken effect, so it must not fail. + // Here the body arrives but its final chunk does not. + constexpr std::string_view kBody = R"({"metadata-location":"s3://bucket/m.json"})"; + const auto response_bytes = std::format( + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n" + "Transfer-Encoding: chunked\r\n\r\n{:x}\r\n{}\r\n", + kBody.size(), kBody); + Listener listener; + ASSERT_NO_FATAL_FAILURE(Listen(listener)); + const auto url = std::format("http://127.0.0.1:{}/v1/commit", listener.port); + + std::thread server([&] { ServeOnce(listener.fd, response_bytes, /*end=*/true); }); + auto response = [&] { + HttpClient client; + auto session = auth::AuthSession::MakeDefault({}); + return client.Post(url, "{}", /*headers=*/{}, *DefaultErrorHandler::Instance(), + *session); + }(); + server.join(); + ASSERT_THAT(response, IsOk()); + EXPECT_EQ(response->body(), kBody); } #endif From 1ece5c9a429cf301860f841a00ff12b56af0ac27 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Wed, 14 Oct 2026 04:37:32 -0700 Subject: [PATCH 3/7] test(rest): keep the HTTP client alive past its responses A response keeps its curl handle, which is attached to the client's connection pool. Destroying the client first left that handle pointing at a freed pool, so the POST test crashed when the response went away. --- src/iceberg/test/http_client_test.cc | 24 ++++++++++-------------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/src/iceberg/test/http_client_test.cc b/src/iceberg/test/http_client_test.cc index e553bfadfe..092c80d6b6 100644 --- a/src/iceberg/test/http_client_test.cc +++ b/src/iceberg/test/http_client_test.cc @@ -135,16 +135,13 @@ TEST(HttpClientTest, GetFailsWhenItsTransferFails) { ASSERT_NO_FATAL_FAILURE(Listen(listener)); const auto url = std::format("http://127.0.0.1:{}/v1/config", listener.port); + HttpClient client; + auto session = auth::AuthSession::MakeDefault({}); std::thread server( [&] { ServeOnce(listener.fd, test_case.response, test_case.end); }); - // Destroy the client before the join, so no pooled connection lingers. - auto response = [&] { - HttpClient client; - auto session = auth::AuthSession::MakeDefault({}); - return client.Get(url, /*params=*/{}, /*headers=*/{}, - *DefaultErrorHandler::Instance(), *session, - std::chrono::milliseconds(500)); - }(); + auto response = + client.Get(url, /*params=*/{}, /*headers=*/{}, *DefaultErrorHandler::Instance(), + *session, std::chrono::milliseconds(500)); server.join(); EXPECT_THAT(response, IsError(ErrorKind::kIOError)); EXPECT_THAT(response, HasErrorMessage(test_case.error)); @@ -163,13 +160,12 @@ TEST(HttpClientTest, PostKeepsASuccessfulStatusWhenTheTransferBreaks) { ASSERT_NO_FATAL_FAILURE(Listen(listener)); const auto url = std::format("http://127.0.0.1:{}/v1/commit", listener.port); + // Outlives the response, which keeps a curl handle in its connection pool. + HttpClient client; + auto session = auth::AuthSession::MakeDefault({}); std::thread server([&] { ServeOnce(listener.fd, response_bytes, /*end=*/true); }); - auto response = [&] { - HttpClient client; - auto session = auth::AuthSession::MakeDefault({}); - return client.Post(url, "{}", /*headers=*/{}, *DefaultErrorHandler::Instance(), - *session); - }(); + auto response = + client.Post(url, "{}", /*headers=*/{}, *DefaultErrorHandler::Instance(), *session); server.join(); ASSERT_THAT(response, IsOk()); EXPECT_EQ(response->body(), kBody); From 6972755d69d2754f86297b8fc382e60db448bf02 Mon Sep 17 00:00:00 2001 From: Li Jiajia Date: Thu, 15 Oct 2026 00:02:54 -0700 Subject: [PATCH 4/7] fix(io): keep new S3 clients from taking the keys they replace The AWS signer asks the credentials provider while the client is built, before SetStorageCredentials installs its keys, so the provider cached the keys being replaced. A handle whose FileIO was released before its first request then signed with them. A client now asks the FileIO for keys only once its own are installed. --- src/iceberg/arrow/s3/arrow_s3_file_io.cc | 36 +++++++++++++++------ src/iceberg/test/arrow_s3_file_io_test.cc | 39 +++++++++++++++++++++++ 2 files changed, 66 insertions(+), 9 deletions(-) diff --git a/src/iceberg/arrow/s3/arrow_s3_file_io.cc b/src/iceberg/arrow/s3/arrow_s3_file_io.cc index 7653b9dd56..3bdf36d999 100644 --- a/src/iceberg/arrow/s3/arrow_s3_file_io.cc +++ b/src/iceberg/arrow/s3/arrow_s3_file_io.cc @@ -546,16 +546,24 @@ class KeySource { // (e.g. another endpoint) are never used; the last keys stay instead. class VendedCredentialsProvider final : public Aws::Auth::AWSCredentialsProvider { public: - VendedCredentialsProvider(std::shared_ptr source, std::string prefix, + VendedCredentialsProvider(std::shared_ptr source, + std::shared_ptr> installed, + std::string prefix, std::unordered_map settings, Aws::Auth::AWSCredentials credentials) : source_(std::move(source)), + installed_(std::move(installed)), prefix_(std::move(prefix)), settings_(std::move(settings)), credentials_(std::move(credentials)) {} Aws::Auth::AWSCredentials GetAWSCredentials() override { - auto current = source_->KeysFor(prefix_, settings_); + std::optional current; + // The SDK also asks while building the client, when the FileIO still holds + // the keys its install replaces. + if (installed_->load()) { + current = source_->KeysFor(prefix_, settings_); + } std::lock_guard lock(mutex_); if (current.has_value()) { credentials_ = std::move(current).value(); @@ -565,6 +573,8 @@ class VendedCredentialsProvider final : public Aws::Auth::AWSCredentialsProvider private: const std::shared_ptr source_; + // Set once the client's keys are installed. + const std::shared_ptr> installed_; const std::string prefix_; const std::unordered_map settings_; std::mutex mutex_; @@ -626,13 +636,15 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { const DelegatesByPrefix& by_prefix, const ExpiriesByPrefix& expiries, std::string_view location); - /// \brief Build a delegate for each credential this FileIO can serve. + /// \brief Build a delegate for each credential this FileIO can serve. Their + /// clients take this FileIO's keys once `installed` is set. /// /// Runs without holding `mutex_`: building an S3 client can block (without /// static keys the AWS SDK may wait on the EC2 metadata service), which would /// stall every concurrent operation. Reads no mutable member state. Result BuildDelegates( - const std::vector& storage_credentials) const; + const std::vector& storage_credentials, + const std::shared_ptr>& installed) const; /// \brief Swap in credentials and, unless null, delegates, handing back the /// retired ones. A refresh keeps the layout, so the delegates stay. @@ -687,23 +699,27 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { Status ArrowS3FileIO::InitializeStorageCredentials( const std::vector& storage_credentials, std::shared_ptr provider) { - ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials)); + auto installed = std::make_shared>(false); + ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials, installed)); auto credentials = storage_credentials; { std::unique_lock lock(mutex_); provider_.swap(provider); InstallCredentials(credentials, &delegates); + installed->store(true); } return {}; } Status ArrowS3FileIO::SetStorageCredentials( const std::vector& storage_credentials) { - ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials)); + auto installed = std::make_shared>(false); + ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials, installed)); auto credentials = storage_credentials; { std::unique_lock lock(mutex_); InstallCredentials(credentials, &delegates); + installed->store(true); } // `credentials` and `delegates` now hold the retired generation and destruct // here, outside the lock. @@ -711,7 +727,8 @@ Status ArrowS3FileIO::SetStorageCredentials( } Result ArrowS3FileIO::BuildDelegates( - const std::vector& storage_credentials) const { + const std::vector& storage_credentials, + const std::shared_ptr>& installed) const { DelegatesByPrefix delegates; delegates.reserve(storage_credentials.size()); for (const auto& credential : storage_credentials) { @@ -727,8 +744,9 @@ Result ArrowS3FileIO::BuildDelegates( std::shared_ptr keys; // Signs with what ConfigureS3Options would, defaults included. if (HasKeys(properties)) { - keys = std::make_shared( - key_source_, prefix, SettingsOf(properties), AwsCredentialsOf(properties)); + keys = std::make_shared(key_source_, installed, prefix, + SettingsOf(properties), + AwsCredentialsOf(properties)); } ICEBERG_ASSIGN_OR_RAISE(auto fs, BuildArrowS3FileSystem(properties, std::move(keys))); delegates.emplace_back(std::move(prefix), diff --git a/src/iceberg/test/arrow_s3_file_io_test.cc b/src/iceberg/test/arrow_s3_file_io_test.cc index ef73622e36..e9b8e23677 100644 --- a/src/iceberg/test/arrow_s3_file_io_test.cc +++ b/src/iceberg/test/arrow_s3_file_io_test.cc @@ -1146,6 +1146,45 @@ TEST_F(ArrowS3FileIOTest, RefreshReachesHandlesOpenedBeforeIt) { } } +TEST_F(ArrowS3FileIOTest, HandleKeepsItsKeysOnceTheFileIOIsGone) { + if (!HasIntegrationEnv()) { + GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test"; + } + + auto properties = PropertiesFromEnv(); + const auto access_key = properties.find(std::string(S3Properties::kAccessKeyId)); + const auto secret_key = properties.find(std::string(S3Properties::kSecretAccessKey)); + if (access_key == properties.end() || secret_key == properties.end()) { + GTEST_SKIP() << "Set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY to enable " + "credential install test"; + } + + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO(properties)); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + const auto prefix = ObjectUri("install_then_release/"); + const auto uri = prefix + "object"; + constexpr std::string_view kContent = "read after its FileIO is gone"; + ASSERT_THAT(io->WriteFile(uri, kContent), IsOk()); + + // Working keys replace bad ones with the same settings. Building their client + // must not leave it with the bad keys once the FileIO cannot supply new ones. + ASSERT_THAT(credentialed->SetStorageCredentials( + {{.prefix = prefix, .config = BadS3Credentials()}}), + IsOk()); + ASSERT_THAT(credentialed->SetStorageCredentials( + {{.prefix = prefix, .config = {*access_key, *secret_key}}}), + IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto file, io->NewInputFile(uri, kContent.size())); + io.reset(); + + EXPECT_THAT(ReadAll(*file, kContent.size()), + HasValue(::testing::Eq(std::string(kContent)))); + ICEBERG_UNWRAP_OR_FAIL(auto cleanup, MakeS3FileIO(properties)); + EXPECT_THAT(cleanup->DeleteFile(uri), IsOk()); +} + TEST_F(ArrowS3FileIOTest, RequestWaitsForARefreshUnderWayOnceExpired) { if (!HasIntegrationEnv()) { GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test"; From 4e6578af9c37e32ae708331e123f0f1dbadfa1aa Mon Sep 17 00:00:00 2001 From: Gang Wu Date: Sat, 17 Oct 2026 23:47:40 +0800 Subject: [PATCH 5/7] feat(rest,s3): simplify vended credential refresh Manage vended credentials independently by storage prefix, preserve scheme-specific routing, and simplify credential resolution and refresh state.\n\nApply REST client connection and request timeouts consistently, and report transport failures during table writes as unknown commit states. --- src/iceberg/arrow/s3/arrow_s3_file_io.cc | 804 ++++++++---------- src/iceberg/arrow/s3/s3_properties.h | 8 +- src/iceberg/catalog/rest/catalog_properties.h | 5 + src/iceberg/catalog/rest/http_client.cc | 48 +- src/iceberg/catalog/rest/http_client.h | 14 +- src/iceberg/catalog/rest/rest_catalog.cc | 64 +- src/iceberg/catalog/rest/rest_table_scan.cc | 3 +- src/iceberg/test/arrow_s3_file_io_test.cc | 43 +- src/iceberg/test/auth_manager_test.cc | 4 +- src/iceberg/test/http_client_test.cc | 17 +- src/iceberg/test/mock_io.h | 6 + src/iceberg/test/rest_arrow_file_io_test.cc | 11 +- .../test/rest_catalog_integration_test.cc | 4 +- src/iceberg/test/rest_catalog_unit_test.cc | 42 +- src/iceberg/test/rest_table_scan_test.cc | 24 +- 15 files changed, 535 insertions(+), 562 deletions(-) diff --git a/src/iceberg/arrow/s3/arrow_s3_file_io.cc b/src/iceberg/arrow/s3/arrow_s3_file_io.cc index 3bdf36d999..a38e864215 100644 --- a/src/iceberg/arrow/s3/arrow_s3_file_io.cc +++ b/src/iceberg/arrow/s3/arrow_s3_file_io.cc @@ -56,6 +56,7 @@ #include "iceberg/util/macros.h" #include "iceberg/util/property_util.h" #include "iceberg/util/string_util.h" +#include "iceberg/util/timepoint.h" namespace iceberg::arrow { @@ -299,8 +300,8 @@ Result<::arrow::fs::S3Options> ConfigureS3Options( namespace { -// `credentials_provider`, if set, signs requests instead of the keys in -// `properties`. +// Build an Arrow S3 filesystem. If a provider is supplied, the AWS signer gets +// access keys from it instead of using the keys in `properties`. Result> BuildArrowS3FileSystem( const std::unordered_map& properties, std::shared_ptr credentials_provider = nullptr) { @@ -314,256 +315,397 @@ Result> BuildArrowS3FileSystem( return std::shared_ptr<::arrow::fs::FileSystem>(std::move(fs)); } -// Rewrites any alias of `s3://` (any case — routing is case-insensitive) to -// exactly that, so locations and credential prefixes compare equal. An alias -// missing from kS3Schemes would silently stop matching its credential. -std::string CanonicalizeS3Scheme(std::string_view location) { +// Make the URI scheme lower-case for matching. Keep different schemes, such as +// `s3` and `oss`, separate because they can use different endpoints and keys. +std::string NormalizeS3Location(std::string_view location) { const auto separator = location.find("://"); if (separator != std::string_view::npos && IsS3Scheme(location.substr(0, separator))) { - return std::string("s3://").append(location.substr(separator + 3)); + // The scheme matched an ASCII alias, so only ASCII letters change. + return StringUtils::ToLower(location.substr(0, separator)) + .append(location.substr(separator)); } return std::string(location); } -// Lead time before expiry, matching Java's VendedCredentialsProvider. +// Start refreshing five minutes before the earliest token expires. constexpr auto kRefreshLeadTime = std::chrono::minutes(5); -// How long to wait after a failed refresh before asking the catalog again. -constexpr auto kRefreshRetryBackoff = std::chrono::seconds(30); +// Wait this long before retrying a failed refresh. +constexpr std::chrono::milliseconds kRefreshRetryBackoff = std::chrono::seconds(30); -// Floor on a backoff shortened to land on the expiry. -constexpr auto kMinRefreshRetryBackoff = std::chrono::seconds(1); +// Never use a retry delay shorter than this value. +constexpr std::chrono::milliseconds kMinRefreshRetryBackoff = std::chrono::seconds(1); -// How long a caller with expired credentials waits for a refresh under way. +// An expired request waits this long for another request's refresh to finish. constexpr auto kExpiredCredentialWait = std::chrono::seconds(10); -// Whether `config` carries an access key pair. +// Return true when the config has both parts of a static key pair. bool HasKeys(const std::unordered_map& config) { return FindProperty(config, S3Properties::kAccessKeyId) != nullptr && FindProperty(config, S3Properties::kSecretAccessKey) != nullptr; } -// The properties an S3 client for `credential` is built from. -std::unordered_map PropertiesFor( +// Convert a complete static key pair to the AWS SDK credential object. +Aws::Auth::AWSCredentials AwsCredentialsOf( + const std::unordered_map& config) { + Aws::Auth::AWSCredentials credentials; + credentials.SetAWSAccessKeyId( + FindProperty(config, S3Properties::kAccessKeyId)->c_str()); + credentials.SetAWSSecretKey( + FindProperty(config, S3Properties::kSecretAccessKey)->c_str()); + if (const auto* session_token = FindProperty(config, S3Properties::kSessionToken); + session_token != nullptr) { + credentials.SetSessionToken(session_token->c_str()); + } + return credentials; +} + +struct ResolvedCredential { + std::string prefix; + std::unordered_map properties; + // These settings identify the client configuration independently of its + // short-lived credentials. They protect a client from receiving keys meant + // for a different endpoint, region, or addressing mode. + std::unordered_map client_settings; + std::optional credentials; +}; + +// Merge FileIO defaults with one vended credential and derive all values used +// by the client and refresh store in one pass. +ResolvedCredential ResolveCredential( const std::unordered_map& default_properties, const StorageCredential& credential) { - auto properties = default_properties; + ResolvedCredential resolved; + resolved.prefix = NormalizeS3Location(credential.prefix); + resolved.properties = default_properties; for (const auto& [key, value] : credential.config) { - properties[key] = value; + resolved.properties[key] = value; } - return properties; -} - -// The client settings in `properties`: everything but the keys. -std::unordered_map SettingsOf( - std::unordered_map properties) { + resolved.client_settings = resolved.properties; for (const auto key : {S3Properties::kAccessKeyId, S3Properties::kSecretAccessKey, S3Properties::kSessionToken, S3Properties::kSessionTokenExpiresAtMs}) { - properties.erase(std::string(key)); + resolved.client_settings.erase(std::string(key)); + } + if (HasKeys(resolved.properties)) { + resolved.credentials = AwsCredentialsOf(resolved.properties); } - return properties; + return resolved; } -// When the session token that `properties` sign with expires: nullopt if they -// sign without one, or an error if the expiry is missing or invalid. -Result> ExpiryOf( +// Read the expiry of the token used by these properties. No session token +// means there is no expiry to track; a bad token expiry is an error. +Result> ParseCredentialExpiry( const std::unordered_map& properties, std::string_view prefix) { - // Only a nonempty token alongside an access key pair signs requests. + // Only a non-empty session token paired with access keys needs an expiry. const auto* token = FindProperty(properties, S3Properties::kSessionToken); if (!HasKeys(properties) || token == nullptr || token->empty()) { return std::nullopt; } const auto* value = FindProperty(properties, S3Properties::kSessionTokenExpiresAtMs); if (value == nullptr) { - return InvalidArgument("Credential \"{}\" has a session token but no \"{}\"", prefix, + return InvalidArgument(R"(Credential "{}" has a session token but no "{}")", prefix, S3Properties::kSessionTokenExpiresAtMs); } auto millis = StringUtils::ParseNumber(*value); if (!millis.has_value()) { return InvalidArgument( - "Credential \"{}\" has a session token but an unparseable \"{}\" value \"{}\"", + R"(Credential "{}" has a session token but an unparseable "{}" value "{}")", prefix, S3Properties::kSessionTokenExpiresAtMs, *value); } - // 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) { - return InvalidArgument("Credential \"{}\" has an out-of-range \"{}\" value \"{}\"", - prefix, S3Properties::kSessionTokenExpiresAtMs, *value); - } - return std::chrono::system_clock::time_point(std::chrono::milliseconds(*millis)); + return TimePointMsFromUnixMs(*millis); } -// The S3 credentials in effect: the first for each canonical prefix, which is -// the one its locations match. +// Keep the first credential for each prefix. This is also the order used when +// a duplicate prefix is returned by the catalog. std::vector S3CredentialsOf( const std::vector& credentials) { std::vector s3_credentials; std::unordered_set prefixes; for (const auto& credential : credentials) { if (IsS3CredentialPrefix(credential.prefix) && - prefixes.insert(CanonicalizeS3Scheme(credential.prefix)).second) { + prefixes.insert(NormalizeS3Location(credential.prefix)).second) { s3_credentials.push_back(&credential); } } return s3_credentials; } -// When installed credentials need replacing. -struct CredentialsExpiry { - // The earliest expiry known, if any. - std::optional expires_at; - // A session token lacks a usable expiry: refresh at once, as Java does, but - // do not treat it as expired. - bool unknown = false; +struct CredentialSlot { + std::string prefix; + std::unordered_map settings; + std::optional credentials; + std::optional expires_at; + bool expiry_unknown = false; + + bool ExpiresBy(TimePointMs time) const { + return expires_at.has_value() && time >= *expires_at; + } }; -CredentialsExpiry ExpiryOf( - const std::unordered_map& default_properties, - const std::vector& credentials) { - CredentialsExpiry expiry; - for (const auto* credential : S3CredentialsOf(credentials)) { - auto expires_at = - ExpiryOf(PropertiesFor(default_properties, *credential), credential->prefix); - if (!expires_at.has_value()) { - ICEBERG_LOG_WARN("{}", expires_at.error().message); - expiry.unknown = true; - } else if (expires_at->has_value() && - (!expiry.expires_at.has_value() || **expires_at < *expiry.expires_at)) { - expiry.expires_at = **expires_at; - } +class CredentialStore { + public: + explicit CredentialStore( + std::unordered_map default_properties) + : default_properties_(std::move(default_properties)) {} + + ResolvedCredential Resolve(const StorageCredential& credential) const { + return ResolveCredential(default_properties_, credential); } - return expiry; -} -// Known expiries by canonical prefix. -std::unordered_map ExpiriesOf( - const std::unordered_map& default_properties, - const std::vector& credentials) { - std::unordered_map expiries; - for (const auto* credential : S3CredentialsOf(credentials)) { - // Invalid expiries are logged on install. - if (auto expires_at = - ExpiryOf(PropertiesFor(default_properties, *credential), credential->prefix); - expires_at.has_value() && expires_at->has_value()) { - expiries.emplace(CanonicalizeS3Scheme(credential->prefix), **expires_at); + /// \brief Make `credentials` current, and `provider` too when one is given. + void Install(std::vector credentials, + std::optional> provider) { + std::unique_lock lock(mutex_); + if (provider.has_value()) { + provider_ = std::move(*provider); } + InstallLocked(std::move(credentials)); } - return expiries; -} - -// The S3 prefixes `credentials` serve, each with its client settings and -// whether it signs with keys. -using Layout = - std::unordered_map, bool>>; - -Layout LayoutOf(const std::unordered_map& default_properties, - const std::vector& credentials) { - Layout layout; - for (const auto* credential : S3CredentialsOf(credentials)) { - auto properties = PropertiesFor(default_properties, *credential); - const bool keys = HasKeys(properties); - layout.try_emplace(CanonicalizeS3Scheme(credential->prefix), - SettingsOf(std::move(properties)), keys); - } - return layout; -} -// A refresh may only rotate keys, so the current clients serve it: it must keep -// the layout of `current`, and its tokens must be usable. -Status ValidateRefreshed( - const std::unordered_map& default_properties, - const std::vector& current, - const std::vector& credentials) { - for (const auto& credential : credentials) { - ICEBERG_RETURN_UNEXPECTED(credential.Validate()); + std::shared_ptr SlotFor( + std::string prefix, std::unordered_map settings, + std::optional initial_credentials) { + std::unique_lock lock(mutex_); + return SlotForLocked(std::move(prefix), std::move(settings), + std::move(initial_credentials)); } - if (LayoutOf(default_properties, credentials) != - LayoutOf(default_properties, current)) { - return InvalidArgument( - "Refreshed credentials change the S3 prefixes, their settings or keys"); - } - const auto now = std::chrono::system_clock::now(); - for (const auto* credential : S3CredentialsOf(credentials)) { - ICEBERG_ASSIGN_OR_RAISE( - auto expires_at, - ExpiryOf(PropertiesFor(default_properties, *credential), credential->prefix)); - if (expires_at.has_value() && *expires_at <= now) { - return InvalidArgument("Credential \"{}\" has already expired", credential->prefix); - } + + std::vector credentials() const { + std::shared_lock lock(mutex_); + return storage_credentials_; } - return {}; -} -// Requires HasKeys(config). -Aws::Auth::AWSCredentials AwsCredentialsOf( - const std::unordered_map& config) { - Aws::Auth::AWSCredentials credentials; - credentials.SetAWSAccessKeyId( - FindProperty(config, S3Properties::kAccessKeyId)->c_str()); - credentials.SetAWSSecretKey( - FindProperty(config, S3Properties::kSecretAccessKey)->c_str()); - if (const auto* session_token = FindProperty(config, S3Properties::kSessionToken); - session_token != nullptr) { - credentials.SetSessionToken(session_token->c_str()); + bool IsExpired(const std::string& prefix) const { + std::shared_lock lock(mutex_); + const auto current = current_slots_.find(prefix); + return current != current_slots_.end() && + current->second->ExpiresBy(CurrentTimePointMs()); } - return credentials; -} -class ArrowS3FileIO; + void MaybeRefreshCredentials() { + // Only one caller loads from the catalog. Other callers keep using the + // current keys while they are still valid. + { + std::shared_lock lock(mutex_); + if (!provider_ || !RefreshDue()) { + return; + } + } -// Gives S3 clients their FileIO's current keys. Clients may outlive the -// FileIO, which detaches on destruction. -class KeySource { - public: - explicit KeySource(ArrowS3FileIO* io) : io_(io) {} + std::unique_lock refresh_lock(refresh_mutex_, std::defer_lock); + if (!refresh_lock.try_lock()) { + { + std::shared_lock lock(mutex_); + if (!Expired()) { + return; + } + } + if (!refresh_lock.try_lock_for(kExpiredCredentialWait)) { + return; + } + } + + std::shared_ptr provider; + uint64_t generation = 0; + { + std::shared_lock lock(mutex_); + if (!provider_ || !RefreshDue()) { + return; + } + provider = provider_; + generation = credential_generation_; + } + + auto refreshed = provider->Load(); - void Detach() { std::unique_lock lock(mutex_); - io_ = nullptr; + if (credential_generation_ != generation) { + return; // An install replaced the credentials this refresh started from. + } + const Status status = refreshed.has_value() + ? ApplyRefreshLocked(*refreshed) + : Status(::iceberg::unexpected(refreshed.error())); + if (status.has_value() && !RefreshDue()) { + return; + } + const auto delay = BackoffDelay(); + retry_refresh_at_ = std::chrono::steady_clock::now() + delay; + if (!status.has_value()) { + ICEBERG_LOG_WARN( + "Failed to refresh vended storage credentials ({}); keeping the current " + "ones and retrying in {}ms", + status.error().message, delay.count()); + } } - /// \brief Refreshes if due, then returns the FileIO's keys for `prefix` and - /// `settings`; nullopt if it has none or is gone. - std::optional KeysFor( - std::string_view prefix, - const std::unordered_map& settings); + std::optional CredentialsFor(const CredentialSlot& slot) { + MaybeRefreshCredentials(); + std::shared_lock lock(mutex_); + return slot.credentials; + } private: - std::shared_mutex mutex_; - ArrowS3FileIO* io_; + void InstallLocked(std::vector credentials) { + std::unordered_map> current_slots; + for (const auto* credential : S3CredentialsOf(credentials)) { + auto resolved = Resolve(*credential); + auto parsed_expiry = ParseCredentialExpiry(resolved.properties, credential->prefix); + if (!parsed_expiry.has_value()) { + ICEBERG_LOG_WARN("{}", parsed_expiry.error().message); + } + auto slot = SlotForLocked(std::move(resolved.prefix), + std::move(resolved.client_settings), std::nullopt); + // Clients from an earlier install may share this slot; they move to these + // keys too. + slot->credentials = std::move(resolved.credentials); + slot->expires_at = parsed_expiry.value_or(std::nullopt); + slot->expiry_unknown = !parsed_expiry.has_value(); + current_slots.emplace(slot->prefix, std::move(slot)); + } + retry_refresh_at_ = {}; + ++credential_generation_; + current_slots_ = std::move(current_slots); + storage_credentials_ = std::move(credentials); + } + + Status ApplyRefreshLocked(const std::vector& credentials) { + // Check every credential, skipped ones included, before changing any slot. + for (const auto& credential : credentials) { + ICEBERG_RETURN_UNEXPECTED(credential.Validate()); + } + bool updated = false; + for (const auto* credential : S3CredentialsOf(credentials)) { + auto resolved = Resolve(*credential); + const auto& prefix = resolved.prefix; + const auto current = current_slots_.find(prefix); + if (current == current_slots_.end()) { + continue; + } + const auto& settings = resolved.client_settings; + const bool has_keys = resolved.credentials.has_value(); + auto& slot = *current->second; + if (slot.settings != settings || slot.credentials.has_value() != has_keys) { + ICEBERG_LOG_WARN( + "Ignoring refreshed credentials for '{}' because their client settings " + "changed", + prefix); + continue; + } + auto parsed_expiry = ParseCredentialExpiry(resolved.properties, credential->prefix); + if (!parsed_expiry.has_value()) { + ICEBERG_LOG_WARN("Ignoring refreshed credentials for '{}': {}", prefix, + parsed_expiry.error().message); + continue; + } + if (parsed_expiry->has_value() && **parsed_expiry <= CurrentTimePointMs()) { + ICEBERG_LOG_WARN("Ignoring already expired refreshed credentials for '{}'", + prefix); + continue; + } + slot.credentials = std::move(resolved.credentials); + slot.expires_at = *parsed_expiry; + slot.expiry_unknown = false; + for (auto& installed : storage_credentials_) { + if (NormalizeS3Location(installed.prefix) == prefix) { + installed = *credential; + break; + } + } + updated = true; + } + if (!updated && !current_slots_.empty()) { + return InvalidArgument("Refreshed credentials did not update any active prefix"); + } + ++credential_generation_; + return {}; + } + + std::shared_ptr SlotForLocked( + std::string prefix, std::unordered_map settings, + std::optional initial_credentials) { + for (const auto& slot : known_slots_) { + if (slot->prefix == prefix && slot->settings == settings) { + return slot; + } + } + auto slot = std::make_shared( + CredentialSlot{.prefix = std::move(prefix), + .settings = std::move(settings), + .credentials = std::move(initial_credentials)}); + known_slots_.push_back(slot); + return slot; + } + + bool RefreshDue() const { + if (std::chrono::steady_clock::now() < retry_refresh_at_) { + return false; + } + const auto refresh_by = CurrentTimePointMs() + kRefreshLeadTime; + return std::ranges::any_of(current_slots_, [refresh_by](const auto& entry) { + return entry.second->expiry_unknown || entry.second->ExpiresBy(refresh_by); + }); + } + + bool Expired() const { + const auto now = CurrentTimePointMs(); + return std::ranges::any_of(current_slots_, [now](const auto& entry) { + return entry.second->ExpiresBy(now); + }); + } + + // Retry after kRefreshRetryBackoff, or sooner if a token expires first, but + // never sooner than kMinRefreshRetryBackoff. + std::chrono::milliseconds BackoffDelay() const { + std::optional earliest; + for (const auto& [_, slot] : current_slots_) { + if (slot->expires_at.has_value() && + (!earliest.has_value() || *slot->expires_at < *earliest)) { + earliest = slot->expires_at; + } + } + const auto now = CurrentTimePointMs(); + if (earliest.has_value() && *earliest > now) { + return std::clamp( + *earliest - now, kMinRefreshRetryBackoff, kRefreshRetryBackoff); + } + return kRefreshRetryBackoff; + } + + const std::unordered_map default_properties_; + mutable std::shared_mutex mutex_; + std::vector storage_credentials_; + std::vector> known_slots_; + std::unordered_map> current_slots_; + std::shared_ptr provider_; + std::chrono::steady_clock::time_point retry_refresh_at_; + uint64_t credential_generation_ = 0; + std::timed_mutex refresh_mutex_; }; -// Supplies an S3 client's keys. The AWS SDK asks before each request, so open -// handles pick up refreshed keys, as in Java. Keys vended with other settings -// (e.g. another endpoint) are never used; the last keys stay instead. +// Supplies one S3 client's keys from its slot in the shared store, which lets +// clients and handles outlive the FileIO. The last keys seen stay in use while +// the slot has none, as after the prefix is reinstalled without keys. class VendedCredentialsProvider final : public Aws::Auth::AWSCredentialsProvider { public: - VendedCredentialsProvider(std::shared_ptr source, + VendedCredentialsProvider(std::shared_ptr store, std::shared_ptr> installed, - std::string prefix, - std::unordered_map settings, + std::shared_ptr slot, Aws::Auth::AWSCredentials credentials) - : source_(std::move(source)), + : store_(std::move(store)), installed_(std::move(installed)), - prefix_(std::move(prefix)), - settings_(std::move(settings)), + slot_(std::move(slot)), credentials_(std::move(credentials)) {} Aws::Auth::AWSCredentials GetAWSCredentials() override { - std::optional current; - // The SDK also asks while building the client, when the FileIO still holds - // the keys its install replaces. - if (installed_->load()) { - current = source_->KeysFor(prefix_, settings_); + // The SDK also asks while building the client. Until its install, the slot + // may hold the keys being replaced, and asking could start their refresh. + if (!installed_->load()) { + std::lock_guard lock(mutex_); + return credentials_; } + auto current = store_->CredentialsFor(*slot_); std::lock_guard lock(mutex_); if (current.has_value()) { credentials_ = std::move(current).value(); @@ -572,11 +714,10 @@ class VendedCredentialsProvider final : public Aws::Auth::AWSCredentialsProvider } private: - const std::shared_ptr source_; - // Set once the client's keys are installed. + const std::shared_ptr store_; + // Set once the install that built this client has made its keys current. const std::shared_ptr> installed_; - const std::string prefix_; - const std::unordered_map settings_; + const std::shared_ptr slot_; std::mutex mutex_; Aws::Auth::AWSCredentials credentials_; }; @@ -587,10 +728,9 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { std::unordered_map default_properties, size_t delete_parallelism) : default_file_io_(std::make_shared(std::move(arrow_fs))), - default_properties_(std::move(default_properties)), - delete_parallelism_(delete_parallelism) {} - - ~ArrowS3FileIO() override { key_source_->Detach(); } + delete_parallelism_(delete_parallelism), + credential_store_( + std::make_shared(std::move(default_properties))) {} Result> NewInputFile(std::string file_location) override; @@ -607,8 +747,7 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { const std::vector& storage_credentials) override; std::vector credentials() const override { - std::shared_lock lock(mutex_); - return storage_credentials_; + return credential_store_->credentials(); } Status InitializeStorageCredentials( @@ -618,26 +757,22 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; } private: - friend class KeySource; - using DelegatesByPrefix = std::vector>>; - using ExpiriesByPrefix = - std::unordered_map; - /// \brief Delegate serving `location`, pinned by the caller against a - /// concurrent credential install. + /// \brief Delegate with the longest prefix of `location`, pinned by the caller + /// against a concurrent credential install. Fails with TokenExpired once that + /// prefix's keys have expired. Result> FileIOForPath(std::string_view location); - /// \brief Longest-prefix match, or `fallback` if none matches. Fails with - /// TokenExpired once the matched prefix's keys have expired. - static Result> MatchDelegate( - const std::shared_ptr& fallback, - const DelegatesByPrefix& by_prefix, const ExpiriesByPrefix& expiries, - std::string_view location); + /// \brief Build delegates for `storage_credentials` and make them current, + /// replacing the refresh provider too when one is given. + Status InstallStorageCredentials( + const std::vector& storage_credentials, + std::optional> provider); /// \brief Build a delegate for each credential this FileIO can serve. Their - /// clients take this FileIO's keys once `installed` is set. + /// clients take keys from the store once `installed` is set. /// /// Runs without holding `mutex_`: building an S3 client can block (without /// static keys the AWS SDK may wait on the EC2 metadata service), which would @@ -646,83 +781,39 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials { const std::vector& storage_credentials, const std::shared_ptr>& installed) const; - /// \brief Swap in credentials and, unless null, delegates, handing back the - /// retired ones. A refresh keeps the layout, so the delegates stay. - /// - /// Callers must hold `mutex_` exclusively and let the returned generation - /// destruct only after releasing it: tearing down an S3 client can block on - /// in-flight requests, which would stall every operation. - void InstallCredentials(std::vector& storage_credentials, - DelegatesByPrefix* delegates); - - /// \brief Whether a refresh is due and not backed off. Requires `mutex_`. - bool RefreshDue() const; - - /// \brief Whether the installed credentials have expired. Requires `mutex_`. - bool Expired() const; - - /// \brief When to retry after a failed refresh; never past the expiry. - /// Requires `mutex_`. - std::chrono::steady_clock::time_point BackoffUntil() const; - - /// \brief Refresh the credentials if they are close to expiring. - /// - /// Called before file operations and by S3 clients before each request. A - /// refresh only rotates keys and builds no clients, so a request can run it - /// or, with expired credentials, wait for one under way. A failure keeps the - /// current credentials. - void MaybeRefreshCredentials(); - - /// \brief The installed keys for `prefix`, if its credential has `settings`. - std::optional KeysFor( - std::string_view prefix, - const std::unordered_map& settings) const; - std::shared_ptr default_file_io_; - std::unordered_map default_properties_; const size_t delete_parallelism_; - const std::shared_ptr key_source_ = std::make_shared(this); + const std::shared_ptr credential_store_; // Guards everything below; shared because reads happen per file operation. mutable std::shared_mutex mutex_; - std::vector storage_credentials_; DelegatesByPrefix file_io_by_prefix_; - ExpiriesByPrefix expiries_by_prefix_; - std::shared_ptr provider_; - CredentialsExpiry expiry_; - std::chrono::steady_clock::time_point retry_refresh_at_; - // Bumped on each install so a refresh cannot overwrite newer credentials. - uint64_t credential_generation_ = 0; - // Serializes refreshes; timed so waits are bounded. - std::timed_mutex refresh_mutex_; }; Status ArrowS3FileIO::InitializeStorageCredentials( const std::vector& storage_credentials, std::shared_ptr provider) { - auto installed = std::make_shared>(false); - ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials, installed)); - auto credentials = storage_credentials; - { - std::unique_lock lock(mutex_); - provider_.swap(provider); - InstallCredentials(credentials, &delegates); - installed->store(true); - } - return {}; + return InstallStorageCredentials(storage_credentials, std::move(provider)); } Status ArrowS3FileIO::SetStorageCredentials( const std::vector& storage_credentials) { + return InstallStorageCredentials(storage_credentials, std::nullopt); +} + +Status ArrowS3FileIO::InstallStorageCredentials( + const std::vector& storage_credentials, + std::optional> provider) { auto installed = std::make_shared>(false); ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials, installed)); auto credentials = storage_credentials; { std::unique_lock lock(mutex_); - InstallCredentials(credentials, &delegates); + file_io_by_prefix_.swap(delegates); + credential_store_->Install(std::move(credentials), std::move(provider)); installed->store(true); } - // `credentials` and `delegates` now hold the retired generation and destruct - // here, outside the lock. + // `delegates` now holds the retired clients and destructs here, outside the + // lock: tearing down an S3 client can block on its in-flight requests. return {}; } @@ -733,28 +824,23 @@ Result ArrowS3FileIO::BuildDelegates( 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; - // non-S3 prefixes are skipped, not rejected (Java S3FileIO filters - // credentials by the "s3" prefix). - if (!IsS3CredentialPrefix(credential.prefix)) { - continue; - } - auto prefix = CanonicalizeS3Scheme(credential.prefix); - const auto properties = PropertiesFor(default_properties_, credential); + } + for (const auto* credential : S3CredentialsOf(storage_credentials)) { + auto resolved = credential_store_->Resolve(*credential); std::shared_ptr keys; - // Signs with what ConfigureS3Options would, defaults included. - if (HasKeys(properties)) { - keys = std::make_shared(key_source_, installed, prefix, - SettingsOf(properties), - AwsCredentialsOf(properties)); + if (resolved.credentials.has_value()) { + auto slot = credential_store_->SlotFor( + resolved.prefix, std::move(resolved.client_settings), resolved.credentials); + keys = std::make_shared( + credential_store_, installed, std::move(slot), + *std::move(resolved.credentials)); } - ICEBERG_ASSIGN_OR_RAISE(auto fs, BuildArrowS3FileSystem(properties, std::move(keys))); - delegates.emplace_back(std::move(prefix), + ICEBERG_ASSIGN_OR_RAISE(auto fs, + BuildArrowS3FileSystem(resolved.properties, std::move(keys))); + delegates.emplace_back(std::move(resolved.prefix), std::make_shared(std::move(fs))); } 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( "None of the {} vended storage credential(s) has an S3-compatible prefix; " "S3 access will use the default credentials", @@ -763,178 +849,31 @@ Result ArrowS3FileIO::BuildDelegates( return delegates; } -void ArrowS3FileIO::InstallCredentials( - std::vector& storage_credentials, DelegatesByPrefix* delegates) { - if (delegates != nullptr) { - file_io_by_prefix_.swap(*delegates); - } - expiry_ = ExpiryOf(default_properties_, storage_credentials); - expiries_by_prefix_ = ExpiriesOf(default_properties_, storage_credentials); - retry_refresh_at_ = {}; - ++credential_generation_; - storage_credentials_.swap(storage_credentials); -} - -bool ArrowS3FileIO::RefreshDue() const { - return (expiry_.unknown || - (expiry_.expires_at.has_value() && - std::chrono::system_clock::now() + kRefreshLeadTime >= *expiry_.expires_at)) && - std::chrono::steady_clock::now() >= retry_refresh_at_; -} - -bool ArrowS3FileIO::Expired() const { - return expiry_.expires_at.has_value() && - std::chrono::system_clock::now() >= *expiry_.expires_at; -} - -std::chrono::steady_clock::time_point ArrowS3FileIO::BackoffUntil() const { - auto delay = - std::chrono::duration_cast(kRefreshRetryBackoff); - if (expiry_.expires_at.has_value()) { - const auto remaining = std::chrono::duration_cast( - *expiry_.expires_at - std::chrono::system_clock::now()); - // Retry sooner while still valid; once expired, use the full backoff. - 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() { - { - // Fast path. - std::shared_lock lock(mutex_); - if (!provider_ || !RefreshDue()) { - return; - } - } - - std::unique_lock refresh_lock(refresh_mutex_, std::defer_lock); - if (!refresh_lock.try_lock()) { - // Another caller is refreshing; keep the current keys unless expired. - { - std::shared_lock lock(mutex_); - if (!Expired()) { - return; - } - } - if (!refresh_lock.try_lock_for(kExpiredCredentialWait)) { - return; - } - } - std::shared_ptr provider; - uint64_t generation = 0; - std::vector current; - { - // The previous holder may have just refreshed. - std::shared_lock lock(mutex_); - if (!provider_ || !RefreshDue()) { - return; - } - provider = provider_; - generation = credential_generation_; - current = storage_credentials_; - } - - // Slow, so done without holding `mutex_`. - auto refreshed = provider->Load(); - Status status; - if (refreshed.has_value()) { - status = ValidateRefreshed(default_properties_, current, *refreshed); - } else { - status = ::iceberg::unexpected(refreshed.error()); - } - - std::unique_lock lock(mutex_); - // Credentials installed meanwhile are newer than what this fetched. - const bool superseded = credential_generation_ != generation; - if (!status.has_value()) { - // Logged even if superseded, so catalog failures stay 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, /*delegates=*/nullptr); - if (RefreshDue()) { - // Tokens shorter-lived than the lead time would be due again at once. - retry_refresh_at_ = BackoffUntil(); - } -} - -std::optional ArrowS3FileIO::KeysFor( - std::string_view prefix, - const std::unordered_map& settings) const { - std::shared_lock lock(mutex_); - for (const auto* credential : S3CredentialsOf(storage_credentials_)) { - if (CanonicalizeS3Scheme(credential->prefix) != prefix) { - continue; - } - const auto properties = PropertiesFor(default_properties_, *credential); - if (!HasKeys(properties) || SettingsOf(properties) != settings) { - return std::nullopt; - } - return AwsCredentialsOf(properties); - } - return std::nullopt; -} - -Result> ArrowS3FileIO::MatchDelegate( - const std::shared_ptr& fallback, - const DelegatesByPrefix& by_prefix, const ExpiriesByPrefix& expiries, +Result> ArrowS3FileIO::FileIOForPath( std::string_view location) { - if (by_prefix.empty()) { - return fallback; - } - const std::string canonical = CanonicalizeS3Scheme(location); - auto best = fallback; - size_t best_len = 0; - const std::string* best_prefix = nullptr; - for (const auto& [prefix, file_io] : by_prefix) { - if (prefix.size() > best_len && canonical.starts_with(prefix)) { - best = file_io; - best_len = prefix.size(); - best_prefix = &prefix; + credential_store_->MaybeRefreshCredentials(); + std::shared_lock lock(mutex_); + if (file_io_by_prefix_.empty()) { + return default_file_io_; + } + const std::string normalized = NormalizeS3Location(location); + const DelegatesByPrefix::value_type* best = nullptr; + for (const auto& delegate : file_io_by_prefix_) { + const auto& prefix = delegate.first; + if (prefix.size() > (best != nullptr ? best->first.size() : 0) && + normalized.starts_with(prefix)) { + best = &delegate; } } - if (best_prefix == nullptr) { - return fallback; + if (best == nullptr) { + return default_file_io_; } - if (const auto expiry = expiries.find(*best_prefix); - expiry != expiries.end() && std::chrono::system_clock::now() >= expiry->second) { + if (credential_store_->IsExpired(best->first)) { return TokenExpired( "Vended credentials for \"{}\" have expired and could not be refreshed", - *best_prefix); + best->first); } - return best; -} - -Result> ArrowS3FileIO::FileIOForPath( - std::string_view location) { - MaybeRefreshCredentials(); - std::shared_lock lock(mutex_); - return MatchDelegate(default_file_io_, file_io_by_prefix_, expiries_by_prefix_, - location); + return best->second; } Result> ArrowS3FileIO::NewInputFile( @@ -969,17 +908,6 @@ Status ArrowS3FileIO::DeleteFiles(const std::vector& file_locations }); } -std::optional KeySource::KeysFor( - std::string_view prefix, - const std::unordered_map& settings) { - std::shared_lock lock(mutex_); - if (io_ == nullptr) { - return std::nullopt; - } - io_->MaybeRefreshCredentials(); - return io_->KeysFor(prefix, settings); -} - } // namespace Result> MakeS3FileIO( diff --git a/src/iceberg/arrow/s3/s3_properties.h b/src/iceberg/arrow/s3/s3_properties.h index 14cd2b4fdc..79609c878e 100644 --- a/src/iceberg/arrow/s3/s3_properties.h +++ b/src/iceberg/arrow/s3/s3_properties.h @@ -61,9 +61,11 @@ struct S3Properties { static constexpr std::string_view kDeleteNumThreads = "s3.delete.num-threads"; }; -/// \brief URI schemes served by the Arrow S3 FileIO, lower-case. +/// \brief URI schemes served by the Arrow S3 FileIO, lower-case. These schemes +/// share the S3 client implementation but remain distinct when routing +/// credentials. /// -/// Single source of truth: registration, IsS3Scheme and alias canonicalization +/// Single source of truth: registration, IsS3Scheme and alias normalization /// all derive from this list, so a new alias only has to be added here. /// /// `oss` is served because the store is S3-compatible; see the FileIO docs. @@ -91,7 +93,7 @@ inline constexpr bool IsS3Scheme(std::string_view scheme) { /// \brief Return whether a storage credential prefix belongs to S3. /// /// Accepts the bare `s3` prefix (exactly: a case variant would be stored as a -/// key no canonicalized location can match) or a URI prefix such as +/// key no normalized location can match) or a URI prefix such as /// `s3a://bucket`. Bare aliases such as `oss` are deliberately rejected, /// matching Java S3FileIO's credential filter. inline constexpr bool IsS3CredentialPrefix(std::string_view prefix) { diff --git a/src/iceberg/catalog/rest/catalog_properties.h b/src/iceberg/catalog/rest/catalog_properties.h index da5bd550aa..a1f205a605 100644 --- a/src/iceberg/catalog/rest/catalog_properties.h +++ b/src/iceberg/catalog/rest/catalog_properties.h @@ -72,6 +72,11 @@ class ICEBERG_REST_EXPORT RestCatalogProperties "rest-scan-planning.poll-num-retries", 10}; /// \brief The prefix for HTTP headers. inline static constexpr std::string_view kHeaderPrefix = "header."; + /// \brief HTTP connection timeout in milliseconds, matching Java REST clients. + inline static Entry kClientConnectionTimeoutMs{ + "rest.client.connection-timeout-ms", 0}; + /// \brief HTTP socket/read timeout in milliseconds, matching Java REST clients. + inline static Entry kClientSocketTimeoutMs{"rest.client.socket-timeout-ms", 0}; /// \brief Create a default RestCatalogProperties instance. static RestCatalogProperties default_properties(); diff --git a/src/iceberg/catalog/rest/http_client.cc b/src/iceberg/catalog/rest/http_client.cc index 42dac26d8c..871a21e67e 100644 --- a/src/iceberg/catalog/rest/http_client.cc +++ b/src/iceberg/catalog/rest/http_client.cc @@ -146,24 +146,27 @@ ErrorResponse BuildDefaultErrorResponse(const cpr::Response& response) { }; } -/// \brief Fails a GET whose transfer failed, even after a success status. -/// -/// Not for writes: one may have taken effect, and a failed commit's files are -/// cleaned up. -Status CheckGetTransfer(const cpr::Response& response) { +/// \brief Fails a response whose transfer failed, even after a success status. +Status CheckTransfer(const cpr::Response& response) { if (response.error.code == cpr::ErrorCode::OPERATION_TIMEDOUT) { return IOError("HTTP request timed out: {}", response.error.message); } - if (response.error && IsSuccessful(response.status_code)) { - return IOError("HTTP request failed after status {}: {}", response.status_code, - response.error.message); + if (response.error) { + if (IsSuccessful(response.status_code)) { + return IOError("HTTP request failed after status {}: {}", response.status_code, + response.error.message); + } + return IOError("HTTP request failed: {}", response.error.message); } return {}; } /// \brief Handles failure responses by invoking the provided error handler. +/// +/// A failed transfer is checked first, as its status code may be 0 or a success. Status HandleFailureResponse(const cpr::Response& response, const ErrorHandler& error_handler) { + ICEBERG_RETURN_UNEXPECTED(CheckTransfer(response)); if (IsSuccessful(response.status_code)) { return {}; } @@ -176,8 +179,12 @@ Status HandleFailureResponse(const cpr::Response& response, } // namespace -HttpClient::HttpClient(std::unordered_map default_headers) +HttpClient::HttpClient(std::unordered_map default_headers, + std::chrono::milliseconds connection_timeout, + std::chrono::milliseconds socket_timeout) : default_headers_{std::move(default_headers)}, + connection_timeout_{connection_timeout}, + socket_timeout_{socket_timeout}, connection_pool_{std::make_unique()} { // Set default Content-Type for all requests (including GET/HEAD/DELETE). // Many systems require that content type is set regardless and will fail, @@ -191,18 +198,15 @@ HttpClient::~HttpClient() = default; Result HttpClient::Get( const std::string& path, const std::unordered_map& params, const std::unordered_map& headers, - const ErrorHandler& error_handler, auth::AuthSession& session, - std::optional timeout) { + const ErrorHandler& error_handler, auth::AuthSession& session) { ICEBERG_ASSIGN_OR_RAISE(auto url, AppendQueryString(path, params)); ICEBERG_ASSIGN_OR_RAISE(auto authenticated, AuthenticateRequest(session, HttpMethod::kGet, std::move(url), MergeHeaders(default_headers_, headers))); - // Zero leaves the request unbounded. cpr::Response response = cpr::Get(cpr::Url{authenticated.url}, ToCprHeader(authenticated), *connection_pool_, - cpr::Timeout{timeout.value_or(std::chrono::milliseconds::zero())}); + cpr::ConnectTimeout{connection_timeout_}, cpr::Timeout{socket_timeout_}); - ICEBERG_RETURN_UNEXPECTED(CheckGetTransfer(response)); ICEBERG_RETURN_UNEXPECTED(HandleFailureResponse(response, error_handler)); HttpResponse http_response; http_response.impl_ = std::make_unique(std::move(response)); @@ -219,7 +223,8 @@ Result HttpClient::Post( MergeHeaders(default_headers_, headers), body)); cpr::Response response = cpr::Post(cpr::Url{authenticated.url}, cpr::Body{authenticated.body}, - ToCprHeader(authenticated), *connection_pool_); + ToCprHeader(authenticated), *connection_pool_, + cpr::ConnectTimeout{connection_timeout_}, cpr::Timeout{socket_timeout_}); ICEBERG_RETURN_UNEXPECTED(HandleFailureResponse(response, error_handler)); HttpResponse http_response; @@ -249,7 +254,8 @@ Result HttpClient::PostForm( std::move(encoded_body))); cpr::Response response = cpr::Post(cpr::Url{authenticated.url}, cpr::Body{authenticated.body}, - ToCprHeader(authenticated), *connection_pool_); + ToCprHeader(authenticated), *connection_pool_, + cpr::ConnectTimeout{connection_timeout_}, cpr::Timeout{socket_timeout_}); ICEBERG_RETURN_UNEXPECTED(HandleFailureResponse(response, error_handler)); HttpResponse http_response; @@ -263,8 +269,9 @@ Result HttpClient::Head( ICEBERG_ASSIGN_OR_RAISE(auto authenticated, AuthenticateRequest(session, HttpMethod::kHead, path, MergeHeaders(default_headers_, headers))); - cpr::Response response = cpr::Head(cpr::Url{authenticated.url}, - ToCprHeader(authenticated), *connection_pool_); + cpr::Response response = cpr::Head( + cpr::Url{authenticated.url}, ToCprHeader(authenticated), *connection_pool_, + cpr::ConnectTimeout{connection_timeout_}, cpr::Timeout{socket_timeout_}); ICEBERG_RETURN_UNEXPECTED(HandleFailureResponse(response, error_handler)); HttpResponse http_response; @@ -281,8 +288,9 @@ Result HttpClient::Delete( auto authenticated, AuthenticateRequest(session, HttpMethod::kDelete, std::move(url), MergeHeaders(default_headers_, headers))); - cpr::Response response = cpr::Delete(cpr::Url{authenticated.url}, - ToCprHeader(authenticated), *connection_pool_); + cpr::Response response = cpr::Delete( + cpr::Url{authenticated.url}, ToCprHeader(authenticated), *connection_pool_, + cpr::ConnectTimeout{connection_timeout_}, cpr::Timeout{socket_timeout_}); ICEBERG_RETURN_UNEXPECTED(HandleFailureResponse(response, error_handler)); HttpResponse http_response; diff --git a/src/iceberg/catalog/rest/http_client.h b/src/iceberg/catalog/rest/http_client.h index 1f2d02eac2..a185064041 100644 --- a/src/iceberg/catalog/rest/http_client.h +++ b/src/iceberg/catalog/rest/http_client.h @@ -22,7 +22,6 @@ #include #include #include -#include #include #include @@ -73,7 +72,10 @@ class ICEBERG_REST_EXPORT HttpResponse { /// \brief HTTP client for making requests to Iceberg REST Catalog API. class ICEBERG_REST_EXPORT HttpClient { public: - explicit HttpClient(std::unordered_map default_headers = {}); + /// \param connection_timeout,socket_timeout Zero leaves requests unbounded. + explicit HttpClient(std::unordered_map default_headers = {}, + std::chrono::milliseconds connection_timeout = {}, + std::chrono::milliseconds socket_timeout = {}); virtual ~HttpClient(); HttpClient(const HttpClient&) = delete; @@ -82,14 +84,10 @@ class ICEBERG_REST_EXPORT HttpClient { HttpClient& operator=(HttpClient&&) = delete; /// \brief Sends a GET request. - /// - /// \param timeout Limit on the whole request, after which it fails with an - /// IOError; nullopt for none. virtual Result Get( const std::string& path, const std::unordered_map& params, const std::unordered_map& headers, - const ErrorHandler& error_handler, auth::AuthSession& session, - std::optional timeout); + const ErrorHandler& error_handler, auth::AuthSession& session); /// \brief Sends a POST request. virtual Result Post( @@ -118,6 +116,8 @@ class ICEBERG_REST_EXPORT HttpClient { private: std::unordered_map default_headers_; + std::chrono::milliseconds connection_timeout_; + std::chrono::milliseconds socket_timeout_; std::unique_ptr connection_pool_; }; diff --git a/src/iceberg/catalog/rest/rest_catalog.cc b/src/iceberg/catalog/rest/rest_catalog.cc index 38728deac7..a945e25d4c 100644 --- a/src/iceberg/catalog/rest/rest_catalog.cc +++ b/src/iceberg/catalog/rest/rest_catalog.cc @@ -19,7 +19,9 @@ #include "iceberg/catalog/rest/rest_catalog.h" +#include #include +#include #include #include #include @@ -70,9 +72,6 @@ namespace iceberg::rest { namespace { -// A refresh runs inline in file operations and S3 requests, so bound it. -constexpr auto kLoadCredentialsTimeout = std::chrono::seconds(10); - class RestStorageCredentialProvider final : public StorageCredentialProvider { public: RestStorageCredentialProvider(std::shared_ptr client, @@ -87,10 +86,9 @@ class RestStorageCredentialProvider final : public StorageCredentialProvider { Result> Load() override { // The provider can be shared by several FileIOs; keep fetches serialized. std::lock_guard lock(mutex_); - ICEBERG_ASSIGN_OR_RAISE( - const auto response, - client_->Get(path_, params_, /*headers=*/{}, *TableErrorHandler::Instance(), - *session_, kLoadCredentialsTimeout)); + ICEBERG_ASSIGN_OR_RAISE(const auto response, + client_->Get(path_, params_, /*headers=*/{}, + *TableErrorHandler::Instance(), *session_)); // Parse errors can contain credential data; return a fixed message. auto json = FromJsonString(response.body()); if (!json.has_value()) { @@ -134,6 +132,18 @@ std::string RestTableName(std::string_view catalog_name, return std::string(catalog_name) + '.' + identifier.ToString(); } +/// \brief Report a write whose response transfer failed as having an unknown +/// outcome, since the server may have applied it. +Result UnknownIfTransferFailed(Result response, + std::string_view operation, + std::string_view state) { + if (!response.has_value() && response.error().kind == ErrorKind::kIOError) { + return CommitStateUnknown("{} response transfer failed; the {} state is unknown: {}", + operation, state, response.error().message); + } + return response; +} + /// \brief Fetch server configuration from the REST catalog server. Result FetchServerConfig(HttpClient& client, const ResourcePaths& paths, const RestCatalogProperties& current_config, @@ -152,8 +162,7 @@ Result FetchServerConfig(HttpClient& client, const ResourcePaths& ICEBERG_ASSIGN_OR_RAISE(const auto response, client.Get(config_path, params, /*headers=*/{}, - *ConfigErrorHandler::Instance(), session, - /*timeout=*/std::nullopt)); + *ConfigErrorHandler::Instance(), session)); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); return CatalogConfigFromJson(json); } @@ -443,7 +452,14 @@ Result> RestCatalog::Make( if (!client_factory) { client_factory = [](const RestCatalogProperties& properties) -> Result> { - return std::make_shared(properties.ExtractHeaders()); + // A value that is not positive leaves requests unbounded. + const auto timeout_of = [&properties](const auto& entry) { + return std::chrono::milliseconds(std::max(0, properties.Get(entry))); + }; + return std::make_shared( + properties.ExtractHeaders(), + timeout_of(RestCatalogProperties::kClientConnectionTimeoutMs), + timeout_of(RestCatalogProperties::kClientSocketTimeoutMs)); }; } ICEBERG_ASSIGN_OR_RAISE(auto uri, config.Uri()); @@ -656,8 +672,7 @@ Result> RestCatalog::ListNamespaces( } ICEBERG_ASSIGN_OR_RAISE(const auto response, client_->Get(path, params, /*headers=*/{}, - *NamespaceErrorHandler::Instance(), session, - /*timeout=*/std::nullopt)); + *NamespaceErrorHandler::Instance(), session)); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); ICEBERG_ASSIGN_OR_RAISE(auto list_response, ListNamespacesResponseFromJson(json)); result.insert(result.end(), list_response.namespaces.begin(), @@ -691,8 +706,7 @@ Result> RestCatalog::GetNamespacePr ICEBERG_ASSIGN_OR_RAISE(auto path, paths_->Namespace_(ns)); ICEBERG_ASSIGN_OR_RAISE(const auto response, client_->Get(path, /*params=*/{}, /*headers=*/{}, - *NamespaceErrorHandler::Instance(), session, - /*timeout=*/std::nullopt)); + *NamespaceErrorHandler::Instance(), session)); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); ICEBERG_ASSIGN_OR_RAISE(auto get_response, GetNamespaceResponseFromJson(json)); return get_response.properties; @@ -751,8 +765,7 @@ Result> RestCatalog::ListTables( } ICEBERG_ASSIGN_OR_RAISE(const auto response, client_->Get(path, params, /*headers=*/{}, - *TableErrorHandler::Instance(), session, - /*timeout=*/std::nullopt)); + *TableErrorHandler::Instance(), session)); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); ICEBERG_ASSIGN_OR_RAISE(auto list_response, ListTablesResponseFromJson(json)); result.insert(result.end(), list_response.identifiers.begin(), @@ -786,9 +799,11 @@ Result RestCatalog::CreateTableInternal( ICEBERG_ASSIGN_OR_RAISE(auto request_json, ToJson(request)); ICEBERG_ASSIGN_OR_RAISE(auto json_request, ToJsonString(request_json)); - ICEBERG_ASSIGN_OR_RAISE(const auto response, - client_->Post(path, json_request, /*headers=*/{}, - *TableErrorHandler::Instance(), session)); + ICEBERG_ASSIGN_OR_RAISE( + const auto response, + UnknownIfTransferFailed(client_->Post(path, json_request, /*headers=*/{}, + *TableErrorHandler::Instance(), session), + "Create table", "table")); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); ICEBERG_ASSIGN_OR_RAISE(auto load_result, LoadTableResultFromJson(json)); @@ -837,7 +852,9 @@ Result RestCatalog::UpdateTableInternal( } ICEBERG_ASSIGN_OR_RAISE( const auto response, - client_->Post(path, json_request, /*headers=*/{}, *error_handler, session)); + UnknownIfTransferFailed( + client_->Post(path, json_request, /*headers=*/{}, *error_handler, session), + "Commit", "commit")); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); ICEBERG_ASSIGN_OR_RAISE(auto commit_response, CommitTableResponseFromJson(json)); @@ -947,10 +964,9 @@ Result RestCatalog::LoadTableInternal(const TableIdentifier& id params["snapshots"] = "all"; } - ICEBERG_ASSIGN_OR_RAISE( - const auto response, - client_->Get(path, params, /*headers=*/{}, *TableErrorHandler::Instance(), session, - /*timeout=*/std::nullopt)); + ICEBERG_ASSIGN_OR_RAISE(const auto response, + client_->Get(path, params, /*headers=*/{}, + *TableErrorHandler::Instance(), session)); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); return LoadTableResultFromJson(json); } diff --git a/src/iceberg/catalog/rest/rest_table_scan.cc b/src/iceberg/catalog/rest/rest_table_scan.cc index 6317b87efe..e7dae0f492 100644 --- a/src/iceberg/catalog/rest/rest_table_scan.cc +++ b/src/iceberg/catalog/rest/rest_table_scan.cc @@ -205,8 +205,7 @@ Result FetchPlanningResult( for (int32_t retry = 0; retry <= poll_config.max_retries; ++retry) { ICEBERG_ASSIGN_OR_RAISE(auto response, ctx.client->Get(path, /*params=*/{}, /*headers=*/{}, - *PlanErrorHandler::Instance(), *ctx.session, - /*timeout=*/std::nullopt)); + *PlanErrorHandler::Instance(), *ctx.session)); ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body())); ICEBERG_ASSIGN_OR_RAISE(auto result, FetchPlanningResultResponseFromJson(json, specs, schema)); diff --git a/src/iceberg/test/arrow_s3_file_io_test.cc b/src/iceberg/test/arrow_s3_file_io_test.cc index e9b8e23677..f663d58e5a 100644 --- a/src/iceberg/test/arrow_s3_file_io_test.cc +++ b/src/iceberg/test/arrow_s3_file_io_test.cc @@ -57,7 +57,9 @@ #include "iceberg/test/executor.h" #include "iceberg/test/logging_test_helpers.h" #include "iceberg/test/matchers.h" +#include "iceberg/test/mock_io.h" #include "iceberg/util/macros.h" +#include "iceberg/util/timepoint.h" namespace { @@ -167,16 +169,8 @@ bool HasWarning(const CapturingLogger& logger, std::string_view substring = {}) }); } -class MockStorageCredentialProvider : public StorageCredentialProvider { - public: - MOCK_METHOD((Result>), Load, (), (override)); -}; - std::string ExpiresAtMs(std::chrono::milliseconds valid_for) { - return std::to_string( - std::chrono::duration_cast( - (std::chrono::system_clock::now() + valid_for).time_since_epoch()) - .count()); + return std::to_string(UnixMsFromTimePointMs(CurrentTimePointMs() + valid_for)); } std::vector ExpiringCredentials(std::chrono::milliseconds valid_for, @@ -839,6 +833,34 @@ TEST_F(ArrowS3FileIOTest, KeepsCredentialsWhenTheRefreshFails) { } } +TEST_F(ArrowS3FileIOTest, RejectsARefreshWithAnInvalidCredential) { + // Usable S3 keys do not make up for an invalid credential beside them, even + // one this FileIO would skip. + const std::vector invalid = { + {.prefix = "", .config = {{"k", "v"}}}, + {.prefix = "gs://bucket", .config = {}}, + {.prefix = "s3://other/", .config = {}}}; + for (const auto& bad : invalid) { + SCOPED_TRACE(bad.prefix); + ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); + auto* credentialed = io->AsSupportsStorageCredentials(); + ASSERT_NE(credentialed, nullptr); + + auto refreshed = ExpiringCredentials(std::chrono::hours(1), "refreshed-key"); + refreshed.push_back(bad); + auto provider = std::make_shared(); + EXPECT_CALL(*provider, Load()).WillOnce(::testing::Return(refreshed)); + const auto expiring = ExpiringCredentials(std::chrono::minutes(1), "expiring-key"); + ASSERT_THAT(credentialed->InitializeStorageCredentials(expiring, provider), IsOk()); + + auto logger = std::make_shared(); + ScopedDefaultLogger scoped(logger); + EXPECT_THAT(io->NewInputFile("s3://bucket/key"), IsOk()); + EXPECT_EQ(credentialed->credentials(), expiring); + EXPECT_TRUE(HasWarning(*logger, "Failed to refresh")); + } +} + TEST_F(ArrowS3FileIOTest, AcceptsRefreshedPrefixesInAnotherOrder) { ICEBERG_UNWRAP_OR_FAIL(auto io, MakeS3FileIO({})); auto* credentialed = io->AsSupportsStorageCredentials(); @@ -858,7 +880,8 @@ TEST_F(ArrowS3FileIOTest, AcceptsRefreshedPrefixesInAnotherOrder) { ASSERT_THAT(credentialed->InitializeStorageCredentials(initial, provider), IsOk()); EXPECT_THAT(io->NewInputFile("s3://first/key"), IsOk()); - EXPECT_EQ(credentialed->credentials(), refreshed); + EXPECT_THAT(credentialed->credentials(), + ::testing::UnorderedElementsAreArray(refreshed)); } TEST_F(ArrowS3FileIOTest, TracksTheExpiryOfTheTokenItSignsWith) { diff --git a/src/iceberg/test/auth_manager_test.cc b/src/iceberg/test/auth_manager_test.cc index 848356e30a..fd9605fc31 100644 --- a/src/iceberg/test/auth_manager_test.cc +++ b/src/iceberg/test/auth_manager_test.cc @@ -24,7 +24,6 @@ #include #include #include -#include #include #include #include @@ -254,8 +253,7 @@ TEST_F(AuthManagerTest, HttpClientRejectsParamsWhenUrlAlreadyHasQuery) { auto session = AuthSession::MakeDefault({}); auto result = client_->Get("http://127.0.0.1/v1/config?existing=true", {{"warehouse", "prod"}}, - /*headers=*/{}, *rest::DefaultErrorHandler::Instance(), *session, - /*timeout=*/std::nullopt); + /*headers=*/{}, *rest::DefaultErrorHandler::Instance(), *session); EXPECT_THAT(result, IsError(ErrorKind::kInvalidArgument)); EXPECT_THAT(result, HasErrorMessage("must not contain a query string")); diff --git a/src/iceberg/test/http_client_test.cc b/src/iceberg/test/http_client_test.cc index 092c80d6b6..4203e74266 100644 --- a/src/iceberg/test/http_client_test.cc +++ b/src/iceberg/test/http_client_test.cc @@ -135,22 +135,21 @@ TEST(HttpClientTest, GetFailsWhenItsTransferFails) { ASSERT_NO_FATAL_FAILURE(Listen(listener)); const auto url = std::format("http://127.0.0.1:{}/v1/config", listener.port); - HttpClient client; + HttpClient client({}, {}, std::chrono::milliseconds(500)); auto session = auth::AuthSession::MakeDefault({}); std::thread server( [&] { ServeOnce(listener.fd, test_case.response, test_case.end); }); - auto response = - client.Get(url, /*params=*/{}, /*headers=*/{}, *DefaultErrorHandler::Instance(), - *session, std::chrono::milliseconds(500)); + auto response = client.Get(url, /*params=*/{}, /*headers=*/{}, + *DefaultErrorHandler::Instance(), *session); server.join(); EXPECT_THAT(response, IsError(ErrorKind::kIOError)); EXPECT_THAT(response, HasErrorMessage(test_case.error)); } } -TEST(HttpClientTest, PostKeepsASuccessfulStatusWhenTheTransferBreaks) { - // A write that got its status may have taken effect, so it must not fail. - // Here the body arrives but its final chunk does not. +TEST(HttpClientTest, PostReportsTransferFailureAfterSuccessfulStatus) { + // A write that got its status may have taken effect, so the caller must treat + // the result as unknown rather than retrying or cleaning up automatically. constexpr std::string_view kBody = R"({"metadata-location":"s3://bucket/m.json"})"; const auto response_bytes = std::format( "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n" @@ -167,8 +166,8 @@ TEST(HttpClientTest, PostKeepsASuccessfulStatusWhenTheTransferBreaks) { auto response = client.Post(url, "{}", /*headers=*/{}, *DefaultErrorHandler::Instance(), *session); server.join(); - ASSERT_THAT(response, IsOk()); - EXPECT_EQ(response->body(), kBody); + EXPECT_THAT(response, IsError(ErrorKind::kIOError)); + EXPECT_THAT(response, HasErrorMessage("failed after status 200")); } #endif diff --git a/src/iceberg/test/mock_io.h b/src/iceberg/test/mock_io.h index f643a852ba..63abb6ea6d 100644 --- a/src/iceberg/test/mock_io.h +++ b/src/iceberg/test/mock_io.h @@ -37,6 +37,7 @@ #include #include "iceberg/file_io.h" +#include "iceberg/storage_credential.h" #include "iceberg/util/macros.h" namespace iceberg { @@ -256,4 +257,9 @@ class MockFileIO : public FileIO { std::shared_ptr files_ = std::make_shared(); }; +class MockStorageCredentialProvider : public StorageCredentialProvider { + public: + MOCK_METHOD((Result>), Load, (), (override)); +}; + } // namespace iceberg diff --git a/src/iceberg/test/rest_arrow_file_io_test.cc b/src/iceberg/test/rest_arrow_file_io_test.cc index 30066f946e..5b2569c30b 100644 --- a/src/iceberg/test/rest_arrow_file_io_test.cc +++ b/src/iceberg/test/rest_arrow_file_io_test.cc @@ -40,7 +40,9 @@ #include "iceberg/catalog/rest/rest_file_io.h" #include "iceberg/storage_credential.h" #include "iceberg/test/matchers.h" +#include "iceberg/test/mock_io.h" #include "iceberg/test/temp_file_test_base.h" +#include "iceberg/util/timepoint.h" namespace iceberg::rest { @@ -68,11 +70,6 @@ TEST_F(RestArrowFileIOTest, ReadsBackWhatItWroteThroughRealLocalFileIO) { #if ICEBERG_S3_ENABLED -class MockStorageCredentialProvider : public StorageCredentialProvider { - public: - MOCK_METHOD((Result>), Load, (), (override)); -}; - std::optional GetEnvIfSet(const char* key) { const char* value = std::getenv(key); if (value == nullptr || std::string_view(value).empty()) { @@ -216,9 +213,7 @@ TEST_F(RestArrowFileIOTest, RefreshesCredentialsThroughTheRealS3FileIO) { config["s3.session-token"] = *token; // A refreshed session token must say when it expires. config["s3.session-token-expires-at-ms"] = std::to_string( - std::chrono::duration_cast( - (std::chrono::system_clock::now() + std::chrono::hours(1)).time_since_epoch()) - .count()); + UnixMsFromTimePointMs(CurrentTimePointMs() + std::chrono::hours(1))); } if (const auto endpoint = GetEnvIfSet("ICEBERG_TEST_S3_ENDPOINT")) { config["s3.endpoint"] = *endpoint; diff --git a/src/iceberg/test/rest_catalog_integration_test.cc b/src/iceberg/test/rest_catalog_integration_test.cc index 1fb174a314..c03c9c2716 100644 --- a/src/iceberg/test/rest_catalog_integration_test.cc +++ b/src/iceberg/test/rest_catalog_integration_test.cc @@ -22,7 +22,6 @@ #include #include #include -#include #include #include #include @@ -405,8 +404,7 @@ TEST_F(RestCatalogIntegrationTest, FetchServerConfigDirect) { ICEBERG_UNWRAP_OR_FAIL(const auto response, client.Get(config_url, {}, /*headers=*/{}, - *DefaultErrorHandler::Instance(), *noop_session, - /*timeout=*/std::nullopt)); + *DefaultErrorHandler::Instance(), *noop_session)); ICEBERG_UNWRAP_OR_FAIL(auto json, FromJsonString(response.body())); EXPECT_TRUE(json.contains("defaults")); diff --git a/src/iceberg/test/rest_catalog_unit_test.cc b/src/iceberg/test/rest_catalog_unit_test.cc index 67a2f04988..86865a6638 100644 --- a/src/iceberg/test/rest_catalog_unit_test.cc +++ b/src/iceberg/test/rest_catalog_unit_test.cc @@ -17,7 +17,6 @@ * under the License. */ -#include #include #include #include @@ -68,8 +67,7 @@ class MockHttpClient : public HttpClient { (const std::string& path, (const std::unordered_map&)params, (const std::unordered_map&)headers, - const ErrorHandler& error_handler, auth::AuthSession& session, - std::optional timeout), + const ErrorHandler& error_handler, auth::AuthSession& session), (override)); MOCK_METHOD(Result, Post, @@ -206,7 +204,7 @@ class RestCatalogLoadTableTest : public ::testing::Test { props["scan-planning-mode"] = client_scan_mode; } CatalogConfig server_config{.endpoints = {endpoints.begin(), endpoints.end()}}; - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/config"), _, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/config"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, ToJson(server_config).dump()))); auto config = RestCatalogProperties::FromMap(std::move(props)); return RestCatalog::Make(config, /*metrics_executor=*/nullptr, @@ -250,7 +248,7 @@ class RestCatalogLabelsTest : public RestCatalogLoadTableTest, }; TEST_P(RestCatalogLabelsTest, LoadExposesLabels) { - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, LabelsResponse()))); ICEBERG_UNWRAP_OR_FAIL(auto catalog, MakeLabelsCatalog()); ICEBERG_UNWRAP_OR_FAIL(auto table, catalog->LoadTable(identifier_)); @@ -295,9 +293,9 @@ TEST_F(RestCatalogLoadTableTest, InitializationPropagatesClientFactoryErrors) { for (int fail_on_call : {1, 2}) { SCOPED_TRACE(fail_on_call); if (fail_on_call == 1) { - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); } else { - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/config"), _, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/config"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, ToJson(CatalogConfig{}).dump()))); } int calls = 0; @@ -324,14 +322,14 @@ TEST_F(RestCatalogLoadTableTest, InitializationUsesMergedServerConfig) { {"scan-planning-mode", "server"}, {"header.X-Shared", "server"}}, .endpoints = {all_plan_endpoints_.begin(), all_plan_endpoints_.end()}}; - EXPECT_CALL(*mock_client_, - Get(::testing::EndsWith("/v1/config"), - ::testing::Contains(::testing::Pair("warehouse", "client-warehouse")), - _, _, _, _)) + EXPECT_CALL( + *mock_client_, + Get(::testing::EndsWith("/v1/config"), + ::testing::Contains(::testing::Pair("warehouse", "client-warehouse")), _, _, _)) .WillOnce(Return(HttpResponse(200, ToJson(server_config).dump()))); EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/v1/catalog/namespaces/default/tables/my_table"), - ::testing::Contains(::testing::Pair("snapshots", "refs")), _, _, _, _)) + ::testing::Contains(::testing::Pair("snapshots", "refs")), _, _, _)) .WillOnce(Return(HttpResponse(200, kLoadTableDefaultResponse))); EXPECT_CALL( *mock_client_, @@ -371,7 +369,7 @@ TEST_F(RestCatalogLoadTableTest, InitializationUsesMergedServerConfig) { // LoadTable returns NotSupported. // -------------------------------------------------------------------------- TEST_F(RestCatalogLoadTableTest, ServerScanWithMissingEndpointReturnsNotSupported) { - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) .WillOnce(Return(HttpResponse(200, kLoadTableServerScanResponse))); ICEBERG_UNWRAP_OR_FAIL(auto catalog, MakeCatalog(no_plan_endpoint_set_)); @@ -387,14 +385,14 @@ TEST_F(RestCatalogLoadTableTest, ServerScanWithMissingEndpointReturnsNotSupporte TEST_F(RestCatalogLoadTableTest, ScanCredentialProviderWithoutPlanIdOmitsPlanId) { RegisterCredentialedFileIO(); EXPECT_CALL(*mock_client_, - Get(::testing::Not(::testing::HasSubstr("/credentials")), _, _, _, _, _)) + Get(::testing::Not(::testing::HasSubstr("/credentials")), _, _, _, _)) .WillOnce(Return(HttpResponse(200, kLoadTableServerScanResponse))); EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) .WillOnce(Return(HttpResponse( 200, R"({"status":"completed","storage-credentials":[{"prefix":"s3://bucket","config":{"k":"v"}}]})"))); - EXPECT_CALL(*mock_client_, Get(::testing::HasSubstr("/credentials"), - ::testing::IsEmpty(), _, _, _, ::testing::Optional(_))) + EXPECT_CALL(*mock_client_, + Get(::testing::HasSubstr("/credentials"), ::testing::IsEmpty(), _, _, _)) .WillOnce(Return(HttpResponse( 200, R"({"storage-credentials":[{"prefix":"s3://bucket","config":{"k":"v2"}}]})"))); @@ -421,7 +419,7 @@ TEST_F(RestCatalogLoadTableTest, ScanCredentialProviderWithoutPlanIdOmitsPlanId) TEST_F(RestCatalogLoadTableTest, ScanCredentialsWithoutCredentialsEndpointHaveNoProvider) { RegisterCredentialedFileIO(); - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) .WillOnce(Return(HttpResponse(200, kLoadTableServerScanResponse))); EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) .WillOnce(Return(HttpResponse(200, std::string(kPlanWithCredentialsResponse)))); @@ -456,7 +454,7 @@ TEST_P(RestCatalogPlanningModeTest, LoadAndCommitKeepTheEffectivePlanningMode) { const auto& config = GetParam(); auto response = nlohmann::json::parse(kLoadTableDefaultResponse); if (!config.table.empty()) response["config"]["scan-planning-mode"] = config.table; - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, response.dump()))); EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/tables/my_table"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, kCommitTableResponse))); @@ -509,13 +507,13 @@ TEST_P(RestCatalogScanLifetimeTest, ScanKeepsCredentialRefreshAlive) { {"timestamp-ms", 2}, {"manifest-list", "s3://bucket/manifests-2.avro"}, {"schema-id", 1}}}); - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/tables/my_table"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, response.dump()))); EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, std::string(kPlanWithCredentialsResponse)))); - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/credentials"), - ::testing::Contains(::testing::Pair("planId", "plan-1")), - _, _, _, ::testing::Optional(_))) + EXPECT_CALL(*mock_client_, + Get(::testing::EndsWith("/credentials"), + ::testing::Contains(::testing::Pair("planId", "plan-1")), _, _, _)) .WillOnce(Return(HttpResponse( 200, R"({"storage-credentials":[{"prefix":"s3://bucket","config":{"k":"v2"}}]})"))); diff --git a/src/iceberg/test/rest_table_scan_test.cc b/src/iceberg/test/rest_table_scan_test.cc index b9f9112f6d..a1b8e6e867 100644 --- a/src/iceberg/test/rest_table_scan_test.cc +++ b/src/iceberg/test/rest_table_scan_test.cc @@ -20,7 +20,6 @@ #include "iceberg/catalog/rest/rest_table_scan.h" #include -#include #include #include #include @@ -115,8 +114,7 @@ class MockHttpClient : public HttpClient { (const std::string& path, (const std::unordered_map&)params, (const std::unordered_map&)headers, - const ErrorHandler& error_handler, auth::AuthSession& session, - std::optional timeout), + const ErrorHandler& error_handler, auth::AuthSession& session), (override)); MOCK_METHOD(Result, Post, @@ -606,7 +604,7 @@ TEST_F(RestTableScanTest, PlanFilesStreamPollNumRetriesZero) { R"({"status":"submitted","plan-id":"plan-no-retry"})"; EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) .WillOnce(Return(HttpResponse(200, std::string(kSubmittedBody)))); - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) .WillOnce(Return(HttpResponse(200, std::string(kSubmittedBody)))); EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)) .WillOnce(Return(HttpResponse(200, "{}"))); @@ -651,7 +649,7 @@ TEST_F(RestTableScanTest, CancelCalledOnFetchPlanningResultGetError) { EXPECT_CALL(*mock_client_, Post(_, _, _, _, _)) .WillOnce(Return(HttpResponse(200, std::string(kSubmittedBody)))); - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)) .WillOnce(Return(IOError("network failure"))); EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)) .WillOnce(Return(IOError("cancel failure"))); @@ -970,7 +968,7 @@ TEST_P(RestScanResidualTest, PreservesFilteringAndDeleteFiles) { if (submitted) { completed.erase("plan-id"); EXPECT_CALL(*mock_client_, - Get(::testing::EndsWith("/plan/residual-plan"), _, _, _, _, _)) + Get(::testing::EndsWith("/plan/residual-plan"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, completed.dump()))); } if (fetch_token) { @@ -1086,7 +1084,7 @@ TEST_P(RestScanRegressionTest, MissingPollingEndpointCancelsSubmittedPlan) { EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) .WillOnce(Return( HttpResponse(200, R"({"status":"submitted","plan-id":"missing-poll"})"))); - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); EXPECT_CALL(*mock_client_, Delete(::testing::EndsWith("/plan/missing-poll"), _, _, _, _)) .WillOnce(Return(HttpResponse(204, ""))); @@ -1166,10 +1164,10 @@ class RestScanStreamTest : public RestTableScanTest, .WillOnce(Return(HttpResponse(200, initial.dump()))); if (GetParam()) { completed.erase("plan-id"); - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/plan/stream"), _, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/plan/stream"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, completed.dump()))); } else { - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); } } }; @@ -1239,10 +1237,10 @@ TEST_P(RestScanRegressionTest, EmptyCompletedPlan) { EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, initial))); if (submitted) { - EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/plan/empty"), _, _, _, _, _)) + EXPECT_CALL(*mock_client_, Get(::testing::EndsWith("/plan/empty"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, R"({"status":"completed"})"))); } else { - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); } EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); ICEBERG_UNWRAP_OR_FAIL(auto tasks, PlanScan(GetParam(), MakeContext(), [](auto&) {})); @@ -1256,7 +1254,7 @@ TEST_P(RestScanRegressionTest, FailedPlanReturnsError) { .WillOnce(Return(HttpResponse( 200, R"({"status":"failed","error":{"message":"server error","type":"ServerError","code":500}})"))); - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); EXPECT_CALL(*mock_client_, Delete(_, _, _, _, _)).Times(0); EXPECT_THAT( PlanScan(GetParam(), MakeContext(), [](auto&) {}), @@ -1305,7 +1303,7 @@ TEST_P(RestScanRegressionTest, InvalidPlanResponseCancelsKnownId) { SCOPED_TRACE(test_case.body); EXPECT_CALL(*mock_client_, Post(::testing::EndsWith("/plan"), _, _, _, _)) .WillOnce(Return(HttpResponse(200, test_case.body))); - EXPECT_CALL(*mock_client_, Get(_, _, _, _, _, _)).Times(0); + EXPECT_CALL(*mock_client_, Get(_, _, _, _, _)).Times(0); if (test_case.cancel) { EXPECT_CALL(*mock_client_, Delete(::testing::EndsWith("/plan/invalid"), _, _, _, _)) .WillOnce(Return(HttpResponse(204, ""))); From 51232b6051843e3645f8f761d8e0a732a3112125 Mon Sep 17 00:00:00 2001 From: Gang Wu Date: Sun, 18 Oct 2026 00:02:02 +0800 Subject: [PATCH 6/7] test(s3): match credentials to compatible URI schemes --- src/iceberg/test/arrow_s3_file_io_test.cc | 14 +++++--------- src/iceberg/test/rest_arrow_file_io_test.cc | 8 ++++---- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/src/iceberg/test/arrow_s3_file_io_test.cc b/src/iceberg/test/arrow_s3_file_io_test.cc index f663d58e5a..9a5eace0dd 100644 --- a/src/iceberg/test/arrow_s3_file_io_test.cc +++ b/src/iceberg/test/arrow_s3_file_io_test.cc @@ -1322,10 +1322,8 @@ TEST_F(ArrowS3FileIOTest, DeleteFilesAttemptsEveryFile) { EXPECT_THAT(io->DeleteFiles({paths[0], paths[2]}), IsOk()); } -// The credential is vended under the oss spelling and the object addressed as -// `s3://`, so they only meet through canonicalization — and every other path -// to authentication is broken. (rest_arrow_file_io_test covers the mirrored -// direction.) +// The credential and object use the same S3-compatible scheme. The scheme is +// preserved for routing, while the underlying endpoint remains S3-compatible. TEST_F(ArrowS3FileIOTest, AppliesOssCredentialInRealRoundTrip) { if (!HasIntegrationEnv()) { GTEST_SKIP() << "Set ICEBERG_TEST_S3_URI to enable S3 IO test"; @@ -1353,15 +1351,13 @@ TEST_F(ArrowS3FileIOTest, AppliesOssCredentialInRealRoundTrip) { const auto scheme_end = object_uri.find("://"); ASSERT_NE(scheme_end, std::string::npos) << "ICEBERG_TEST_S3_URI must carry a scheme"; // Both spellings are forced, so they cross whatever scheme the env URI uses. - const auto oss_spelling = std::string("oss").append(object_uri.substr(scheme_end)); - const auto oss_prefix = - oss_spelling.substr(0, oss_spelling.size() - object_name.size()); - const auto s3_uri = std::string("s3").append(object_uri.substr(scheme_end)); + const auto oss_uri = std::string("oss").append(object_uri.substr(scheme_end)); + const auto oss_prefix = oss_uri.substr(0, oss_uri.size() - object_name.size()); EXPECT_THAT(credentialed->SetStorageCredentials( {{.prefix = oss_prefix, .config = std::move(properties)}}), IsOk()); - EXPECT_THAT(CheckReadWrite(*io, s3_uri, "hello oss with vended credentials"), IsOk()); + EXPECT_THAT(CheckReadWrite(*io, oss_uri, "hello oss with vended credentials"), IsOk()); } TEST_F(ArrowS3FileIOTest, DeleteFilesReachesEveryCredentialPrefix) { diff --git a/src/iceberg/test/rest_arrow_file_io_test.cc b/src/iceberg/test/rest_arrow_file_io_test.cc index 5b2569c30b..81152ee3f0 100644 --- a/src/iceberg/test/rest_arrow_file_io_test.cc +++ b/src/iceberg/test/rest_arrow_file_io_test.cc @@ -151,8 +151,8 @@ std::string AsOssUri(std::string_view uri) { } // Resolution, credential matching and real I/O for `oss://`. The credential -// env vars are scrubbed, so only the vended `s3`-scoped credential matching -// the canonicalized location can authenticate. +// env vars are scrubbed, so only the vended `oss`-scoped credential matching +// the location can authenticate. TEST_F(RestArrowFileIOTest, ReadsBackWhatItWroteThroughAnOssLocation) { const auto base_uri = GetEnvIfSet("ICEBERG_TEST_S3_URI"); if (!base_uri.has_value()) { @@ -179,10 +179,10 @@ TEST_F(RestArrowFileIOTest, ReadsBackWhatItWroteThroughAnOssLocation) { ScopedScrubbedAwsCredentialEnv scrubbed; - // Scoped to `s3`, while the data it grants access to is addressed as `oss://`. + // The scheme is preserved for routing, while the endpoint is S3-compatible. auto io = MakeTableFileIO({{"warehouse", "logical_warehouse_name"}}, /*table_config=*/{}, - {{.prefix = "s3", .config = std::move(credential_config)}}); + {{.prefix = "oss", .config = std::move(credential_config)}}); ASSERT_THAT(io, IsOk()); auto object_uri = AsOssUri(*base_uri); From d3ea98f32c0454ae3712bb60fc6472d96ec8cd7c Mon Sep 17 00:00:00 2001 From: Gang Wu Date: Sun, 18 Oct 2026 00:19:55 +0800 Subject: [PATCH 7/7] test(s3): use a valid OSS credential prefix --- src/iceberg/test/rest_arrow_file_io_test.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/iceberg/test/rest_arrow_file_io_test.cc b/src/iceberg/test/rest_arrow_file_io_test.cc index 81152ee3f0..7553e30d6a 100644 --- a/src/iceberg/test/rest_arrow_file_io_test.cc +++ b/src/iceberg/test/rest_arrow_file_io_test.cc @@ -180,9 +180,9 @@ TEST_F(RestArrowFileIOTest, ReadsBackWhatItWroteThroughAnOssLocation) { ScopedScrubbedAwsCredentialEnv scrubbed; // The scheme is preserved for routing, while the endpoint is S3-compatible. - auto io = MakeTableFileIO({{"warehouse", "logical_warehouse_name"}}, - /*table_config=*/{}, - {{.prefix = "oss", .config = std::move(credential_config)}}); + auto io = + MakeTableFileIO({{"warehouse", "logical_warehouse_name"}}, /*table_config=*/{}, + {{.prefix = "oss://", .config = std::move(credential_config)}}); ASSERT_THAT(io, IsOk()); auto object_uri = AsOssUri(*base_uri);