feat(io): refresh vended storage credentials before they expire

Author: plusplusjiajia

src/iceberg/arrow/s3/arrow_s3_file_io.cc

@@ -40,6 +40,9 @@ struct S3Properties {
   static constexpr std::string_view kSecretAccessKey = "s3.secret-access-key";
   /// AWS session token (for temporary credentials)
   static constexpr std::string_view kSessionToken = "s3.session-token";
+  /// Epoch milliseconds at which a vended session token stops being valid
+  static constexpr std::string_view kSessionTokenExpiresAtMs =
+      "s3.session-token-expires-at-ms";
   /// AWS region, standard Iceberg client property.
   static constexpr std::string_view kClientRegion = "client.region";
   /// Custom endpoint override (for MinIO, LocalStack, etc.)

src/iceberg/arrow/s3/s3_properties.h

@@ -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.cc

@@ -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/json_serde_internal.h

@@ -42,6 +42,7 @@
 #include "iceberg/catalog/rest/rest_util.h"
 #include "iceberg/catalog/rest/types.h"
 #include "iceberg/json_serde_internal.h"
+#include "iceberg/logging/log_macros.h"
 #include "iceberg/metrics/metrics_reporters.h"
 #include "iceberg/partition_spec.h"
 #include "iceberg/result.h"
@@ -508,12 +509,56 @@ Result<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));
+    ICEBERG_ASSIGN_OR_RAISE(auto json, FromJsonString(response.body()));
+    ICEBERG_ASSIGN_OR_RAISE(auto result, LoadCredentialsResponseFromJson(json));
+    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 +817,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 +936,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);

src/iceberg/catalog/rest/rest_catalog.cc

@@ -85,9 +85,16 @@ class ICEBERG_REST_EXPORT RestCatalog final
       std::shared_ptr<auth::AuthSession> contextual_session);
 
   Result<std::shared_ptr<FileIO>> 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;
+
+  /// \brief Callback that reloads this table's vended credentials, or nullptr
+  /// when the catalog does not serve the LoadCredentials endpoint.
+  StorageCredentialRefresher MakeCredentialRefresher(
+      const TableIdentifier& identifier,
+      std::shared_ptr<auth::AuthSession> table_session) const;
 
   Result<std::vector<Namespace>> ListNamespaces(const Namespace& ns,
                                                 auth::AuthSession& session) const;

src/iceberg/catalog/rest/rest_catalog.h

@@ -21,6 +21,7 @@
 
 #include <string>
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
 #include "iceberg/catalog/rest/types.h"
@@ -61,14 +62,19 @@ Result<std::unique_ptr<FileIO>> MakeCatalogFileIO(const RestCatalogProperties& c
 Result<std::unique_ptr<FileIO>> MakeTableFileIO(
     const std::unordered_map<std::string, std::string>& catalog_config,
     const std::unordered_map<std::string, std::string>& table_config,
-    const std::vector<StorageCredential>& storage_credentials) {
+    const std::vector<StorageCredential>& storage_credentials,
+    StorageCredentialRefresher refresher) {
   const auto default_properties = MergeFileIOProperties(catalog_config, table_config);
   ICEBERG_ASSIGN_OR_RAISE(
       auto io, MakeCatalogFileIO(RestCatalogProperties::FromMap(default_properties)));
 
   if (storage_credentials.empty()) {
     return io;
   } else if (auto* credentialed = io->AsSupportsStorageCredentials()) {
+    // First, so the FileIO never briefly holds credentials it cannot replace.
+    if (refresher) {
+      credentialed->SetCredentialRefresher(std::move(refresher));
+    }
     ICEBERG_RETURN_UNEXPECTED(credentialed->SetStorageCredentials(storage_credentials));
   } else {
     return NotSupported("Configured FileIO does not support vended storage credentials");

src/iceberg/catalog/rest/rest_file_io.cc

@@ -41,9 +41,13 @@ ICEBERG_REST_EXPORT Result<std::unique_ptr<FileIO>> MakeCatalogFileIO(
     const RestCatalogProperties& config);
 
 /// \brief Build the configured table FileIO and apply storage credentials if present.
+///
+/// \param refresher Optional callback used to replace the vended credentials
+/// before they expire; ignored when the FileIO cannot tell when they do.
 ICEBERG_REST_EXPORT Result<std::unique_ptr<FileIO>> MakeTableFileIO(
     const std::unordered_map<std::string, std::string>& catalog_config,
     const std::unordered_map<std::string, std::string>& table_config,
-    const std::vector<StorageCredential>& storage_credentials);
+    const std::vector<StorageCredential>& storage_credentials,
+    StorageCredentialRefresher refresher = nullptr);
 
 }  // namespace iceberg::rest

src/iceberg/catalog/rest/rest_file_io.h

@@ -209,6 +209,23 @@ using CreateTableResponse = LoadTableResult;
 /// \brief Alias of LoadTableResult used as the body of LoadTableResponse
 using LoadTableResponse = LoadTableResult;
 
+/// \brief Response body of the LoadCredentials API.
+///
+/// Used to replace the credentials vended alongside a table before they expire.
+struct ICEBERG_REST_EXPORT LoadCredentialsResponse {
+  std::vector<StorageCredential> storage_credentials;
+
+  /// \brief Validates the LoadCredentialsResponse.
+  Status Validate() const {
+    for (const auto& credential : storage_credentials) {
+      ICEBERG_RETURN_UNEXPECTED(credential.Validate());
+    }
+    return {};
+  }
+
+  bool operator==(const LoadCredentialsResponse& other) const = default;
+};
+
 /// \brief Response body for listing namespaces.
 struct ICEBERG_REST_EXPORT ListNamespacesResponse {
   PageToken next_page_token;

src/iceberg/catalog/rest/types.h

@@ -193,8 +193,16 @@ class ICEBERG_EXPORT SupportsStorageCredentials {
   virtual Status SetStorageCredentials(
       const std::vector<StorageCredential>& storage_credentials) = 0;
 
-  /// \brief Return currently installed storage credentials.
-  virtual const std::vector<StorageCredential>& credentials() const = 0;
+  /// \brief Return the storage credentials this FileIO holds.
+  ///
+  /// By value because a refresh may replace them concurrently. An
+  /// implementation that delegates may report what was installed on it.
+  virtual std::vector<StorageCredential> credentials() const = 0;
+
+  /// \brief Install a callback that re-fetches credentials before they expire.
+  ///
+  /// Ignored by implementations that cannot tell when theirs expire.
+  virtual void SetCredentialRefresher(StorageCredentialRefresher /*refresher*/) {}
 };
 
 }  // namespace iceberg