Skip to content

Navigation Menu

Sign in
Sign up

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
plusplusjiajia wants to merge 3 commits into apache:main
base: main
Choose a base branch
Loading
from plusplusjiajia:feat-vended-credential-refresh
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
View file Open in desktop

Large diffs are not rendered by default.

3 changes: 3 additions & 0 deletions src/iceberg/arrow/s3/s3_properties.h
View file Open in desktop
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 MinIO, LocalStack, etc.)
Expand Down
42 changes: 33 additions & 9 deletions src/iceberg/catalog/rest/json_serde.cc
View file Open in desktop
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
View file Open in desktop
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
View file Open in desktop
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
View file Open in desktop
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
View file Open in desktop
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
View file Open in desktop
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
View file Open in desktop
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
View file Open in desktop
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
37 changes: 31 additions & 6 deletions src/iceberg/resolving_file_io.cc
View file Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,12 @@ Result<std::shared_ptr<FileIO>> ResolvingFileIO::FileIOForPath(
ICEBERG_ASSIGN_OR_RAISE(auto io, FileIORegistry::Load(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 @@ -103,14 +107,35 @@ Status ResolvingFileIO::SetStorageCredentials(
const std::vector<StorageCredential>& storage_credentials) {
// Rebuild delegates lazily with the new credentials. Updating live delegates
// instead would leave the resolver inconsistent if one of them rejected them.
std::unique_lock lock(mutex_);
storage_credentials_ = storage_credentials;
io_by_name_.clear();
// Retired outside the lock: tearing down a delegate can block.
decltype(io_by_name_) retired;
{
std::unique_lock lock(mutex_);
storage_credentials_ = storage_credentials;
retired.swap(io_by_name_);
}
return {};
}

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

void ResolvingFileIO::SetCredentialRefresher(StorageCredentialRefresher refresher) {
// Drop the cached delegates so they are rebuilt with the refresher. Retired
// outside the lock: teardown can block, and the outgoing callback's captures
// must not destruct under `mutex_`.
decltype(io_by_name_) retired;
// Holds the incoming callback going in and the outgoing one coming out; a
// pure swap never destroys a target under the lock, which std::exchange's
// move is permitted to do.
StorageCredentialRefresher handoff = std::move(refresher);
{
std::unique_lock lock(mutex_);
refresher_.swap(handoff);
retired.swap(io_by_name_);
}
}

} // namespace iceberg
20 changes: 17 additions & 3 deletions src/iceberg/resolving_file_io.h
View file Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,9 @@
namespace iceberg {

/// \brief FileIO that resolves and caches implementations by registry name.
///
/// Vended credentials and any refresher are forwarded to every resolved
/// implementation that supports them; each applies what it understands.
class ICEBERG_EXPORT ResolvingFileIO final : public FileIO,
public SupportsStorageCredentials {
public:
Expand All @@ -58,7 +61,17 @@ 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;

/// \brief Install a refresher and rebuild cached delegates with it.
///
/// A delegate dropped here is rebuilt from the credentials installed on this
/// resolver, which may be older than what it had refreshed itself to.
void SetCredentialRefresher(StorageCredentialRefresher refresher) override;

SupportsStorageCredentials* AsSupportsStorageCredentials() override { return this; }

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

std::unordered_map<std::string, std::string> properties_;
// Guards lazy resolution and credential refresh.
std::shared_mutex mutex_;
// Guards lazy resolution and credential state.
mutable std::shared_mutex mutex_;
std::vector<StorageCredential> storage_credentials_;
StorageCredentialRefresher refresher_;
std::unordered_map<std::string, std::shared_ptr<FileIO>, StringHash, StringEqual>
io_by_name_;
};
Expand Down
Loading

AltStyle によって変換されたページ (->オリジナル) /