Skip to content
Open
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
332 changes: 305 additions & 27 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 @@ -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.)
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
61 changes: 54 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,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_;
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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);
Expand Down
11 changes: 9 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,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;
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 @@ -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");
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
17 changes: 17 additions & 0 deletions src/iceberg/catalog/rest/types.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
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 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
18 changes: 15 additions & 3 deletions src/iceberg/resolving_file_io.cc
Original file line number Diff line number Diff line change
Expand Up @@ -61,8 +61,12 @@ Result<FileIO*> ResolvingFileIO::FileIOForPath(std::string_view location) {
FileIORegistry::Load(std::string(name), properties_));
// Forward all credentials; each implementation applies the prefixes it
// understands.
if (!storage_credentials_.empty()) {
if (auto* credentialed = io->AsSupportsStorageCredentials()) {
if (auto* credentialed = io->AsSupportsStorageCredentials()) {
// Before the credentials, so the delegate can always replace them.
if (refresher_) {
credentialed->SetCredentialRefresher(refresher_);
}
if (!storage_credentials_.empty()) {
ICEBERG_RETURN_UNEXPECTED(
credentialed->SetStorageCredentials(storage_credentials_));
}
Expand Down Expand Up @@ -117,8 +121,16 @@ Status ResolvingFileIO::SetStorageCredentials(
return {};
}

const std::vector<StorageCredential>& ResolvingFileIO::credentials() const {
std::vector<StorageCredential> ResolvingFileIO::credentials() const {
std::lock_guard lock(mutex_);
return storage_credentials_;
}

void ResolvingFileIO::SetCredentialRefresher(StorageCredentialRefresher refresher) {
// Drop the cached delegates so they are rebuilt with the refresher.
std::lock_guard lock(mutex_);
refresher_ = std::move(refresher);
io_by_name_.clear();
}

} // namespace iceberg
20 changes: 14 additions & 6 deletions src/iceberg/resolving_file_io.h
Original file line number Diff line number Diff line change
Expand Up @@ -47,9 +47,10 @@ namespace iceberg {
/// them; each applies the prefixes it understands and ignores the rest.
///
/// Lazy resolution is internally synchronized, so file operations may run
/// concurrently. Credentials are not: install them before sharing the instance,
/// since credentials() hands out a reference that SetStorageCredentials
/// replaces.
/// concurrently -- provided credentials and any refresher were installed before
/// the instance was shared, as the REST catalog does when it binds a table's
/// FileIO. Installing either drops the resolved implementations so they are
/// rebuilt, which would free one that a concurrent operation is still using.
class ICEBERG_EXPORT ResolvingFileIO final : public FileIO,
public SupportsStorageCredentials {
public:
Expand All @@ -70,7 +71,13 @@ class ICEBERG_EXPORT ResolvingFileIO final : public FileIO,
Status SetStorageCredentials(
const std::vector<StorageCredential>& storage_credentials) override;

const std::vector<StorageCredential>& credentials() const override;
/// \brief Return the credentials installed on this resolver.
///
/// Not necessarily the ones in use: a resolved implementation refreshes its
/// own without reporting back.
std::vector<StorageCredential> credentials() const override;

void SetCredentialRefresher(StorageCredentialRefresher refresher) override;

SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; }

Expand All @@ -79,9 +86,10 @@ class ICEBERG_EXPORT ResolvingFileIO final : public FileIO,
Result<FileIO*> FileIOForPath(std::string_view location);

std::unordered_map<std::string, std::string> properties_;
// Guards lazy resolution; set credentials before sharing across threads.
std::mutex mutex_;
// Guards lazy resolution and credential state.
mutable std::mutex mutex_;
std::vector<StorageCredential> storage_credentials_;
StorageCredentialRefresher refresher_;
std::unordered_map<std::string, std::unique_ptr<FileIO>, StringHash, StringEqual>
io_by_name_;
};
Expand Down
9 changes: 9 additions & 0 deletions src/iceberg/storage_credential.h
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,10 @@
/// \file iceberg/storage_credential.h
/// \brief Define storage credential metadata.

#include <functional>
#include <string>
#include <unordered_map>
#include <vector>

#include "iceberg/iceberg_export.h"
#include "iceberg/result.h"
Expand All @@ -48,4 +50,11 @@ struct ICEBERG_EXPORT StorageCredential {
bool operator==(const StorageCredential& other) const = default;
};

/// \brief Re-fetches the storage credentials that are currently valid.
///
/// Lets a FileIO replace expiring credentials without knowing how they are
/// delivered. Returns the whole vended list.
using StorageCredentialRefresher =
std::function<Result<std::vector<StorageCredential>>()>;

} // namespace iceberg
Loading
Loading