Skip to content

Commit 79fdb86

Browse files
feat(rest): load table credentials from the LoadCredentials endpoint (#899)
1 parent 85f8a84 commit 79fdb86

14 files changed

Lines changed: 362 additions & 29 deletions

‎src/iceberg/catalog/rest/json_serde.cc‎

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,25 @@ Result<StorageCredential> StorageCredentialFromJson(const nlohmann::json& json)
152152
return credential;
153153
}
154154

155+
/// \brief Reads the optional `storage-credentials` array shared by the
156+
/// LoadTable and LoadCredentials responses.
157+
Result<std::vector<StorageCredential>> StorageCredentialsFromJson(
158+
const nlohmann::json& json) {
159+
std::vector<StorageCredential> credentials;
160+
auto it = json.find(kStorageCredentials);
161+
if (it == json.end() || it->is_null()) {
162+
return credentials;
163+
}
164+
if (!it->is_array()) {
165+
return JsonParseError("Cannot parse storage credentials from non-array");
166+
}
167+
for (const auto& entry : *it) {
168+
ICEBERG_ASSIGN_OR_RAISE(auto credential, StorageCredentialFromJson(entry));
169+
credentials.push_back(std::move(credential));
170+
}
171+
return credentials;
172+
}
173+
155174
template <typename Value>
156175
Result<std::map<int32_t, Value>> KeyValueMapFromJson(const nlohmann::json& json,
157176
std::string_view key) {
@@ -738,19 +757,24 @@ Result<LoadTableResult> LoadTableResultFromJson(const nlohmann::json& json) {
738757
ICEBERG_ASSIGN_OR_RAISE(result.metadata, TableMetadataFromJson(metadata_json));
739758
ICEBERG_ASSIGN_OR_RAISE(result.config,
740759
GetJsonValueOrDefault<decltype(result.config)>(json, kConfig));
741-
if (auto it = json.find(kStorageCredentials); it != json.end() && !it->is_null()) {
742-
if (!it->is_array()) {
743-
return JsonParseError("Cannot parse storage credentials from non-array");
744-
}
745-
for (const auto& entry : *it) {
746-
ICEBERG_ASSIGN_OR_RAISE(auto cred, StorageCredentialFromJson(entry));
747-
result.storage_credentials.push_back(std::move(cred));
748-
}
749-
}
760+
ICEBERG_ASSIGN_OR_RAISE(result.storage_credentials, StorageCredentialsFromJson(json));
750761
ICEBERG_RETURN_UNEXPECTED(result.Validate());
751762
return result;
752763
}
753764

765+
Result<LoadCredentialsResponse> LoadCredentialsResponseFromJson(
766+
const nlohmann::json& json) {
767+
// Required here, unlike in LoadTable: reading a malformed response as "no
768+
// credentials" would look like a refresh that succeeded and dropped them.
769+
if (auto it = json.find(kStorageCredentials); it == json.end() || it->is_null()) {
770+
return JsonParseError("Missing '{}'", kStorageCredentials);
771+
}
772+
LoadCredentialsResponse response;
773+
ICEBERG_ASSIGN_OR_RAISE(response.storage_credentials, StorageCredentialsFromJson(json));
774+
ICEBERG_RETURN_UNEXPECTED(response.Validate());
775+
return response;
776+
}
777+
754778
nlohmann::json ToJson(const ListNamespacesResponse& response) {
755779
nlohmann::json json;
756780
SetOptionalStringField(json, kNextPageToken, response.next_page_token);

‎src/iceberg/catalog/rest/json_serde_internal.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,10 @@ template <>
7171
ICEBERG_REST_EXPORT Result<LoadTableResult> FromJson(const nlohmann::json& json);
7272
ICEBERG_REST_EXPORT Result<nlohmann::json> ToJson(const LoadTableResult& model);
7373

74+
// Response-only model: a client never serializes it, so no ToJson.
75+
ICEBERG_REST_EXPORT Result<LoadCredentialsResponse> LoadCredentialsResponseFromJson(
76+
const nlohmann::json& json);
77+
7478
ICEBERG_REST_EXPORT Result<CreateTableRequest> CreateTableRequestFromJson(
7579
const nlohmann::json& json);
7680
template <>

‎src/iceberg/catalog/rest/rest_catalog.cc‎

Lines changed: 80 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020
#include "iceberg/catalog/rest/rest_catalog.h"
2121

2222
#include <memory>
23+
#include <mutex>
24+
#include <string>
2325
#include <string_view>
2426
#include <tuple>
2527
#include <unordered_map>
@@ -42,6 +44,7 @@
4244
#include "iceberg/catalog/rest/rest_util.h"
4345
#include "iceberg/catalog/rest/types.h"
4446
#include "iceberg/json_serde_internal.h"
47+
#include "iceberg/logging/log_macros.h"
4548
#include "iceberg/metrics/metrics_reporters.h"
4649
#include "iceberg/partition_spec.h"
4750
#include "iceberg/result.h"
@@ -60,6 +63,41 @@ namespace iceberg::rest {
6063

6164
namespace {
6265

66+
class RestStorageCredentialProvider final : public StorageCredentialProvider {
67+
public:
68+
RestStorageCredentialProvider(std::shared_ptr<HttpClient> client,
69+
std::shared_ptr<auth::AuthSession> session,
70+
std::string path)
71+
: client_(std::move(client)),
72+
session_(std::move(session)),
73+
path_(std::move(path)) {}
74+
75+
Result<std::vector<StorageCredential>> Load() override {
76+
// The provider can be shared by several FileIOs; keep fetches serialized.
77+
std::lock_guard lock(mutex_);
78+
ICEBERG_ASSIGN_OR_RAISE(const auto response,
79+
client_->Get(path_, /*params=*/{}, /*headers=*/{},
80+
*TableErrorHandler::Instance(), *session_));
81+
// Parse errors can contain credential data; return a fixed message.
82+
auto json = FromJsonString(response.body());
83+
if (!json.has_value()) {
84+
return JsonParseError("Malformed LoadCredentials response");
85+
}
86+
auto result = LoadCredentialsResponseFromJson(*json);
87+
if (!result.has_value()) {
88+
return std::unexpected<Error>(
89+
{.kind = result.error().kind, .message = "Malformed LoadCredentials response"});
90+
}
91+
return std::move(result->storage_credentials);
92+
}
93+
94+
private:
95+
std::shared_ptr<HttpClient> client_;
96+
std::shared_ptr<auth::AuthSession> session_;
97+
std::string path_;
98+
std::mutex mutex_;
99+
};
100+
63101
/// \brief Get the default set of endpoints for backwards compatibility according to the
64102
/// iceberg rest spec.
65103
std::unordered_set<Endpoint> GetDefaultEndpoints() {
@@ -508,12 +546,45 @@ Result<std::shared_ptr<auth::AuthSession>> RestCatalog::TableAuthSession(
508546
std::move(contextual_session));
509547
}
510548

549+
std::shared_ptr<StorageCredentialProvider> RestCatalog::MakeStorageCredentialProvider(
550+
const TableIdentifier& identifier,
551+
std::shared_ptr<auth::AuthSession> table_session) const {
552+
if (!supported_endpoints_.contains(Endpoint::TableCredentials())) {
553+
// Not an error, but it surfaces much later as credentials expiring.
554+
ICEBERG_LOG_DEBUG(
555+
"Catalog does not advertise {}; vended credentials for '{}' will not be "
556+
"refreshed",
557+
Endpoint::TableCredentials().ToString(), ToString(identifier));
558+
return nullptr;
559+
}
560+
auto path = paths_->Credentials(identifier);
561+
if (!path.has_value()) {
562+
ICEBERG_LOG_WARN(
563+
"Cannot build the credentials path for '{}' ({}); its vended credentials "
564+
"will not be refreshed",
565+
ToString(identifier), path.error().message);
566+
return nullptr;
567+
}
568+
auto client = client_;
569+
auto credentials_path = std::move(path.value());
570+
auto session = std::move(table_session);
571+
return std::make_shared<RestStorageCredentialProvider>(
572+
std::move(client), std::move(session), std::move(credentials_path));
573+
}
574+
511575
Result<std::shared_ptr<FileIO>> RestCatalog::TableFileIO(
512-
const SessionContext& /*context*/,
576+
const SessionContext& /*context*/, const TableIdentifier& identifier,
513577
const std::unordered_map<std::string, std::string>& table_config,
514-
const std::vector<StorageCredential>& storage_credentials) const {
578+
const std::vector<StorageCredential>& storage_credentials,
579+
std::shared_ptr<auth::AuthSession> table_session) const {
515580
if (!table_config.empty() || !storage_credentials.empty()) {
516-
return MakeTableFileIO(config_.configs(), table_config, storage_credentials);
581+
// Only vended credentials expire, so only they need a provider.
582+
std::shared_ptr<StorageCredentialProvider> provider;
583+
if (!storage_credentials.empty()) {
584+
provider = MakeStorageCredentialProvider(identifier, std::move(table_session));
585+
}
586+
return MakeTableFileIO(config_.configs(), table_config, storage_credentials,
587+
std::move(provider));
517588
}
518589

519590
return file_io_;
@@ -772,11 +843,12 @@ Result<std::shared_ptr<Transaction>> RestCatalog::StageCreateTable(
772843
/*stage_create=*/true, *contextual_session));
773844
auto table_config = std::move(result.config);
774845
auto storage_credentials = std::move(result.storage_credentials);
775-
ICEBERG_ASSIGN_OR_RAISE(auto table_io,
776-
TableFileIO(context, table_config, storage_credentials));
846+
// Before the FileIO: refreshing its credentials reuses the table session.
777847
ICEBERG_ASSIGN_OR_RAISE(
778848
auto table_session,
779849
TableAuthSession(identifier, table_config, std::move(contextual_session)));
850+
ICEBERG_ASSIGN_OR_RAISE(auto table_io, TableFileIO(context, identifier, table_config,
851+
storage_credentials, table_session));
780852
ICEBERG_ASSIGN_OR_RAISE(auto reporter, MakeTableReporter(identifier, table_session));
781853
auto table_catalog = std::make_shared<TableScopedCatalog>(
782854
shared_from_this(), context, identifier, table_config, std::move(table_session),
@@ -890,11 +962,12 @@ Result<std::shared_ptr<Table>> RestCatalog::MakeTableFromLoadResult(
890962
std::shared_ptr<auth::AuthSession> contextual_session) {
891963
auto table_config = std::move(result.config);
892964
auto storage_credentials = std::move(result.storage_credentials);
893-
ICEBERG_ASSIGN_OR_RAISE(auto table_io,
894-
TableFileIO(context, table_config, storage_credentials));
965+
// Before the FileIO: refreshing its credentials reuses the table session.
895966
ICEBERG_ASSIGN_OR_RAISE(
896967
auto table_session,
897968
TableAuthSession(identifier, table_config, std::move(contextual_session)));
969+
ICEBERG_ASSIGN_OR_RAISE(auto table_io, TableFileIO(context, identifier, table_config,
970+
storage_credentials, table_session));
898971
ICEBERG_ASSIGN_OR_RAISE(auto reporter, MakeTableReporter(identifier, table_session));
899972
auto table_catalog = std::make_shared<TableScopedCatalog>(
900973
shared_from_this(), context, identifier, table_config, table_session, table_io);

‎src/iceberg/catalog/rest/rest_catalog.h‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -85,9 +85,17 @@ class ICEBERG_REST_EXPORT RestCatalog final
8585
std::shared_ptr<auth::AuthSession> contextual_session);
8686

8787
Result<std::shared_ptr<FileIO>> TableFileIO(
88-
const SessionContext& context,
88+
const SessionContext& context, const TableIdentifier& identifier,
8989
const std::unordered_map<std::string, std::string>& table_config,
90-
const std::vector<StorageCredential>& storage_credentials) const;
90+
const std::vector<StorageCredential>& storage_credentials,
91+
std::shared_ptr<auth::AuthSession> table_session) const;
92+
93+
/// \brief Build a provider for this table's vended credentials.
94+
///
95+
/// Returns nullptr when the catalog does not serve LoadCredentials.
96+
std::shared_ptr<StorageCredentialProvider> MakeStorageCredentialProvider(
97+
const TableIdentifier& identifier,
98+
std::shared_ptr<auth::AuthSession> table_session) const;
9199

92100
Result<std::vector<Namespace>> ListNamespaces(const Namespace& ns,
93101
auth::AuthSession& session) const;

‎src/iceberg/catalog/rest/rest_file_io.cc‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,13 @@
2121

2222
#include <string>
2323
#include <unordered_map>
24+
#include <utility>
2425
#include <vector>
2526

2627
#include "iceberg/catalog/rest/types.h"
2728
#include "iceberg/file_io.h"
2829
#include "iceberg/file_io_registry.h"
30+
#include "iceberg/logging/log_macros.h"
2931
#include "iceberg/resolving_file_io.h"
3032
#include "iceberg/util/macros.h"
3133

@@ -57,15 +59,25 @@ Result<std::unique_ptr<FileIO>> MakeCatalogFileIO(const RestCatalogProperties& c
5759
Result<std::unique_ptr<FileIO>> MakeTableFileIO(
5860
const std::unordered_map<std::string, std::string>& catalog_config,
5961
const std::unordered_map<std::string, std::string>& table_config,
60-
const std::vector<StorageCredential>& storage_credentials) {
62+
const std::vector<StorageCredential>& storage_credentials,
63+
std::shared_ptr<StorageCredentialProvider> provider) {
6164
const auto default_properties = MergeFileIOProperties(catalog_config, table_config);
6265
ICEBERG_ASSIGN_OR_RAISE(
6366
auto io, MakeCatalogFileIO(RestCatalogProperties::FromMap(default_properties)));
6467

6568
if (storage_credentials.empty()) {
6669
return io;
6770
} else if (auto* credentialed = io->AsSupportsStorageCredentials()) {
68-
ICEBERG_RETURN_UNEXPECTED(credentialed->SetStorageCredentials(storage_credentials));
71+
// First, so the FileIO never briefly holds credentials it cannot replace.
72+
auto status = credentialed->InitializeStorageCredentials(storage_credentials,
73+
std::move(provider));
74+
if (!status && !storage_credentials.empty() &&
75+
status.error().kind == ErrorKind::kNotSupported) {
76+
ICEBERG_LOG_WARN("Configured FileIO cannot refresh vended storage credentials: {}",
77+
status.error().message);
78+
status = credentialed->SetStorageCredentials(storage_credentials);
79+
}
80+
ICEBERG_RETURN_UNEXPECTED(status);
6981
} else {
7082
return NotSupported("Configured FileIO does not support vended storage credentials");
7183
}

‎src/iceberg/catalog/rest/rest_file_io.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ ICEBERG_REST_EXPORT Result<std::unique_ptr<FileIO>> MakeCatalogFileIO(
4444
ICEBERG_REST_EXPORT Result<std::unique_ptr<FileIO>> MakeTableFileIO(
4545
const std::unordered_map<std::string, std::string>& catalog_config,
4646
const std::unordered_map<std::string, std::string>& table_config,
47-
const std::vector<StorageCredential>& storage_credentials);
47+
const std::vector<StorageCredential>& storage_credentials,
48+
std::shared_ptr<StorageCredentialProvider> provider = nullptr);
4849

4950
} // namespace iceberg::rest

‎src/iceberg/catalog/rest/types.h‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,21 @@ using CreateTableResponse = LoadTableResult;
209209
/// \brief Alias of LoadTableResult used as the body of LoadTableResponse
210210
using LoadTableResponse = LoadTableResult;
211211

212+
/// \brief Response body of the LoadCredentials API.
213+
struct ICEBERG_REST_EXPORT LoadCredentialsResponse {
214+
std::vector<StorageCredential> storage_credentials;
215+
216+
/// \brief Validates the LoadCredentialsResponse.
217+
Status Validate() const {
218+
for (const auto& credential : storage_credentials) {
219+
ICEBERG_RETURN_UNEXPECTED(credential.Validate());
220+
}
221+
return {};
222+
}
223+
224+
bool operator==(const LoadCredentialsResponse& other) const = default;
225+
};
226+
212227
/// \brief Response body for listing namespaces.
213228
struct ICEBERG_REST_EXPORT ListNamespacesResponse {
214229
PageToken next_page_token;

‎src/iceberg/file_io.h‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
#include <span>
3030
#include <string>
3131
#include <string_view>
32+
#include <utility>
3233
#include <vector>
3334

3435
#include "iceberg/iceberg_export.h"
@@ -198,6 +199,18 @@ class ICEBERG_EXPORT SupportsStorageCredentials {
198199
/// By value because a concurrent install may replace them. An implementation
199200
/// that delegates may report what was installed on it.
200201
virtual std::vector<StorageCredential> credentials() const = 0;
202+
203+
/// \brief Install initial credentials and their refresh source before first use.
204+
///
205+
/// The default implementation supports static credentials only.
206+
virtual Status InitializeStorageCredentials(
207+
const std::vector<StorageCredential>& storage_credentials,
208+
std::shared_ptr<StorageCredentialProvider> provider) {
209+
if (provider) {
210+
return NotSupported("Credential refresh is not supported");
211+
}
212+
return SetStorageCredentials(storage_credentials);
213+
}
201214
};
202215

203216
} // namespace iceberg

0 commit comments

Comments
 (0)