Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
368 changes: 337 additions & 31 deletions src/iceberg/arrow/s3/arrow_s3_file_io.cc

Large diffs are not rendered by default.

3 changes: 3 additions & 0 deletions src/iceberg/arrow/s3/s3_properties.h
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
42 changes: 33 additions & 9 deletions src/iceberg/catalog/rest/json_serde.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
Expand Down
4 changes: 4 additions & 0 deletions src/iceberg/catalog/rest/json_serde_internal.h
Original file line number Diff line number Diff line change
Expand Up @@ -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 <>
Expand Down
70 changes: 63 additions & 7 deletions src/iceberg/catalog/rest/rest_catalog.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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_;
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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);
Expand Down
14 changes: 12 additions & 2 deletions src/iceberg/catalog/rest/rest_catalog.h
Original file line number Diff line number Diff line change
Expand Up @@ -85,9 +85,19 @@ 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.
///
/// Captures the catalog, so anything documented to outlive it -- the metrics
/// executor included -- must also outlive every FileIO holding a refresher.
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;
Expand Down
8 changes: 7 additions & 1 deletion src/iceberg/catalog/rest/rest_file_io.cc
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@

#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include "iceberg/catalog/rest/types.h"
Expand Down Expand Up @@ -57,14 +58,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");
Expand Down
6 changes: 5 additions & 1 deletion src/iceberg/catalog/rest/rest_file_io.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
15 changes: 15 additions & 0 deletions src/iceberg/catalog/rest/types.h
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,21 @@ using CreateTableResponse = LoadTableResult;
/// \brief Alias of LoadTableResult used as the body of LoadTableResponse
using LoadTableResponse = LoadTableResult;

/// \brief Response body of the LoadCredentials API.
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;
Expand Down
12 changes: 10 additions & 2 deletions src/iceberg/file_io.h
Original file line number Diff line number Diff line change
Expand Up @@ -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 concurrent install may replace them. 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
Loading