feat(rest): load table credentials from the LoadCredentials endpoint

src/iceberg/arrow/s3/arrow_s3_file_io.cc

@@ -17,9 +17,12 @@
  * under the License.
  */
 
+#include <algorithm>
 #include <cstdlib>
 #include <memory>
+#include <mutex>
 #include <optional>
+#include <shared_mutex>
 #include <string>
 #include <string_view>
 #include <unordered_map>
@@ -190,7 +193,7 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials {
  public:
   ArrowS3FileIO(std::shared_ptr<::arrow::fs::FileSystem> arrow_fs,
                 std::unordered_map<std::string, std::string> default_properties)
-      : default_file_io_(std::move(arrow_fs)),
+      : default_file_io_(std::make_shared<ArrowFileSystemFileIO>(std::move(arrow_fs))),
         default_properties_(std::move(default_properties)) {}
 
   Result<std::unique_ptr<InputFile>> NewInputFile(std::string file_location) override;
@@ -207,27 +210,67 @@ class ArrowS3FileIO final : public FileIO, public SupportsStorageCredentials {
   Status SetStorageCredentials(
       const std::vector<StorageCredential>& storage_credentials) override;
 
-  const std::vector<StorageCredential>& credentials() const override {
+  std::vector<StorageCredential> credentials() const override {
+    std::shared_lock lock(mutex_);
     return storage_credentials_;
   }
 
   SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; }
 
  private:
-  ArrowFileSystemFileIO& FileIOForPath(std::string_view location);
-
-  ArrowFileSystemFileIO default_file_io_;
+  /// \brief Delegate serving `location`, pinned by the caller against a
+  /// concurrent credential install.
+  std::shared_ptr<ArrowFileSystemFileIO> FileIOForPath(std::string_view location);
+
+  using DelegatesByPrefix =
+      std::vector<std::pair<std::string, std::shared_ptr<ArrowFileSystemFileIO>>>;
+
+  /// \brief Longest-prefix match against one consistent view of the delegates.
+  static std::shared_ptr<ArrowFileSystemFileIO> MatchDelegate(
+      const std::shared_ptr<ArrowFileSystemFileIO>& fallback,
+      const DelegatesByPrefix& by_prefix, std::string_view location);
+
+  /// \brief Build a delegate for each credential this FileIO can serve.
+  ///
+  /// Lock-free on purpose: building an S3 client can reach out to discover a
+  /// bucket region, which would stall every concurrent operation. Reads no
+  /// mutable member state.
+  Result<DelegatesByPrefix> BuildDelegates(
+      const std::vector<StorageCredential>& storage_credentials) const;
+
+  /// \brief Swap in credentials and delegates, handing back the retired ones.
+  ///
+  /// 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<StorageCredential>& storage_credentials,
+                          DelegatesByPrefix& delegates);
+
+  std::shared_ptr<ArrowFileSystemFileIO> default_file_io_;
   std::unordered_map<std::string, std::string> default_properties_;
+  // Guards everything below; shared because reads happen per file operation.
+  mutable std::shared_mutex mutex_;
   std::vector<StorageCredential> storage_credentials_;
-  std::vector<std::pair<std::string, std::unique_ptr<ArrowFileSystemFileIO>>>
-      file_io_by_prefix_;
+  DelegatesByPrefix file_io_by_prefix_;
 };
 
 Status ArrowS3FileIO::SetStorageCredentials(
     const std::vector<StorageCredential>& storage_credentials) {
-  std::vector<std::pair<std::string, std::unique_ptr<ArrowFileSystemFileIO>>>
-      file_io_by_prefix;
-  file_io_by_prefix.reserve(storage_credentials.size());
+  ICEBERG_ASSIGN_OR_RAISE(auto delegates, BuildDelegates(storage_credentials));
+  auto credentials = storage_credentials;
+  {
+    std::unique_lock lock(mutex_);
+    InstallCredentials(credentials, delegates);
+  }
+  // `credentials` and `delegates` now hold the retired generation and destruct
+  // here, outside the lock.
+  return {};
+}
+
+Result<ArrowS3FileIO::DelegatesByPrefix> ArrowS3FileIO::BuildDelegates(
+    const std::vector<StorageCredential>& 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());
@@ -242,62 +285,91 @@ Status ArrowS3FileIO::SetStorageCredentials(
       properties[key] = value;
     }
     ICEBERG_ASSIGN_OR_RAISE(auto fs, BuildArrowS3FileSystem(properties));
-    file_io_by_prefix.emplace_back(
-        CanonicalizeS3Scheme(credential.prefix),
-        std::make_unique<ArrowFileSystemFileIO>(std::move(fs)));
+    delegates.emplace_back(CanonicalizeS3Scheme(credential.prefix),
+                           std::make_shared<ArrowFileSystemFileIO>(std::move(fs)));
   }
-  if (file_io_by_prefix.empty() && !storage_credentials.empty()) {
+  if (delegates.empty() && !storage_credentials.empty()) {
     // Silent skipping of every vended credential is hard to diagnose: S3 access
     // would proceed with the default credentials and fail only at IO time.
     ICEBERG_LOG_WARN(
         "None of the {} vended storage credential(s) has an S3-compatible prefix; "
         "S3 access will use the default credentials",
         storage_credentials.size());
   }
-  file_io_by_prefix_ = std::move(file_io_by_prefix);
-  storage_credentials_ = storage_credentials;
-  return {};
+  return delegates;
 }
 
-ArrowFileSystemFileIO& ArrowS3FileIO::FileIOForPath(std::string_view location) {
-  if (file_io_by_prefix_.empty()) {
-    return default_file_io_;
+void ArrowS3FileIO::InstallCredentials(
+    std::vector<StorageCredential>& storage_credentials, DelegatesByPrefix& delegates) {
+  file_io_by_prefix_.swap(delegates);
+  storage_credentials_.swap(storage_credentials);
+}
+
+std::shared_ptr<ArrowFileSystemFileIO> ArrowS3FileIO::MatchDelegate(
+    const std::shared_ptr<ArrowFileSystemFileIO>& fallback,
+    const DelegatesByPrefix& by_prefix, std::string_view location) {
+  if (by_prefix.empty()) {
+    return fallback;
   }
   const std::string canonical = CanonicalizeS3Scheme(location);
-  ArrowFileSystemFileIO* best = &default_file_io_;
+  auto best = fallback;
   size_t best_len = 0;
-  for (const auto& [prefix, file_io] : file_io_by_prefix_) {
+  for (const auto& [prefix, file_io] : by_prefix) {
     if (prefix.size() > best_len && canonical.starts_with(prefix)) {
-      best = file_io.get();
+      best = file_io;
       best_len = prefix.size();
     }
   }
-  return *best;
+  return best;
+}
+
+std::shared_ptr<ArrowFileSystemFileIO> ArrowS3FileIO::FileIOForPath(
+    std::string_view location) {
+  std::shared_lock lock(mutex_);
+  return MatchDelegate(default_file_io_, file_io_by_prefix_, location);
 }
 
 Result<std::unique_ptr<InputFile>> ArrowS3FileIO::NewInputFile(
     std::string file_location) {
-  return FileIOForPath(file_location).NewInputFile(std::move(file_location));
+  return FileIOForPath(file_location)->NewInputFile(std::move(file_location));
 }
 
 Result<std::unique_ptr<InputFile>> ArrowS3FileIO::NewInputFile(std::string file_location,
                                                                size_t length) {
-  return FileIOForPath(file_location).NewInputFile(std::move(file_location), length);
+  return FileIOForPath(file_location)->NewInputFile(std::move(file_location), length);
 }
 
 Result<std::unique_ptr<OutputFile>> ArrowS3FileIO::NewOutputFile(
     std::string file_location) {
-  return FileIOForPath(file_location).NewOutputFile(std::move(file_location));
+  return FileIOForPath(file_location)->NewOutputFile(std::move(file_location));
 }
 
 Status ArrowS3FileIO::DeleteFile(const std::string& file_location) {
-  return FileIOForPath(file_location).DeleteFile(file_location);
+  return FileIOForPath(file_location)->DeleteFile(file_location);
 }
 
 Status ArrowS3FileIO::DeleteFiles(const std::vector<std::string>& file_locations) {
-  std::unordered_map<ArrowFileSystemFileIO*, std::vector<std::string>> locations_by_io;
+  // One snapshot so the whole batch matches the same delegate generation; only
+  // ever a handful of delegates, so a linear scan beats hashing.
+  std::shared_ptr<ArrowFileSystemFileIO> fallback;
+  DelegatesByPrefix by_prefix;
+  {
+    std::shared_lock lock(mutex_);
+    fallback = default_file_io_;
+    by_prefix = file_io_by_prefix_;
+  }
+  std::vector<std::pair<std::shared_ptr<ArrowFileSystemFileIO>, std::vector<std::string>>>
+      locations_by_io;
   for (const auto& file_location : file_locations) {
-    locations_by_io[&FileIOForPath(file_location)].push_back(file_location);
+    auto file_io = MatchDelegate(fallback, by_prefix, file_location);
+    auto it = std::ranges::find_if(
+        locations_by_io, [&](const auto& entry) { return entry.first == file_io; });
+    if (it == locations_by_io.end()) {
+      locations_by_io.emplace_back(std::move(file_io),
+                                   std::vector<std::string>{file_location});
+    } else {
+      it->second.push_back(file_location);
+    }
   }
   for (auto& [file_io, locations] : locations_by_io) {
     ICEBERG_RETURN_UNEXPECTED(file_io->DeleteFiles(locations));

src/iceberg/catalog/rest/json_serde.cc

@@ -152,6 +152,25 @@ Result<StorageCredential> StorageCredentialFromJson(const nlohmann::json& json)
   return credential;
 }
 
+/// \brief Reads the optional `storage-credentials` array shared by the
+/// LoadTable and LoadCredentials responses.
+Result<std::vector<StorageCredential>> StorageCredentialsFromJson(
+    const nlohmann::json& json) {
+  std::vector<StorageCredential> credentials;
+  auto it = json.find(kStorageCredentials);
+  if (it == json.end() || it->is_null()) {
+    return credentials;
+  }
+  if (!it->is_array()) {
+    return JsonParseError("Cannot parse storage credentials from non-array");
+  }
+  for (const auto& entry : *it) {
+    ICEBERG_ASSIGN_OR_RAISE(auto credential, StorageCredentialFromJson(entry));
+    credentials.push_back(std::move(credential));
+  }
+  return credentials;
+}
+
 template <typename Value>
 Result<std::map<int32_t, Value>> KeyValueMapFromJson(const nlohmann::json& json,
                                                      std::string_view key) {
@@ -738,19 +757,24 @@ Result<LoadTableResult> LoadTableResultFromJson(const nlohmann::json& json) {
   ICEBERG_ASSIGN_OR_RAISE(result.metadata, TableMetadataFromJson(metadata_json));
   ICEBERG_ASSIGN_OR_RAISE(result.config,
                           GetJsonValueOrDefault<decltype(result.config)>(json, kConfig));
-  if (auto it = json.find(kStorageCredentials); it != json.end() && !it->is_null()) {
-    if (!it->is_array()) {
-      return JsonParseError("Cannot parse storage credentials from non-array");
-    }
-    for (const auto& entry : *it) {
-      ICEBERG_ASSIGN_OR_RAISE(auto cred, StorageCredentialFromJson(entry));
-      result.storage_credentials.push_back(std::move(cred));
-    }
-  }
+  ICEBERG_ASSIGN_OR_RAISE(result.storage_credentials, StorageCredentialsFromJson(json));
   ICEBERG_RETURN_UNEXPECTED(result.Validate());
   return result;
 }
 
+Result<LoadCredentialsResponse> LoadCredentialsResponseFromJson(
+    const nlohmann::json& json) {
+  // Required here, unlike in LoadTable: reading a malformed response as "no
+  // credentials" would look like a refresh that succeeded and dropped them.
+  if (auto it = json.find(kStorageCredentials); it == json.end() || it->is_null()) {
+    return JsonParseError("Missing '{}'", kStorageCredentials);
+  }
+  LoadCredentialsResponse response;
+  ICEBERG_ASSIGN_OR_RAISE(response.storage_credentials, StorageCredentialsFromJson(json));
+  ICEBERG_RETURN_UNEXPECTED(response.Validate());
+  return response;
+}
+
 nlohmann::json ToJson(const ListNamespacesResponse& response) {
   nlohmann::json json;
   SetOptionalStringField(json, kNextPageToken, response.next_page_token);

src/iceberg/catalog/rest/json_serde_internal.h

@@ -71,6 +71,10 @@ template <>
 ICEBERG_REST_EXPORT Result<LoadTableResult> FromJson(const nlohmann::json& json);
 ICEBERG_REST_EXPORT Result<nlohmann::json> ToJson(const LoadTableResult& model);
 
+// Response-only model: a client never serializes it, so no ToJson.
+ICEBERG_REST_EXPORT Result<LoadCredentialsResponse> LoadCredentialsResponseFromJson(
+    const nlohmann::json& json);
+
 ICEBERG_REST_EXPORT Result<CreateTableRequest> CreateTableRequestFromJson(
     const nlohmann::json& json);
 template <>

src/iceberg/catalog/rest/rest_catalog.cc

@@ -42,6 +42,7 @@
 #include "iceberg/catalog/rest/rest_util.h"
 #include "iceberg/catalog/rest/types.h"
 #include "iceberg/json_serde_internal.h"
+#include "iceberg/logging/log_macros.h"
 #include "iceberg/metrics/metrics_reporters.h"
 #include "iceberg/partition_spec.h"
 #include "iceberg/result.h"
@@ -508,12 +509,65 @@ Result<std::shared_ptr<auth::AuthSession>> RestCatalog::TableAuthSession(
                                      std::move(contextual_session));
 }
 
+StorageCredentialRefresher RestCatalog::MakeCredentialRefresher(
+    const TableIdentifier& identifier,
+    std::shared_ptr<auth::AuthSession> table_session) const {
+  if (!supported_endpoints_.contains(Endpoint::TableCredentials())) {
+    // Not an error, but it surfaces much later as credentials expiring.
+    ICEBERG_LOG_DEBUG(
+        "Catalog does not advertise {}; vended credentials for '{}' will not be "
+        "refreshed",
+        Endpoint::TableCredentials().ToString(), ToString(identifier));
+    return nullptr;
+  }
+  auto path = paths_->Credentials(identifier);
+  if (!path.has_value()) {
+    ICEBERG_LOG_WARN(
+        "Cannot build the credentials path for '{}' ({}); its vended credentials "
+        "will not be refreshed",
+        ToString(identifier), path.error().message);
+    return nullptr;
+  }
+  auto client = client_;
+  auto credentials_path = std::move(path.value());
+  auto session = std::move(table_session);
+  // The catalog's destructor closes the session, and a table's FileIO can
+  // outlive the table keeping the catalog alive. No cycle: the catalog's own
+  // FileIO never gets a refresher.
+  auto catalog = shared_from_this();
+  return [catalog, client, credentials_path,
+          session]() -> Result<std::vector<StorageCredential>> {
+    ICEBERG_ASSIGN_OR_RAISE(const auto response,
+                            client->Get(credentials_path, /*params=*/{}, /*headers=*/{},
+                                        *TableErrorHandler::Instance(), *session));
+    // Parse errors embed the offending input, and this body carries
+    // credentials; strip the message so it can never reach a log.
+    auto json = FromJsonString(response.body());
+    if (!json.has_value()) {
+      return JsonParseError("Malformed LoadCredentials response");
+    }
+    auto result = LoadCredentialsResponseFromJson(*json);
+    if (!result.has_value()) {
+      return std::unexpected<Error>(
+          {.kind = result.error().kind, .message = "Malformed LoadCredentials response"});
+    }
+    return std::move(result->storage_credentials);
+  };
+}
+
 Result<std::shared_ptr<FileIO>> RestCatalog::TableFileIO(
-    const SessionContext& /*context*/,
+    const SessionContext& /*context*/, const TableIdentifier& identifier,
     const std::unordered_map<std::string, std::string>& table_config,
-    const std::vector<StorageCredential>& storage_credentials) const {
+    const std::vector<StorageCredential>& storage_credentials,
+    std::shared_ptr<auth::AuthSession> table_session) const {
   if (!table_config.empty() || !storage_credentials.empty()) {
-    return MakeTableFileIO(config_.configs(), table_config, storage_credentials);
+    // Only vended credentials expire, so only they need a refresher.
+    StorageCredentialRefresher refresher;
+    if (!storage_credentials.empty()) {
+      refresher = MakeCredentialRefresher(identifier, std::move(table_session));
+    }
+    return MakeTableFileIO(config_.configs(), table_config, storage_credentials,
+                           std::move(refresher));
   }
 
   return file_io_;
@@ -772,11 +826,12 @@ Result<std::shared_ptr<Transaction>> RestCatalog::StageCreateTable(
                           /*stage_create=*/true, *contextual_session));
   auto table_config = std::move(result.config);
   auto storage_credentials = std::move(result.storage_credentials);
-  ICEBERG_ASSIGN_OR_RAISE(auto table_io,
-                          TableFileIO(context, table_config, storage_credentials));
+  // Before the FileIO: refreshing its credentials reuses the table session.
   ICEBERG_ASSIGN_OR_RAISE(
       auto table_session,
       TableAuthSession(identifier, table_config, std::move(contextual_session)));
+  ICEBERG_ASSIGN_OR_RAISE(auto table_io, TableFileIO(context, identifier, table_config,
+                                                     storage_credentials, table_session));
   ICEBERG_ASSIGN_OR_RAISE(auto reporter, MakeTableReporter(identifier, table_session));
   auto table_catalog = std::make_shared<TableScopedCatalog>(
       shared_from_this(), context, identifier, table_config, std::move(table_session),
@@ -890,11 +945,12 @@ Result<std::shared_ptr<Table>> RestCatalog::MakeTableFromLoadResult(
     std::shared_ptr<auth::AuthSession> contextual_session) {
   auto table_config = std::move(result.config);
   auto storage_credentials = std::move(result.storage_credentials);
-  ICEBERG_ASSIGN_OR_RAISE(auto table_io,
-                          TableFileIO(context, table_config, storage_credentials));
+  // Before the FileIO: refreshing its credentials reuses the table session.
   ICEBERG_ASSIGN_OR_RAISE(
       auto table_session,
       TableAuthSession(identifier, table_config, std::move(contextual_session)));
+  ICEBERG_ASSIGN_OR_RAISE(auto table_io, TableFileIO(context, identifier, table_config,
+                                                     storage_credentials, table_session));
   ICEBERG_ASSIGN_OR_RAISE(auto reporter, MakeTableReporter(identifier, table_session));
   auto table_catalog = std::make_shared<TableScopedCatalog>(
       shared_from_this(), context, identifier, table_config, table_session, table_io);