From 06a1aa56659eb814c1e14ea70ab178f6954e2bd9 Mon Sep 17 00:00:00 2001 From: Andrey Zvonov <32552679+zvonand@users.noreply.github.com> Date: Thu, 30 Jul 2026 14:00:44 +0200 Subject: [PATCH 1/2] Cherry-pick of https://github.com/ClickHouse/ClickHouse/pull/107960 with unresolved conflict markers (resolution in next commit) --- docs/en/engines/database-engines/datalake.md | 1 + src/Common/ProfileEvents.cpp | 20 + src/Databases/DataLake/DatabaseDataLake.cpp | 28 + .../DataLake/DatabaseDataLakeSettings.cpp | 1 + src/Databases/DataLake/ICatalog.h | 28 + src/Databases/DataLake/RestCatalog.cpp | 586 ++++++++++++++++-- src/Databases/DataLake/RestCatalog.h | 92 ++- src/Databases/DataLake/StorageCredentials.h | 11 + .../AzureBlobStorage/AzureBlobStorageCommon.h | 4 + .../AzureBlobStorage/AzureObjectStorage.cpp | 138 +++-- .../AzureBlobStorage/AzureObjectStorage.h | 15 +- .../IO/ReadBufferFromAzureBlobStorage.cpp | 134 +++- src/Disks/IO/ReadBufferFromAzureBlobStorage.h | 12 +- .../IO/WriteBufferFromAzureBlobStorage.cpp | 75 ++- .../IO/WriteBufferFromAzureBlobStorage.h | 23 +- .../WriteBufferFromAzureDataLakeStorage.cpp | 26 +- .../IO/WriteBufferFromAzureDataLakeStorage.h | 13 +- .../isRetryableAzureException.cpp | 6 + .../isRetryableAzureException.h | 3 + .../ObjectStorage/Azure/Configuration.cpp | 31 +- ...ker_compose_iceberg_lakekeeper_catalog.yml | 4 +- .../test.py | 164 +++++ 22 files changed, 1291 insertions(+), 124 deletions(-) diff --git a/docs/en/engines/database-engines/datalake.md b/docs/en/engines/database-engines/datalake.md index b37fc38f790d..6bf01f0cf1b3 100644 --- a/docs/en/engines/database-engines/datalake.md +++ b/docs/en/engines/database-engines/datalake.md @@ -54,6 +54,7 @@ The following settings are supported: | `storage_endpoint` | Endpoint URL for the underlying storage | | `oauth_server_uri` | URI of the OAuth2 authorization server for authentication | | `vended_credentials` | Boolean indicating whether to use vended credentials from the catalog (supports AWS S3 and Azure ADLS Gen2) | +| `vended_credentials_cache_ttl` | Maximum cache entry lifetime (in seconds) for vended credentials (REST catalogs only). Default `300`; `0` disables caching. | | `aws_access_key_id` | AWS access key ID for S3/Glue access (if not using vended credentials) | | `aws_secret_access_key` | AWS secret access key for S3/Glue access (if not using vended credentials) | | `region` | AWS region for the service (e.g., `us-east-1`) | diff --git a/src/Common/ProfileEvents.cpp b/src/Common/ProfileEvents.cpp index 7f35ef869e92..59b79ae83811 100644 --- a/src/Common/ProfileEvents.cpp +++ b/src/Common/ProfileEvents.cpp @@ -1503,6 +1503,26 @@ The server successfully detected this situation and will download merged part fr M(AIRowsProcessed, "Number of rows that received an AI result.", ValueType::Number) \ M(AIRowsSkipped, "Number of rows that received a default value due to quota or error.", ValueType::Number) \ \ +<<<<<<< HEAD +======= + M(DataLakeRestCatalogCredentialsVended, "Number of table metadata requests to REST catalog asking to vend storage credentials.", ValueType::Number) \ + M(DataLakeRestCatalogCredentialsCacheHits, "Number of table metadata requests to REST catalog reusing cached storage credentials.", ValueType::Number) \ + \ + M(StatelessWorkerRequested, "Number of stateless workers requested by queries for distributed query execution.", ValueType::Number) \ + M(StatelessWorkerProvided, "Number of stateless workers provided to queries for distributed query execution.", ValueType::Number) \ + M(StatelessWorkerProvisioningMicroseconds, "Total time queries spent waiting for stateless workers to be provisioned.", ValueType::Microseconds) \ + M(StatelessWorkerProvisioningWaits, "Number of times a query had to wait for the stateless worker discovery service to provide workers.", ValueType::Number) \ + M(StatelessWorkerCreateLeaseRequests, "Number of create_lease requests sent to the stateless worker discovery service to acquire workers.", ValueType::Number) \ + M(StatelessWorkerCreateLeaseMicroseconds, "Total time spent in create_lease requests to the stateless worker discovery service.", ValueType::Microseconds) \ + M(StatelessWorkerCreateLeaseErrors, "Number of failed create_lease requests to the stateless worker discovery service (any error or timeout).", ValueType::Number) \ + M(StatelessWorkerGetLeaseRequests, "Number of get_lease requests sent to the stateless worker discovery service to refresh the list of leased workers.", ValueType::Number) \ + M(StatelessWorkerGetLeaseMicroseconds, "Total time spent in get_lease requests to the stateless worker discovery service.", ValueType::Microseconds) \ + M(StatelessWorkerGetLeaseErrors, "Number of failed get_lease requests to the stateless worker discovery service (any error or timeout).", ValueType::Number) \ + M(StatelessWorkerUpdateLeaseRequests, "Number of update_lease requests sent to the stateless worker discovery service to renew worker leases.", ValueType::Number) \ + M(StatelessWorkerUpdateLeaseMicroseconds, "Total time spent in update_lease requests to the stateless worker discovery service.", ValueType::Microseconds) \ + M(StatelessWorkerUpdateLeaseErrors, "Number of failed update_lease requests to the stateless worker discovery service (any error or timeout).", ValueType::Number) \ + \ +>>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) #ifdef APPLY_FOR_EXTERNAL_EVENTS #define APPLY_FOR_EVENTS(M) APPLY_FOR_BUILTIN_EVENTS(M) APPLY_FOR_EXTERNAL_EVENTS(M) diff --git a/src/Databases/DataLake/DatabaseDataLake.cpp b/src/Databases/DataLake/DatabaseDataLake.cpp index 12fbb051ba4a..f31cbf5dc22a 100644 --- a/src/Databases/DataLake/DatabaseDataLake.cpp +++ b/src/Databases/DataLake/DatabaseDataLake.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include #include @@ -63,6 +64,7 @@ namespace DatabaseDataLakeSetting extern const DatabaseDataLakeSettingsString oauth_server_uri; extern const DatabaseDataLakeSettingsBool oauth_server_use_request_body; extern const DatabaseDataLakeSettingsBool vended_credentials; + extern const DatabaseDataLakeSettingsUInt64 vended_credentials_cache_ttl; extern const DatabaseDataLakeSettingsString aws_access_key_id; extern const DatabaseDataLakeSettingsString aws_secret_access_key; extern const DatabaseDataLakeSettingsString region; @@ -318,6 +320,31 @@ std::shared_ptr DatabaseDataLake::getCatalog() const std::lock_guard lock(catalog_mutex); /// Lazily build the catalog on first access for databases attached at startup (see ctor). if (!catalog_impl) +<<<<<<< HEAD +======= + { + initializeOrLeaveUnavailable(); + if (!catalog_impl) + throw Exception( + ErrorCodes::ACCESS_DENIED, + "DataLakeCatalog database is inaccessible: its catalog uses server-managed credentials that are " + "restricted for user queries and could not be resolved when the database was loaded from metadata. " + "Provide explicit credentials, or enable `s3_allow_server_credentials_in_user_queries`. Reason: {}", + catalog_unavailable_reason); + } + + const auto settings_version = database_settings.get(); + catalog_impl->setVendedCredentialsCacheTTL( + std::chrono::seconds((*settings_version)[DatabaseDataLakeSetting::vended_credentials_cache_ttl].value)); + + return catalog_impl; +} + +void DatabaseDataLake::initializeOrLeaveUnavailable() const +{ + try + { +>>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) initialize(); return catalog_impl; } @@ -1208,6 +1235,7 @@ The following settings are supported: | `storage_endpoint` | Endpoint URL for the underlying storage | | `oauth_server_uri` | URI of the OAuth2 authorization server for authentication | | `vended_credentials` | Boolean indicating whether to use vended credentials from the catalog (supports AWS S3 and Azure ADLS Gen2) | +| `vended_credentials_cache_ttl` | Maximum cache entry lifetime (in seconds) for vended credentials (REST catalogs only). Default `300`; `0` disables caching. | | `aws_access_key_id` | AWS access key ID for S3/Glue access (if not using vended credentials) | | `aws_secret_access_key` | AWS secret access key for S3/Glue access (if not using vended credentials) | | `region` | AWS region for the service (e.g., `us-east-1`) | diff --git a/src/Databases/DataLake/DatabaseDataLakeSettings.cpp b/src/Databases/DataLake/DatabaseDataLakeSettings.cpp index 969b0769d13a..1b1b6e58fe1b 100644 --- a/src/Databases/DataLake/DatabaseDataLakeSettings.cpp +++ b/src/Databases/DataLake/DatabaseDataLakeSettings.cpp @@ -20,6 +20,7 @@ namespace ErrorCodes DECLARE(DatabaseDataLakeCatalogType, catalog_type, DatabaseDataLakeCatalogType::NONE, "Catalog type", 0) \ DECLARE(String, catalog_credential, "", "", 0) \ DECLARE(Bool, vended_credentials, true, "Use vended credentials (storage credentials) from catalog", 0) \ + DECLARE(UInt64, vended_credentials_cache_ttl, 300, "Maximum cache entry lifetime (in seconds) for vended credentials. '0' disables caching.", 0) \ DECLARE(String, auth_scope, "PRINCIPAL_ROLE:ALL", "Authorization scope for client credentials or token exchange", 0) \ DECLARE(String, oauth_server_uri, "", "OAuth server uri", 0) \ DECLARE(Bool, oauth_server_use_request_body, true, "Put parameters into request body or query params", 0) \ diff --git a/src/Databases/DataLake/ICatalog.h b/src/Databases/DataLake/ICatalog.h index e14b00ac3732..499c62dd331b 100644 --- a/src/Databases/DataLake/ICatalog.h +++ b/src/Databases/DataLake/ICatalog.h @@ -1,4 +1,5 @@ #pragma once +#include #include #include #include @@ -212,6 +213,33 @@ class ICatalog return std::nullopt; } +<<<<<<< HEAD +======= + virtual void setVendedCredentialsCacheTTL(std::chrono::seconds /*ttl*/) {} + + /// Result of `prepareSettingsChanges`: the new catalog state built off to the side, + /// ready to be published by `commitSettingsChanges`. + struct PreparedSettingsChanges + { + virtual ~PreparedSettingsChanges() = default; + }; + using PreparedSettingsChangesPtr = std::unique_ptr; + + /// Validate `ALTER DATABASE ... MODIFY SETTING` changes and build the new catalog + /// state without publishing anything (may throw, may do network I/O). The state + /// becomes visible only after `commitSettingsChanges`, so the caller can persist + /// the changes in between and abandon the prepared state on failure. + virtual PreparedSettingsChangesPtr prepareSettingsChanges(const DB::SettingsChanges & changes); + + /// Publish the state built by `prepareSettingsChanges`. Must not fail. + virtual void commitSettingsChanges(PreparedSettingsChangesPtr prepared); + + void applySettingsChanges(const DB::SettingsChanges & changes) + { + commitSettingsChanges(prepareSettingsChanges(changes)); + } + +>>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) protected: /// Name of the warehouse, /// which is sometimes also called "catalog name". diff --git a/src/Databases/DataLake/RestCatalog.cpp b/src/Databases/DataLake/RestCatalog.cpp index 28c1195082e4..321b3aec8056 100644 --- a/src/Databases/DataLake/RestCatalog.cpp +++ b/src/Databases/DataLake/RestCatalog.cpp @@ -2,6 +2,7 @@ #include #include #include +#include #include #include #include @@ -48,6 +49,15 @@ #include #include #include +<<<<<<< HEAD +======= +#include +#include +#include +#include +#include +#include +>>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) namespace DB::ErrorCodes @@ -68,6 +78,25 @@ namespace DB::FailPoints extern const char check_database_datalake_negative[]; } +<<<<<<< HEAD +======= +namespace ProfileEvents +{ + extern const Event DataLakeRestCatalogCredentialsVended; + extern const Event DataLakeRestCatalogCredentialsCacheHits; +} + +namespace DB::DatabaseDataLakeSetting +{ + extern const DatabaseDataLakeSettingsString catalog_credential; + extern const DatabaseDataLakeSettingsString auth_header; + extern const DatabaseDataLakeSettingsString onelake_tenant_id; + extern const DatabaseDataLakeSettingsString onelake_bearer_token; + extern const DatabaseDataLakeSettingsString onelake_client_id; + extern const DatabaseDataLakeSettingsString onelake_client_secret; +} + +>>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) namespace DataLake { @@ -77,6 +106,13 @@ static constexpr auto NAMESPACES_ENDPOINT = "namespaces"; namespace { +String parseTableUuid(const Poco::JSON::Object::Ptr & metadata_object) +{ + if (metadata_object && metadata_object->has("table-uuid")) + return metadata_object->get("table-uuid").extract(); + return {}; +} + std::pair parseCatalogCredential(const std::string & catalog_credential) { /// Parse a string of format ":" @@ -312,7 +348,229 @@ OneLakeCatalog::OneLakeCatalog( config = loadConfig(); } +<<<<<<< HEAD AccessToken RestCatalog::retrieveAccessToken() const +======= +void RestCatalog::validateSettingsChangesImpl( + const DB::SettingsChanges & changes, + const std::unordered_set & alterable_settings, + const std::string & auth_mode_description) +{ + for (const auto & change : changes) + { + if (alterable_settings.empty()) + throw DB::Exception( + DB::ErrorCodes::BAD_ARGUMENTS, + "Setting `{}` cannot be altered for a {}: the database was created without authentication settings", + change.name, + auth_mode_description); + + if (!alterable_settings.contains(change.name)) + throw DB::Exception( + DB::ErrorCodes::BAD_ARGUMENTS, + "Setting `{}` cannot be altered for a {} " + "(the authentication mode is fixed when the database is created; " + "alterable settings are: {})", + change.name, + auth_mode_description, + fmt::join(alterable_settings, ", ")); + + if (change.value.getType() != DB::Field::Types::String) + throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Setting `{}` must be a string", change.name); + + if (change.value.safeGet().empty()) + throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Setting `{}` cannot be set to an empty value", change.name); + } +} + +void RestCatalog::validateSettingsChanges(const DB::SettingsChanges & changes, bool credential_mode, bool header_mode) +{ + static const std::unordered_set credential_mode_settings = { + DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::catalog_credential)}; + static const std::unordered_set header_mode_settings = { + DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::auth_header)}; + static const std::unordered_set no_auth_settings = {}; + + if (credential_mode) + validateSettingsChangesImpl(changes, credential_mode_settings, "REST catalog with catalog credential authentication"); + else if (header_mode) + validateSettingsChangesImpl(changes, header_mode_settings, "REST catalog with auth header authentication"); + else + validateSettingsChangesImpl(changes, no_auth_settings, "REST catalog"); +} + +struct RestCatalog::PreparedAuthChanges : ICatalog::PreparedSettingsChanges +{ + std::unique_ptr new_state; + /// Set only when the OAuth credentials changed. + std::unique_ptr new_access_token; +}; + +ICatalog::PreparedSettingsChangesPtr RestCatalog::prepareSettingsChanges(const DB::SettingsChanges & changes) +{ + const auto old_state = state.get(); + CatalogState new_state = *old_state; + + auto prepared = std::make_unique(); + std::optional new_auth_headers; + applySettingsChangesToState(changes, *old_state, new_state, new_auth_headers, prepared->new_access_token); + + /// The config was loaded with the old credentials; the new ones may resolve the + /// warehouse to a different prefix or base location, so reload it before publishing. + new_state.config = loadConfig(new_state, new_auth_headers); + prepared->new_state = std::make_unique(std::move(new_state)); + return prepared; +} + +void RestCatalog::commitSettingsChanges(ICatalog::PreparedSettingsChangesPtr prepared) +{ + auto * prepared_auth = dynamic_cast(prepared.get()); + if (!prepared_auth || !prepared_auth->new_state) + throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Settings changes to commit were not prepared by this catalog"); + + state.set(std::move(prepared_auth->new_state)); + if (prepared_auth->new_access_token) + access_token.set(std::move(prepared_auth->new_access_token)); + + /// Cached storage credentials were vended under the old auth identity; do not reuse them. + { + std::lock_guard lock(credentials_cache_mutex); + credentials_cache.clear(); + } +} + +void RestCatalog::applySettingsChangesToState( + const DB::SettingsChanges & changes, + const CatalogState & old_state, + CatalogState & new_state, + std::optional & new_auth_headers, + std::unique_ptr & new_access_token) +{ + const bool credential_mode = !old_state.client_id.empty(); + const bool header_mode = old_state.auth_header.has_value(); + + validateSettingsChanges(changes, credential_mode, header_mode); + + for (const auto & change : changes) + { + if (change.name == DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::catalog_credential)) + { + std::tie(new_state.client_id, new_state.client_secret) = parseCatalogCredential(change.value.safeGet()); + } + else if (change.name == DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::auth_header)) + { + new_state.auth_header = parseAuthHeader(change.value.safeGet()); + validateAuthHeaders(new_state.auth_header.value()); + } + else + throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Unexpected setting `{}` after validation", change.name); + } + + if (credential_mode && (new_state.client_id != old_state.client_id || new_state.client_secret != old_state.client_secret)) + { + /// Eagerly fetch a token with the not-yet-published credentials: wrong credentials + /// fail the ALTER right here, and the config reload authenticates with that token + /// instead of the cached one. + new_access_token = std::make_unique(retrieveAccessToken(new_state.client_id, new_state.client_secret)); + new_auth_headers = DB::HTTPHeaderEntries{{"Authorization", "Bearer " + new_access_token->token}}; + } +} + +DB::HTTPHeaderEntries OneLakeCatalog::getAuthHeaders(const CatalogState & catalog_state, bool update_token) const +{ + auto headers = RestCatalog::getAuthHeaders(catalog_state, update_token); + headers.emplace_back("User-Agent", fmt::format("ClickHouse/{}{} OneLake-Catalog", VERSION_STRING, VERSION_OFFICIAL)); + return headers; +} + +void OneLakeCatalog::validateSettingsChanges(const DB::SettingsChanges & changes, bool bearer_mode) +{ + static const std::unordered_set bearer_mode_settings = { + DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_tenant_id), + DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_bearer_token)}; + static const std::unordered_set client_mode_settings = { + DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_tenant_id), + DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_client_id), + DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_client_secret)}; + + RestCatalog::validateSettingsChangesImpl( + changes, + bearer_mode ? bearer_mode_settings : client_mode_settings, + bearer_mode ? "OneLake catalog with bearer token authentication" : "OneLake catalog with client credentials authentication"); +} + +void OneLakeCatalog::applySettingsChangesToState( + const DB::SettingsChanges & changes, + const CatalogState & old_state, + CatalogState & new_state, + std::optional & new_auth_headers, + std::unique_ptr & new_access_token) +{ + const bool bearer_mode = !old_state.bearer_token.empty(); + + validateSettingsChanges(changes, bearer_mode); + + for (const auto & change : changes) + { + if (change.name == DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_tenant_id)) + new_state.tenant_id = change.value.safeGet(); + else if (change.name == DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_bearer_token)) + new_state.bearer_token = change.value.safeGet(); + else if (change.name == DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_client_id)) + new_state.client_id = change.value.safeGet(); + else if (change.name == DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_client_secret)) + new_state.client_secret = change.value.safeGet(); + else + throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Unexpected setting `{}` after validation", change.name); + } + + if (bearer_mode) + { + new_state.auth_header = DB::HTTPHeaderEntry("Authorization", "Bearer " + new_state.bearer_token); + validateAuthHeaders(new_state.auth_header.value()); + } + else if (new_state.client_id != old_state.client_id || new_state.client_secret != old_state.client_secret) + { + /// Eagerly fetch a token with the not-yet-published credentials: wrong credentials + /// fail the ALTER right here, and the config reload authenticates with that token + /// instead of the cached one. + new_access_token = std::make_unique(retrieveAccessToken(new_state.client_id, new_state.client_secret)); + new_auth_headers = DB::HTTPHeaderEntries{{"Authorization", "Bearer " + new_access_token->token}}; + } +} + +namespace +{ + +[[maybe_unused]] const bool rest_settings_alter_validator_registered = [] +{ + CatalogSettingsAlterValidatorFactory::instance().registerValidator( + DB::DatabaseDataLakeCatalogType::ICEBERG_REST, + [](const DB::DatabaseDataLakeSettings & current_settings, const DB::SettingsChanges & changes) + { + const bool credential_mode = !current_settings[DB::DatabaseDataLakeSetting::catalog_credential].value.empty(); + const bool header_mode = !current_settings[DB::DatabaseDataLakeSetting::auth_header].value.empty(); + RestCatalog::validateSettingsChanges(changes, credential_mode, header_mode); + }); + return true; +}(); + +[[maybe_unused]] const bool onelake_settings_alter_validator_registered = [] +{ + CatalogSettingsAlterValidatorFactory::instance().registerValidator( + DB::DatabaseDataLakeCatalogType::ICEBERG_ONELAKE, + [](const DB::DatabaseDataLakeSettings & current_settings, const DB::SettingsChanges & changes) + { + const bool bearer_mode = !current_settings[DB::DatabaseDataLakeSetting::onelake_bearer_token].value.empty(); + OneLakeCatalog::validateSettingsChanges(changes, bearer_mode); + }); + return true; +}(); + +} + +AccessToken RestCatalog::retrieveAccessToken(const std::string & client_id, const std::string & client_secret) const +>>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) { static constexpr auto oauth_tokens_endpoint = "oauth/tokens"; @@ -1024,23 +1282,118 @@ void RestCatalog::getTableMetadata( throw DB::Exception(DB::ErrorCodes::DATALAKE_DATABASE_ERROR, "No response from iceberg catalog"); } +namespace +{ + +/// Effective vended-credentials config of a LoadTableResult: "config" overlaid with the +/// "storage-credentials" entry whose prefix is the longest match of the location. +/// Returns nullptr when the response contains neither. +Poco::JSON::Object::Ptr effectiveVendedConfig(const Poco::JSON::Object::Ptr & load_table_result, const std::string & location) +{ + static constexpr auto storage_credentials_str = "storage-credentials"; + + Poco::JSON::Object::Ptr config_object; + if (load_table_result->has("config")) + { + config_object = load_table_result->getObject("config"); + if (!config_object) + throw DB::Exception(DB::ErrorCodes::DATALAKE_DATABASE_ERROR, "Cannot parse config result"); + } + + const auto entries + = load_table_result->isArray(storage_credentials_str) ? load_table_result->getArray(storage_credentials_str) : nullptr; + + Poco::JSON::Object::Ptr best_config; + size_t best_prefix_size = 0; + for (size_t i = 0; entries && i < entries->size(); ++i) + { + const auto entry = entries->getObject(static_cast(i)); + if (!entry) + continue; + const auto prefix_var = entry->get("prefix"); + if (!prefix_var.isString()) + continue; + const auto & prefix = prefix_var.extract(); + if (!location.starts_with(prefix)) + continue; + const auto entry_config = entry->getObject("config"); + if (!entry_config) + continue; + if (!best_config || prefix.size() > best_prefix_size) + { + best_config = entry_config; + best_prefix_size = prefix.size(); + } + } + + if (!best_config) + return config_object; + if (!config_object) + config_object = new Poco::JSON::Object(); + + Poco::JSON::Object::Ptr merged = new Poco::JSON::Object(*config_object); + std::vector names; + best_config->getNames(names); + + /// An entry that supplies any key of a credential group replaces that whole group. + static const std::vector> credential_groups = { + {"s3.access-key-id", "s3.secret-access-key", "s3.session-token", "s3.session-token-expires-at-ms"}, + {"gcs.oauth2.token", "gcs.oauth2.token-expires-at"}, + }; + for (const auto & group : credential_groups) + if (std::any_of(group.begin(), group.end(), [&](const auto & key) { return best_config->has(key); })) + for (const auto & key : group) + merged->remove(key); + + /// Azure SAS tokens form one group keyed by account name (adls.sas-token....). + static constexpr auto sas_prefix = "adls.sas-token."; + if (std::any_of(names.begin(), names.end(), [](const auto & name) { return name.starts_with(sas_prefix); })) + { + std::vector base_names; + merged->getNames(base_names); + for (const auto & name : base_names) + if (name.starts_with(sas_prefix)) + merged->remove(name); + } + + for (const auto & name : names) + merged->set(name, best_config->get(name)); + + return merged; +} + +} + bool RestCatalog::getTableMetadataImpl( const std::string & namespace_name, const std::string & table_name, - TableMetadata & result) const + TableMetadata & result, + bool allow_credentials_cache) const { LOG_DEBUG(log, "Checking table {} in namespace {}", table_name, namespace_name); DB::HTTPHeaderEntries headers; - if (result.requiresCredentials()) + + const bool want_credentials = result.requiresCredentials(); + + /// Reuse previously vended credentials is possible + std::optional cached_credentials; + if (want_credentials) { + if (allow_credentials_cache) + cached_credentials = tryGetCachedCredentials(namespace_name, table_name); + /// Header `X-Iceberg-Access-Delegation` tells catalog to include storage credentials in LoadTableResponse. /// Value can be one of the two: /// 1. `vended-credentials` /// 2. `remote-signing` /// Currently we support only the first. /// https://github.com/apache/iceberg/blob/3badfe0c1fcf0c0adfc7aa4a10f0b50365c48cf9/open-api/rest-catalog-open-api.yaml#L1832 - headers.emplace_back("X-Iceberg-Access-Delegation", "vended-credentials"); + if (!cached_credentials) + { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogCredentialsVended); + headers.emplace_back("X-Iceberg-Access-Delegation", "vended-credentials"); + } } const std::string endpoint = std::filesystem::path(NAMESPACES_ENDPOINT) / encodeNamespaceForURI(namespace_name) / "tables" / table_name; @@ -1069,6 +1422,8 @@ bool RestCatalog::getTableMetadataImpl( if (!metadata_object) throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Cannot parse result"); + const std::string table_uuid = parseTableUuid(metadata_object); + std::string location; if (result.requiresLocation()) { @@ -1094,16 +1449,36 @@ bool RestCatalog::getTableMetadataImpl( result.setSchema(*schema); } - if (result.isDefaultReadableTable() && result.requiresCredentials() && object->has("config")) + if (want_credentials && result.isDefaultReadableTable()) { - auto config_object = object->get("config").extract(); - if (!config_object) - throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Cannot parse config result"); - auto [parsed_credentials, parsed_endpoint] = getCredentialsAndEndpoint(config_object, location); - if (parsed_credentials) - result.setStorageCredentials(parsed_credentials); - if (!parsed_endpoint.empty()) - result.setEndpoint(parsed_endpoint); + if (cached_credentials) + { + /// Reuse the cached credentials only for the very same table with the very same UUID. + if (table_uuid.empty() || cached_credentials->table_uuid != table_uuid) + { + { + std::lock_guard lock(credentials_cache_mutex); + credentials_cache.erase({namespace_name, table_name}); + } + return getTableMetadataImpl(namespace_name, table_name, result, /* allow_credentials_cache */ false); + } + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogCredentialsCacheHits); + result.setStorageCredentials(cached_credentials->credentials); + if (!cached_credentials->endpoint.empty()) + result.setEndpoint(cached_credentials->endpoint); + } + else if (const auto config_object = effectiveVendedConfig(object, location)) + { + auto parsed = getCredentialsAndEndpoint(config_object, location); + parsed.table_uuid = table_uuid; + if (parsed.credentials) + { + result.setStorageCredentials(parsed.credentials); + cacheCredentials(namespace_name, table_name, parsed, state_snapshot); + } + if (!parsed.endpoint.empty()) + result.setEndpoint(parsed.endpoint); + } } if (result.requiresDataLakeSpecificProperties()) @@ -1115,8 +1490,8 @@ bool RestCatalog::getTableMetadataImpl( } } - if (metadata_object->has("table-uuid")) - result.setTableUUID(metadata_object->get("table-uuid").extract()); + if (!table_uuid.empty()) + result.setTableUUID(table_uuid); return true; } @@ -1374,7 +1749,69 @@ void RestCatalog::dropTable(const String & namespace_name, const String & table_ } } -std::pair, String> RestCatalog::getCredentialsAndEndpoint(Poco::JSON::Object::Ptr object, const String & location) const +namespace +{ +/// Parse a "...-expires-at-ms" value (ms since epoch); nullopt if absent, epoch (= don't cache) +/// if invalid; values beyond the representable range are clamped to the maximum time point. +std::optional +parseExpiresAtMs(const Poco::JSON::Object::Ptr & object, const std::string & key) +{ + if (!object->has(key)) + return std::nullopt; + try + { + static constexpr Int64 max_representable_sec + = std::chrono::duration_cast(std::chrono::system_clock::duration::max()).count(); + const Int64 expires_at_ms = object->get(key).convert(); + if (expires_at_ms <= 0) + return std::chrono::system_clock::time_point{}; + if (expires_at_ms / 1000 < max_representable_sec) + return std::chrono::system_clock::from_time_t(static_cast(expires_at_ms / 1000)); + return std::chrono::system_clock::time_point::max(); + } + catch (...) // NOLINT(bugprone-empty-catch) Ok: fail close below + { + } + return std::chrono::system_clock::time_point{}; +} + +std::chrono::system_clock::time_point parseSasTokenExpiry(const std::string & sas_token) +{ + std::string token = sas_token; + if (!token.empty() && token.front() == '?') + token.erase(0, 1); + + Poco::StringTokenizer params(token, "&", Poco::StringTokenizer::TOK_IGNORE_EMPTY | Poco::StringTokenizer::TOK_TRIM); + for (const auto & param : params) + { + if (!param.starts_with("se=")) + continue; + + try + { + std::string decoded; + Poco::URI::decode(param.substr(3), decoded); + + int time_zone_differential = 0; + Poco::DateTime date_time; + if (Poco::DateTimeParser::tryParse(Poco::DateTimeFormat::ISO8601_FORMAT, decoded, date_time, time_zone_differential)) + { + date_time.makeUTC(time_zone_differential); + return std::chrono::system_clock::from_time_t(date_time.timestamp().epochTime()); + } + } + catch (...) // NOLINT(bugprone-empty-catch) Ok: handled by the fail-close return below + { + } + + break; + } + /// Absent or unparseable 'se': do not cache. + return std::chrono::system_clock::time_point{}; +} +} + +VendedStorageCredentials RestCatalog::getCredentialsAndEndpoint(Poco::JSON::Object::Ptr object, const String & location) const { auto storage_type = parseStorageTypeFromLocation(location); switch (storage_type) @@ -1382,22 +1819,29 @@ std::pair, String> RestCatalog::getCredenti case StorageType::S3: { static constexpr auto gcs_token_str = "gcs.oauth2.token"; + static constexpr auto gcs_token_expires_at_str = "gcs.oauth2.token-expires-at"; static constexpr auto access_key_id_str = "s3.access-key-id"; static constexpr auto secret_access_key_str = "s3.secret-access-key"; static constexpr auto session_token_str = "s3.session-token"; static constexpr auto storage_endpoint_str = "s3.endpoint"; + static constexpr auto session_token_expires_at_ms_str = "s3.session-token-expires-at-ms"; - if (object->has(gcs_token_str)) + /// gs:// also maps to StorageType::S3, and the merged config may carry a warehouse-level + /// GCS token alongside per-table S3 keys, so select GCS by scheme, not by key presence. + if (location.starts_with("gs://") && object->has(gcs_token_str)) { auto gcs_token = object->get(gcs_token_str).extract(); LOG_DEBUG(log, "Using GCS OAuth2 token for location {}", location); - return {std::make_shared(gcs_token), ""}; + /// Do not cache if expiry was not parsed. + auto expires_at = parseExpiresAtMs(object, gcs_token_expires_at_str).value_or(std::chrono::system_clock::time_point{}); + return {std::make_shared(gcs_token), "", expires_at}; } std::string access_key_id; std::string secret_access_key; std::string session_token; std::string storage_endpoint; + std::optional expires_at; if (object->has(access_key_id_str)) access_key_id = object->get(access_key_id_str).extract(); if (object->has(secret_access_key_str)) @@ -1406,9 +1850,13 @@ std::pair, String> RestCatalog::getCredenti session_token = object->get(session_token_str).extract(); if (object->has(storage_endpoint_str)) storage_endpoint = object->get(storage_endpoint_str).extract(); + expires_at = parseExpiresAtMs(object, session_token_expires_at_ms_str); + /// Temporary credentials (session token) with unreported expiry must not be cached (fail close). + if (!expires_at.has_value() && !session_token.empty()) + expires_at = std::chrono::system_clock::time_point{}; LOG_DEBUG(log, "get tokens for location {}", location); - return {std::make_shared(access_key_id, secret_access_key, session_token), storage_endpoint}; + return {std::make_shared(access_key_id, secret_access_key, session_token), storage_endpoint, expires_at}; } case StorageType::Azure: { @@ -1430,15 +1878,70 @@ std::pair, String> RestCatalog::getCredenti } if (!sas_token.empty()) - { - return {std::make_shared(sas_token), ""}; - } + return {std::make_shared(sas_token), "", parseSasTokenExpiry(sas_token)}; break; } default: break; } - return {nullptr, ""}; + return {nullptr, "", std::nullopt}; +} + +std::optional RestCatalog::tryGetCachedCredentials( + const std::string & namespace_name, const std::string & table_name) const +{ + if (vended_credentials_cache_ttl.load(std::memory_order_relaxed) <= std::chrono::seconds::zero()) + return std::nullopt; + + std::lock_guard lock(credentials_cache_mutex); + auto it = credentials_cache.find({namespace_name, table_name}); + if (it == credentials_cache.end()) + return std::nullopt; + if (std::chrono::system_clock::now() >= it->second.expires_at.value()) + { + credentials_cache.erase(it); /// Drop the stale entry. + return std::nullopt; + } + + return it->second; +} + +void RestCatalog::cacheCredentials( + const std::string & namespace_name, + const std::string & table_name, + const VendedStorageCredentials & parsed, + const CatalogStateVersion & state_snapshot) const +{ + const auto ttl = vended_credentials_cache_ttl.load(std::memory_order_relaxed); + if (ttl <= std::chrono::seconds::zero()) + return; + + if (!parsed.credentials || parsed.credentials->isEmpty()) + return; + + const auto now = std::chrono::system_clock::now(); + + /// Cap at the configured TTL so an entry never outlives the documented maximum lifetime. + auto refresh_after = now + ttl; + if (parsed.expires_at) + { + const auto safe_expiry = parsed.expires_at.value() - credentials_expiry_safety_window; + if (safe_expiry < refresh_after) + refresh_after = safe_expiry; + } + if (refresh_after <= now) + return; + + std::lock_guard lock(credentials_cache_mutex); + + /// Do not cache credentials vended under an outdated auth state. + if (state.get() != state_snapshot) + return; + + if (credentials_cache.size() >= credentials_cache_cleanup_threshold) + std::erase_if(credentials_cache, [&now](const auto & entry) { return now >= entry.second.expires_at.value(); }); + credentials_cache[{namespace_name, table_name}] + = VendedStorageCredentials{parsed.credentials, parsed.endpoint, refresh_after, parsed.table_uuid}; } ICatalog::CredentialsRefreshCallback RestCatalog::getCredentialsConfigurationCallback(const DB::StorageID & storage_id) @@ -1468,32 +1971,33 @@ ICatalog::CredentialsRefreshCallback RestCatalog::getCredentialsConfigurationCal Poco::Dynamic::Var json = parser.parse(json_str); const Poco::JSON::Object::Ptr & object = json.extract(); - if (!object->has("config")) - { - LOG_DEBUG(log, "No 'config' in response for table {} – catalog does not support credential vending", table_name); - return nullptr; - } - - auto config_object = object->get("config").extract(); - if (!config_object) - { - LOG_DEBUG(log, "Empty 'config' in response for table {}", table_name); - return nullptr; - } + Poco::JSON::Object::Ptr metadata_object; + if (object->has("metadata")) + metadata_object = object->get("metadata").extract(); + /// Prefix matching uses the table location; metadata may live outside it with a custom `write.metadata.path`. std::string location; - if (object->has("metadata-location")) - { + if (metadata_object && metadata_object->has("location")) + location = metadata_object->get("location").extract(); + else if (object->has("metadata-location")) location = object->get("metadata-location").extract(); - LOG_DEBUG(log, "Location for table {}: {}", table_name, location); - } else + throw DB::Exception(DB::ErrorCodes::DATALAKE_DATABASE_ERROR, "Cannot read table {}, because no location in response", table_name); + LOG_DEBUG(log, "Location for table {}: {}", table_name, location); + + const auto config_object = effectiveVendedConfig(object, location); + if (!config_object) { - throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Cannot read table {}, because no 'metadata-location' in response", table_name); + LOG_DEBUG(log, "No credentials in response for table {} – catalog does not support credential vending", table_name); + return nullptr; } - auto [new_credentials, _] = getCredentialsAndEndpoint(config_object, location); - return new_credentials; + auto parsed = getCredentialsAndEndpoint(config_object, location); + if (metadata_object) + parsed.table_uuid = parseTableUuid(metadata_object); + /// Refresh the per-table cache so subsequent queries reuse these freshly vended credentials. + cacheCredentials(namespace_name, table_name, parsed, state_snapshot); + return parsed.credentials; }; } diff --git a/src/Databases/DataLake/RestCatalog.h b/src/Databases/DataLake/RestCatalog.h index 982475ee2c96..25035c61fba0 100644 --- a/src/Databases/DataLake/RestCatalog.h +++ b/src/Databases/DataLake/RestCatalog.h @@ -8,7 +8,17 @@ #include #include #include +#include +#include +#include #include +<<<<<<< HEAD +======= +#include +#include +#include +#include +>>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) #include namespace DB @@ -32,6 +42,14 @@ struct AccessToken } }; +struct VendedStorageCredentials +{ + std::shared_ptr credentials; + std::string endpoint; + std::optional expires_at; + std::string table_uuid = {}; +}; + class RestCatalog : public ICatalog, public DB::WithContext { public: @@ -87,8 +105,54 @@ class RestCatalog : public ICatalog, public DB::WithContext ICatalog::CredentialsRefreshCallback getCredentialsConfigurationCallback(const DB::StorageID & storage_id) override; +<<<<<<< HEAD String getClientId() const { return client_id; } String getClientSecret() const { return client_secret; } +======= + void setVendedCredentialsCacheTTL(std::chrono::seconds ttl) override { vended_credentials_cache_ttl.store(ttl, std::memory_order_relaxed); } + + struct Config + { + /// Prefix is a path of the catalog endpoint, + /// e.g. /v1/{prefix}/namespaces/{namespace}/tables/{table} + std::filesystem::path prefix; + /// Base location is location of data in storage + /// (in filesystem or object storage). + std::string default_base_location; + + std::string toString() const; + }; + + /// Credentials together with the catalog configuration they resolve to + /// (the /v1/config response depends on the credentials), published as one + /// atomic snapshot so readers never see a torn combination of them. + struct CatalogState + { + std::optional auth_header; + std::string client_id; + std::string client_secret; + std::string tenant_id; + std::string bearer_token; + Config config; + }; + using CatalogStateVersion = MultiVersion::Version; + + CatalogStateVersion getStateSnapshot() const { return state.get(); } + + ICatalog::PreparedSettingsChangesPtr prepareSettingsChanges(const DB::SettingsChanges & changes) override; + + void commitSettingsChanges(ICatalog::PreparedSettingsChangesPtr prepared) override; + + /// Check that we actually support these settings alter + static void validateSettingsChangesImpl( + const DB::SettingsChanges & changes, + const std::unordered_set & alterable_settings, + const std::string & auth_mode_description); + + /// `credential_mode` means the catalog authenticates with `catalog_credential`, + /// `header_mode` with `auth_header`. The mode is fixed when the database is created. + static void validateSettingsChanges(const DB::SettingsChanges & changes, bool credential_mode, bool header_mode); +>>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) protected: RestCatalog( @@ -131,6 +195,18 @@ class RestCatalog : public ICatalog, public DB::WithContext bool oauth_server_use_request_body; mutable MultiVersion access_token; + /// TTL for caching vended credentials per table (0 means no caching). + std::atomic vended_credentials_cache_ttl{std::chrono::seconds::zero()}; + + /// Sweep trigger threshold, not capacity! + static constexpr size_t credentials_cache_cleanup_threshold = 1000; + + static constexpr std::chrono::seconds credentials_expiry_safety_window{60}; + mutable std::mutex credentials_cache_mutex; + + mutable std::map, VendedStorageCredentials> credentials_cache + TSA_GUARDED_BY(credentials_cache_mutex); + Poco::Net::HTTPBasicCredentials credentials{}; DB::ReadWriteBufferFromHTTPPtr createReadBuffer( @@ -160,7 +236,8 @@ class RestCatalog : public ICatalog, public DB::WithContext bool getTableMetadataImpl( const std::string & namespace_name, const std::string & table_name, - TableMetadata & result) const; + TableMetadata & result, + bool allow_credentials_cache = true) const; Config loadConfig(); virtual DB::HTTPHeaderEntries getAuthHeaders(bool update_token) const; @@ -172,7 +249,18 @@ class RestCatalog : public ICatalog, public DB::WithContext const String & method = Poco::Net::HTTPRequest::HTTP_POST, bool ignore_result = false) const; - std::pair, String> getCredentialsAndEndpoint(Poco::JSON::Object::Ptr object, const String & location) const; + VendedStorageCredentials getCredentialsAndEndpoint(Poco::JSON::Object::Ptr object, const String & location) const; + + std::optional tryGetCachedCredentials( + const std::string & namespace_name, const std::string & table_name) const; + + /// `state_snapshot` is the state the credentials were vended under: if the auth settings + /// changed since, the entry is not inserted. + void cacheCredentials( + const std::string & namespace_name, + const std::string & table_name, + const VendedStorageCredentials & parsed, + const CatalogStateVersion & state_snapshot) const; AccessToken retrieveAccessToken() const; }; diff --git a/src/Databases/DataLake/StorageCredentials.h b/src/Databases/DataLake/StorageCredentials.h index 3a2f6f793e89..bed159e072b1 100644 --- a/src/Databases/DataLake/StorageCredentials.h +++ b/src/Databases/DataLake/StorageCredentials.h @@ -18,6 +18,9 @@ class IStorageCredentials virtual ~IStorageCredentials() = default; virtual void addCredentialsToEngineArgs(DB::ASTs & engine_args) const = 0; + + /// True when the credentials are unusable (mandatory fields empty); such credentials are not cached. + virtual bool isEmpty() const = 0; }; class S3Credentials final : public IStorageCredentials @@ -32,6 +35,8 @@ class S3Credentials final : public IStorageCredentials , session_token(session_token_) {} + bool isEmpty() const override { return access_key_id.empty() || secret_access_key.empty(); } + void addCredentialsToEngineArgs(DB::ASTs & engine_args) const override { if (engine_args.size() != 1) @@ -87,6 +92,8 @@ class GCSCredentials final : public IStorageCredentials DB::make_intrusive("Bearer " + oauth_token)))); } + bool isEmpty() const override { return oauth_token.empty(); } + const std::string & getToken() const { return oauth_token; } private: @@ -109,6 +116,10 @@ class AzureCredentials final : public IStorageCredentials engine_args.push_back(DB::make_intrusive(sas_token)); } + bool isEmpty() const override { return sas_token.empty(); } + + const std::string & getToken() const { return sas_token; } + private: std::string sas_token; }; diff --git a/src/Disks/DiskObjectStorage/ObjectStorages/AzureBlobStorage/AzureBlobStorageCommon.h b/src/Disks/DiskObjectStorage/ObjectStorages/AzureBlobStorage/AzureBlobStorageCommon.h index 0131c80ef5ec..8fb6b0f739d8 100644 --- a/src/Disks/DiskObjectStorage/ObjectStorages/AzureBlobStorage/AzureBlobStorageCommon.h +++ b/src/Disks/DiskObjectStorage/ObjectStorages/AzureBlobStorage/AzureBlobStorageCommon.h @@ -153,6 +153,10 @@ struct ConnectionParams std::unique_ptr createForContainer() const; }; +/// Both return an empty value if no fresh credentials are available. +using ConnectionParamsRefreshCallback = std::function()>; +using ContainerClientRefreshCallback = std::function()>; + void processURL(const String & url, const String & container_name, Endpoint & endpoint, AuthMethod & auth_method); std::unique_ptr getContainerClient(const ConnectionParams & params, bool readonly); diff --git a/src/Disks/DiskObjectStorage/ObjectStorages/AzureBlobStorage/AzureObjectStorage.cpp b/src/Disks/DiskObjectStorage/ObjectStorages/AzureBlobStorage/AzureObjectStorage.cpp index aac0ef6a094a..0e4b08c49545 100644 --- a/src/Disks/DiskObjectStorage/ObjectStorages/AzureBlobStorage/AzureObjectStorage.cpp +++ b/src/Disks/DiskObjectStorage/ObjectStorages/AzureBlobStorage/AzureObjectStorage.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #include @@ -121,6 +122,37 @@ class AzureIteratorAsync final : public IObjectStorageIteratorAsync Azure::Storage::Blobs::ListBlobsOptions options; }; +/// The refreshers below hold copies of all they need, so a buffer may outlive the storage that handed them out. +AzureBlobStorage::ContainerClientRefreshCallback makeContainerClientRefresher( + const AzureObjectStorage::AzureCredentialsRefreshCallback & credentials_refresh_callback) +{ + if (!credentials_refresh_callback) + return {}; + + return [credentials_refresh_callback]() -> std::unique_ptr + { + auto params = credentials_refresh_callback(); + if (!params) + return nullptr; + return params->createForContainer(); + }; +} + +WriteBufferFromAzureDataLakeStorage::FileClientRefreshCallback makeDataLakeFileClientRefresher( + const AzureObjectStorage::AzureCredentialsRefreshCallback & credentials_refresh_callback, const String & blob_path) +{ + if (!credentials_refresh_callback) + return {}; + + return [credentials_refresh_callback, blob_path]() -> std::optional + { + auto params = credentials_refresh_callback(); + if (!params) + return {}; + return makeAdlsGen2FileClient(params->endpoint, params->auth_method, params->client_options, blob_path); + }; +} + } @@ -132,7 +164,8 @@ AzureObjectStorage::AzureObjectStorage( const AzureBlobStorage::ConnectionParams & connection_params_, const String & object_namespace_, const String & description_, - const String & common_key_prefix_) + const String & common_key_prefix_, + AzureCredentialsRefreshCallback credentials_refresh_callback_) : name(name_) , auth_method(std::move(auth_method_)) , client(std::move(client_)) @@ -141,10 +174,25 @@ AzureObjectStorage::AzureObjectStorage( , description(description_) , common_key_prefix(common_key_prefix_) , connection_params(connection_params_) + , credentials_refresh_callback(std::move(credentials_refresh_callback_)) , log(getLogger("AzureObjectStorage")) { } +bool AzureObjectStorage::tryRefreshClient(const Azure::Core::RequestFailedException & e) const +{ + if (!credentials_refresh_callback || !isAzureAccessTokenExpiredError(e)) + return false; + + auto params = credentials_refresh_callback(); + if (!params) + return false; + + client.set(params->createForContainer()); + LOG_DEBUG(log, "Refreshed Azure credentials after an authentication failure"); + return true; +} + ObjectStorageKeyGeneratorPtr AzureObjectStorage::createKeyGenerator() const { return createObjectStorageKeyGeneratorByTemplate("[a-z]{32}"); @@ -152,23 +200,28 @@ ObjectStorageKeyGeneratorPtr AzureObjectStorage::createKeyGenerator() const bool AzureObjectStorage::exists(const StoredObject & object) const { - auto client_ptr = client.get(); + for (size_t attempt = 0; ; ++attempt) + { + auto client_ptr = client.get(); - ProfileEvents::increment(ProfileEvents::AzureGetProperties); - if (client_ptr->IsClientForDisk()) - ProfileEvents::increment(ProfileEvents::DiskAzureGetProperties); + ProfileEvents::increment(ProfileEvents::AzureGetProperties); + if (client_ptr->IsClientForDisk()) + ProfileEvents::increment(ProfileEvents::DiskAzureGetProperties); - try - { - auto blob_client = client_ptr->GetBlobClient(object.remote_path); - blob_client.GetProperties(); - return true; - } - catch (const Azure::Storage::StorageException & e) - { - if (e.StatusCode == Azure::Core::Http::HttpStatusCode::NotFound) - return false; - throw; + try + { + auto blob_client = client_ptr->GetBlobClient(object.remote_path); + blob_client.GetProperties(); + return true; + } + catch (const Azure::Storage::StorageException & e) + { + if (e.StatusCode == Azure::Core::Http::HttpStatusCode::NotFound) + return false; + if (attempt == 0 && tryRefreshClient(e)) + continue; + throw; + } } } @@ -251,7 +304,8 @@ std::unique_ptr AzureObjectStorage::readObject( /// NOLI restrict_seek, /* read_until_position */0, std::move(blob_storage_log), - connection_params.getContainer()); + connection_params.getContainer(), + makeContainerClientRefresher(credentials_refresh_callback)); } SmallObjectDataWithMetadata AzureObjectStorage::readSmallObjectAndGetObjectMetadata( /// NOLINT @@ -316,7 +370,8 @@ std::unique_ptr AzureObjectStorage::writeObject( /// NO patchSettings(write_settings), settings.get(), connection_params.getContainer(), - std::move(blob_storage_log)); + std::move(blob_storage_log), + makeDataLakeFileClientRefresher(credentials_refresh_callback, object.remote_path)); } return std::make_unique( @@ -327,7 +382,8 @@ std::unique_ptr AzureObjectStorage::writeObject( /// NO settings.get(), connection_params.getContainer(), std::move(blob_storage_log), - std::move(scheduler)); + std::move(scheduler), + makeContainerClientRefresher(credentials_refresh_callback)); } void AzureObjectStorage::removeObjectImpl( @@ -554,25 +610,37 @@ void AzureObjectStorage::tagObjects(const StoredObjects & objects, const std::st ObjectMetadata AzureObjectStorage::getObjectMetadata(const std::string & path, bool) const { - auto client_ptr = client.get(); - auto blob_client = client_ptr->GetBlobClient(path); - auto properties = blob_client.GetProperties().Value; + for (size_t attempt = 0; ; ++attempt) + { + auto client_ptr = client.get(); + try + { + auto blob_client = client_ptr->GetBlobClient(path); + auto properties = blob_client.GetProperties().Value; - ProfileEvents::increment(ProfileEvents::AzureGetProperties); - if (client_ptr->IsClientForDisk()) - ProfileEvents::increment(ProfileEvents::DiskAzureGetProperties); + ProfileEvents::increment(ProfileEvents::AzureGetProperties); + if (client_ptr->IsClientForDisk()) + ProfileEvents::increment(ProfileEvents::DiskAzureGetProperties); - ObjectMetadata result; - result.size_bytes = properties.BlobSize; - result.etag = properties.ETag.ToString(); - if (!properties.Metadata.empty()) - { - result.attributes.emplace(); - for (const auto & [key, value] : properties.Metadata) - result.attributes[key] = value; + ObjectMetadata result; + result.size_bytes = properties.BlobSize; + result.etag = properties.ETag.ToString(); + if (!properties.Metadata.empty()) + { + result.attributes.emplace(); + for (const auto & [key, value] : properties.Metadata) + result.attributes[key] = value; + } + result.last_modified = static_cast(properties.LastModified).time_since_epoch().count(); + return result; + } + catch (const Azure::Core::RequestFailedException & e) + { + if (attempt == 0 && tryRefreshClient(e)) + continue; + throw; + } } - result.last_modified = static_cast(properties.LastModified).time_since_epoch().count(); - return result; } std::optional AzureObjectStorage::tryGetObjectMetadata(const std::string & path, bool with_tags) const diff --git a/src/Disks/DiskObjectStorage/ObjectStorages/AzureBlobStorage/AzureObjectStorage.h b/src/Disks/DiskObjectStorage/ObjectStorages/AzureBlobStorage/AzureObjectStorage.h index 88adfe903284..0ea5d48edf4c 100644 --- a/src/Disks/DiskObjectStorage/ObjectStorages/AzureBlobStorage/AzureObjectStorage.h +++ b/src/Disks/DiskObjectStorage/ObjectStorages/AzureBlobStorage/AzureObjectStorage.h @@ -4,6 +4,7 @@ #if USE_AZURE_BLOB_STORAGE #include +#include #include #include #include @@ -25,6 +26,7 @@ class AzureObjectStorage : public IObjectStorage public: using ClientPtr = std::unique_ptr; using SettingsPtr = std::unique_ptr; + using AzureCredentialsRefreshCallback = AzureBlobStorage::ConnectionParamsRefreshCallback; AzureObjectStorage( const String & name_, @@ -34,7 +36,8 @@ class AzureObjectStorage : public IObjectStorage const AzureBlobStorage::ConnectionParams & connection_params_, const String & object_namespace_, const String & description_, - const String & common_key_prefix_); + const String & common_key_prefix_, + AzureCredentialsRefreshCallback credentials_refresh_callback_ = {}); void listObjects(const std::string & path, RelativePathsWithMetadata & children, size_t max_keys) const override; @@ -143,10 +146,13 @@ class AzureObjectStorage : public IObjectStorage std::unique_ptr buildDataLakeFileClient(const String & blob_path) const; + /// On an auth failure, swap in a client rebuilt with refreshed credentials; returns true to retry. + bool tryRefreshClient(const Azure::Core::RequestFailedException & e) const; + const String name; AzureBlobStorage::AuthMethod auth_method; - /// client used to access the files in the Blob Storage cloud - MultiVersion client; + /// client used to access the files in the Blob Storage cloud (mutable: tryRefreshClient swaps it). + mutable MultiVersion client; MultiVersion settings; const String object_namespace; /// container + prefix @@ -157,6 +163,9 @@ class AzureObjectStorage : public IObjectStorage const AzureBlobStorage::ConnectionParams connection_params; + /// Empty unless the storage was created for a catalog that vends refreshable credentials. + const AzureCredentialsRefreshCallback credentials_refresh_callback; + LoggerPtr log; }; diff --git a/src/Disks/IO/ReadBufferFromAzureBlobStorage.cpp b/src/Disks/IO/ReadBufferFromAzureBlobStorage.cpp index 9cbed6ada133..a1eab503523c 100644 --- a/src/Disks/IO/ReadBufferFromAzureBlobStorage.cpp +++ b/src/Disks/IO/ReadBufferFromAzureBlobStorage.cpp @@ -49,9 +49,11 @@ ReadBufferFromAzureBlobStorage::ReadBufferFromAzureBlobStorage( bool restricted_seek_, size_t read_until_position_, BlobStorageLogWriterPtr blob_storage_log_, - String container_for_logging_) + String container_for_logging_, + AzureClientRefreshCallback credentials_refresh_callback_) : ReadBufferFromFileBase() , blob_container_client(blob_container_client_) + , credentials_refresh_callback(std::move(credentials_refresh_callback_)) , path(path_) , max_single_read_retries(max_single_read_retries_) , max_single_download_retries(max_single_download_retries_) @@ -72,6 +74,36 @@ ReadBufferFromAzureBlobStorage::ReadBufferFromAzureBlobStorage( } } +std::pair +ReadBufferFromAzureBlobStorage::tryGetRefreshedClient(const Azure::Core::RequestFailedException & e) const +{ + if (!credentials_refresh_callback || !isAzureAccessTokenExpiredError(e)) + return {}; + + auto new_container = credentials_refresh_callback(); + if (!new_container) + return {}; + + BlobClientPtr new_blob = std::make_unique(new_container->GetBlobClient(path)); + return {std::move(new_container), std::move(new_blob)}; +} + +bool ReadBufferFromAzureBlobStorage::tryRefreshCredentials(const Azure::Core::RequestFailedException & e) +{ + if (credentials_refreshed) + return false; + + auto [new_container, new_blob] = tryGetRefreshedClient(e); + if (!new_container) + return false; + + blob_container_client = std::move(new_container); + blob_client = std::move(new_blob); + credentials_refreshed = true; + LOG_DEBUG(log, "Refreshed Azure credentials for {} after an authentication failure", path); + return true; +} + void ReadBufferFromAzureBlobStorage::setReadUntilEnd() { if (read_until_position) @@ -134,6 +166,14 @@ bool ReadBufferFromAzureBlobStorage::nextImpl() ProfileEvents::increment(ProfileEvents::ReadBufferFromAzureRequestsErrors); LOG_DEBUG(log, "Exception caught during Azure Read for file {} at attempt {}/{}: {}", path, i + 1, max_single_read_retries, e.Message); + if (tryRefreshCredentials(e)) + { + initialized = false; + initialize(i + 1); + --i; /// Don't count the refreshed retry against the budget (refresh happens at most once). + continue; + } + if (i + 1 == max_single_read_retries || !isRetryableAzureException(e)) throw; @@ -293,6 +333,12 @@ void ReadBufferFromAzureBlobStorage::initialize(size_t attempt) ProfileEvents::increment(ProfileEvents::ReadBufferFromAzureRequestsErrors); LOG_DEBUG(log, "Exception caught during Azure Download for file {} at offset {} at attempt {}/{}: {}", path, offset, i + 1, max_single_download_retries, e.Message); + if (tryRefreshCredentials(e)) + { + --i; /// Don't count the refreshed retry against the budget (refresh happens at most once). + continue; + } + if (i + 1 == max_single_download_retries || !isRetryableAzureException(e)) throw; @@ -335,24 +381,65 @@ void ReadBufferFromAzureBlobStorage::initialize(size_t attempt) std::optional ReadBufferFromAzureBlobStorage::tryGetFileSize() { - if (!blob_client) - blob_client = std::make_unique(blob_container_client->GetBlobClient(path)); + if (file_size) + return file_size; - if (!file_size) - file_size = blob_client->GetProperties().Value.BlobSize; + for (size_t attempt = 0; ; ++attempt) + { + if (!blob_client) + blob_client = std::make_unique(blob_container_client->GetBlobClient(path)); - return file_size; + try + { + file_size = blob_client->GetProperties().Value.BlobSize; + return file_size; + } + catch (const Azure::Core::RequestFailedException & e) + { + /// tryRefreshCredentials swaps blob_client for one built with fresh credentials, so the retry uses it. + if (attempt == 0 && tryRefreshCredentials(e)) + continue; + throw; + } + } } std::optional ReadBufferFromAzureBlobStorage::getRemoteFileMetadata() const { - const auto properties = blob_container_client->GetBlobClient(path).GetProperties().Value; - const auto last_modification_time = std::chrono::duration_cast( - static_cast(properties.LastModified).time_since_epoch()) - .count(); - return RemoteFileMetadata{ - .size = static_cast(properties.BlobSize), - .last_modification_time = static_cast(last_modification_time)}; + /// This method is const, so on an auth failure we refresh into a local client (as in readBigAt) + /// rather than mutating the shared members. + ContainerClientPtr refreshed_container_client; + BlobClientPtr refreshed_blob_client; + auto initial_blob_client = blob_container_client->GetBlobClient(path); + const AzureBlobStorage::BlobClient * current_blob_client = &initial_blob_client; + + for (size_t attempt = 0; ; ++attempt) + { + try + { + const auto properties = current_blob_client->GetProperties().Value; + const auto last_modification_time = std::chrono::duration_cast( + static_cast(properties.LastModified).time_since_epoch()) + .count(); + return RemoteFileMetadata{ + .size = static_cast(properties.BlobSize), + .last_modification_time = static_cast(last_modification_time)}; + } + catch (const Azure::Core::RequestFailedException & e) + { + if (attempt == 0) + { + if (auto [new_container, new_blob] = tryGetRefreshedClient(e); new_container) + { + refreshed_container_client = std::move(new_container); + refreshed_blob_client = std::move(new_blob); + current_blob_client = refreshed_blob_client.get(); + continue; + } + } + throw; + } + } } size_t ReadBufferFromAzureBlobStorage::readBigAt(char * to, size_t n, size_t range_begin, const std::function & /*progress_callback*/) const @@ -362,6 +449,11 @@ size_t ReadBufferFromAzureBlobStorage::readBigAt(char * to, size_t n, size_t ran ProfileEventTimeIncrement watch(ProfileEvents::ReadBufferFromAzureMicroseconds); + ContainerClientPtr refreshed_container_client; + BlobClientPtr refreshed_blob_client; + const AzureBlobStorage::BlobClient * current_blob_client = blob_client.get(); + bool credentials_refreshed_locally = false; + for (size_t i = 0; i < max_single_download_retries && n > 0; ++i) { size_t bytes_copied = 0; @@ -377,7 +469,7 @@ size_t ReadBufferFromAzureBlobStorage::readBigAt(char * to, size_t n, size_t ran download_options.Range = {static_cast(range_begin), n}; Azure::Core::Context azure_context = Azure::Core::Context().WithValue(PocoAzureHTTPClient::getSDKContextKeyForBufferRetry(), size_t{0}); - auto download_response = blob_client->Download(download_options, azure_context); + auto download_response = current_blob_client->Download(download_options, azure_context); if (blob_storage_log) { blob_storage_log->addEvent( @@ -413,6 +505,20 @@ size_t ReadBufferFromAzureBlobStorage::readBigAt(char * to, size_t n, size_t ran ProfileEvents::increment(ProfileEvents::ReadBufferFromAzureRequestsErrors); LOG_DEBUG(log, "Exception caught during Azure Download for file {} at offset {} at attempt {}/{}: {}", path, offset, i + 1, max_single_download_retries, e.Message); + if (!credentials_refreshed_locally) + { + if (auto [new_container, new_blob] = tryGetRefreshedClient(e); new_container) + { + refreshed_container_client = std::move(new_container); + refreshed_blob_client = std::move(new_blob); + current_blob_client = refreshed_blob_client.get(); + credentials_refreshed_locally = true; + LOG_DEBUG(log, "Refreshed Azure credentials for {} after an authentication failure in readBigAt", path); + --i; /// Don't count the refreshed retry against the budget. + continue; + } + } + if (i + 1 == max_single_download_retries || !isRetryableAzureException(e)) throw; diff --git a/src/Disks/IO/ReadBufferFromAzureBlobStorage.h b/src/Disks/IO/ReadBufferFromAzureBlobStorage.h index 81ee89ac1985..6b7d36d0ad18 100644 --- a/src/Disks/IO/ReadBufferFromAzureBlobStorage.h +++ b/src/Disks/IO/ReadBufferFromAzureBlobStorage.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include "config.h" @@ -23,6 +24,7 @@ class ReadBufferFromAzureBlobStorage : public ReadBufferFromFileBase public: using ContainerClientPtr = std::shared_ptr; using BlobClientPtr = std::unique_ptr; + using AzureClientRefreshCallback = AzureBlobStorage::ContainerClientRefreshCallback; ReadBufferFromAzureBlobStorage( ContainerClientPtr blob_container_client_, @@ -34,7 +36,8 @@ class ReadBufferFromAzureBlobStorage : public ReadBufferFromFileBase bool restricted_seek_ = false, size_t read_until_position_ = 0, BlobStorageLogWriterPtr blob_storage_log_ = {}, - String container_for_logging_ = {}); + String container_for_logging_ = {}, + AzureClientRefreshCallback credentials_refresh_callback_ = {}); off_t seek(off_t off, int whence) override; @@ -67,9 +70,16 @@ class ReadBufferFromAzureBlobStorage : public ReadBufferFromFileBase void initialize(size_t attempt); void setMetadataFromResponse(const Azure::Storage::Blobs::Models::DownloadBlobDetails & details, size_t blob_size) const; + std::pair tryGetRefreshedClient(const Azure::Core::RequestFailedException & e) const; + + /// On an auth failure, swap in refreshed credentials and retry (sequential path only, once per buffer). + bool tryRefreshCredentials(const Azure::Core::RequestFailedException & e); + std::unique_ptr data_stream; ContainerClientPtr blob_container_client; BlobClientPtr blob_client; + const AzureClientRefreshCallback credentials_refresh_callback; + bool credentials_refreshed = false; const String path; size_t max_single_read_retries; diff --git a/src/Disks/IO/WriteBufferFromAzureBlobStorage.cpp b/src/Disks/IO/WriteBufferFromAzureBlobStorage.cpp index 2102daa68969..28c2e81594c4 100644 --- a/src/Disks/IO/WriteBufferFromAzureBlobStorage.cpp +++ b/src/Disks/IO/WriteBufferFromAzureBlobStorage.cpp @@ -60,7 +60,8 @@ WriteBufferFromAzureBlobStorage::WriteBufferFromAzureBlobStorage( std::shared_ptr settings_, const String & container_for_logging_, BlobStorageLogWriterPtr blob_log_, - ThreadPoolCallbackRunnerUnsafe schedule_) + ThreadPoolCallbackRunnerUnsafe schedule_, + AzureBlobStorage::ContainerClientRefreshCallback credentials_refresh_callback_) : WriteBufferFromFileBase(std::min(buf_size_, static_cast(DBMS_DEFAULT_BUFFER_SIZE)), nullptr, 0) , log(getLogger("WriteBufferFromAzureBlobStorage")) , buffer_allocation_policy(createBufferAllocationPolicy(*settings_)) @@ -69,6 +70,7 @@ WriteBufferFromAzureBlobStorage::WriteBufferFromAzureBlobStorage( , blob_path(blob_path_) , write_settings(write_settings_) , blob_container_client(blob_container_client_) + , credentials_refresh_callback(std::move(credentials_refresh_callback_)) , task_tracker( std::make_unique( std::move(schedule_), @@ -112,20 +114,59 @@ WriteBufferFromAzureBlobStorage::~WriteBufferFromAzureBlobStorage() task_tracker->safeWaitAll(); } -void WriteBufferFromAzureBlobStorage::execWithRetry(std::function func, size_t num_tries, size_t cost) +WriteBufferFromAzureBlobStorage::AzureClientPtr WriteBufferFromAzureBlobStorage::getClient() const +{ + std::lock_guard lock(client_mutex); + return blob_container_client; +} + +bool WriteBufferFromAzureBlobStorage::tryRefreshCredentials( + const Azure::Core::RequestFailedException & e, const AzureClientPtr & used_client) +{ + if (!credentials_refresh_callback || !isAzureAccessTokenExpiredError(e)) + return false; + + std::lock_guard lock(client_mutex); + + /// Another part upload already refreshed the credentials while this attempt was in flight. + if (blob_container_client != used_client) + return true; + + if (credentials_refreshed) + return false; + + auto new_client = credentials_refresh_callback(); + if (!new_client) + return false; + + blob_container_client = std::move(new_client); + credentials_refreshed = true; + LOG_DEBUG(log, "Refreshed Azure credentials for blob `{}` after an authentication failure", blob_path); + return true; +} + +void WriteBufferFromAzureBlobStorage::execWithRetry(std::function func, size_t num_tries, size_t cost) { size_t sleep_time_with_backoff_milliseconds = 100; for (size_t i = 0; i < num_tries; ++i) { + auto client_ptr = getClient(); try { ResourceGuard rlock(ResourceGuard::Metrics::getIOWrite(), write_settings.io_scheduling.write_resource_link, cost); // Note that zero-cost requests are ignored - func(i); + func(i, client_ptr); rlock.unlock(cost); break; } catch (const Azure::Core::RequestFailedException & e) { + /// Credentials are refreshed at most once, so the retry after it costs no attempt. + if (tryRefreshCredentials(e, client_ptr)) + { + --i; + continue; + } + if (i == num_tries - 1 || !isRetryableAzureException(e, /* may_be_provisioning_access */ write_settings.is_initial_access_check)) throw; @@ -166,11 +207,9 @@ void WriteBufferFromAzureBlobStorage::preFinalize() if (block_ids.empty()) { ProfileEvents::increment(ProfileEvents::AzureUpload); - if (blob_container_client->IsClientForDisk()) + if (getClient()->IsClientForDisk()) ProfileEvents::increment(ProfileEvents::DiskAzureUpload); - auto block_blob_client = blob_container_client->GetBlockBlobClient(blob_path); - /// If there is only one block and size is less than or equal to max_single_part_upload_size /// then we use single part upload instead of multi part upload if (detached_part_data.size() == 1 && detached_part_data.front().data_size <= max_single_part_upload_size) @@ -185,7 +224,7 @@ void WriteBufferFromAzureBlobStorage::preFinalize() try { execWithRetry( - [&](size_t retry_attempt) + [&](size_t retry_attempt, const AzureClientPtr & client_ptr) { Azure::Storage::Blobs::UploadBlockBlobOptions options; @@ -195,7 +234,7 @@ void WriteBufferFromAzureBlobStorage::preFinalize() if (!write_settings.object_storage_write_if_match.empty()) options.AccessConditions.IfMatch = Azure::ETag(write_settings.object_storage_write_if_match); - block_blob_client.Upload( + client_ptr->GetBlockBlobClient(blob_path).Upload( memory_stream, options, azure_context.WithValue(PocoAzureHTTPClient::getSDKContextKeyForBufferRetry(), retry_attempt)); @@ -248,7 +287,7 @@ void WriteBufferFromAzureBlobStorage::preFinalize() try { execWithRetry( - [&](size_t retry_attempt) + [&](size_t retry_attempt, const AzureClientPtr & client_ptr) { Azure::Storage::Blobs::UploadBlockBlobOptions options; @@ -258,7 +297,7 @@ void WriteBufferFromAzureBlobStorage::preFinalize() if (!write_settings.object_storage_write_if_match.empty()) options.AccessConditions.IfMatch = Azure::ETag(write_settings.object_storage_write_if_match); - block_blob_client.Upload( + client_ptr->GetBlockBlobClient(blob_path).Upload( memory_stream, options, azure_context.WithValue(PocoAzureHTTPClient::getSDKContextKeyForBufferRetry(), retry_attempt)); @@ -317,9 +356,8 @@ void WriteBufferFromAzureBlobStorage::finalizeImpl() if (!block_ids.empty()) { - auto block_blob_client = blob_container_client->GetBlockBlobClient(blob_path); ProfileEvents::increment(ProfileEvents::AzureCommitBlockList); - if (blob_container_client->IsClientForDisk()) + if (getClient()->IsClientForDisk()) ProfileEvents::increment(ProfileEvents::DiskAzureCommitBlockList); Stopwatch watch; @@ -328,7 +366,7 @@ void WriteBufferFromAzureBlobStorage::finalizeImpl() try { execWithRetry( - [&](size_t retry_attetmpt) + [&](size_t retry_attetmpt, const AzureClientPtr & client_ptr) { Azure::Storage::Blobs::CommitBlockListOptions options; @@ -339,7 +377,7 @@ void WriteBufferFromAzureBlobStorage::finalizeImpl() options.AccessConditions.IfMatch = Azure::ETag(write_settings.object_storage_write_if_match); - block_blob_client.CommitBlockList( + client_ptr->GetBlockBlobClient(blob_path).CommitBlockList( block_ids, options, azure_context.WithValue(PocoAzureHTTPClient::getSDKContextKeyForBufferRetry(), retry_attetmpt)); @@ -382,7 +420,7 @@ void WriteBufferFromAzureBlobStorage::finalizeImpl() { try { - auto blob_client = blob_container_client->GetBlobClient(blob_path); + auto blob_client = getClient()->GetBlobClient(blob_path); blob_client.GetProperties(); } catch (const Azure::Storage::StorageException & e) @@ -505,10 +543,9 @@ void WriteBufferFromAzureBlobStorage::writePart(WriteBufferFromAzureBlobStorage: { auto & data_size = std::get<1>(*worker_data).data_size; auto & data_block_id = std::get<0>(*worker_data); - auto block_blob_client = blob_container_client->GetBlockBlobClient(blob_path); ProfileEvents::increment(ProfileEvents::AzureStageBlock); - if (blob_container_client->IsClientForDisk()) + if (getClient()->IsClientForDisk()) ProfileEvents::increment(ProfileEvents::DiskAzureStageBlock); Azure::Core::IO::MemoryBodyStream memory_stream(reinterpret_cast(std::get<1>(*worker_data).memory.data()), data_size); @@ -519,9 +556,9 @@ void WriteBufferFromAzureBlobStorage::writePart(WriteBufferFromAzureBlobStorage: try { execWithRetry( - [&](size_t retry_attempt) + [&](size_t retry_attempt, const AzureClientPtr & client_ptr) { - block_blob_client.StageBlock( + client_ptr->GetBlockBlobClient(blob_path).StageBlock( data_block_id, memory_stream, Azure::Storage::Blobs::StageBlockOptions{}, diff --git a/src/Disks/IO/WriteBufferFromAzureBlobStorage.h b/src/Disks/IO/WriteBufferFromAzureBlobStorage.h index 5ddda389d467..0e0d26a40966 100644 --- a/src/Disks/IO/WriteBufferFromAzureBlobStorage.h +++ b/src/Disks/IO/WriteBufferFromAzureBlobStorage.h @@ -5,7 +5,9 @@ #if USE_AZURE_BLOB_STORAGE #include +#include +#include #include #include #include @@ -39,7 +41,8 @@ class WriteBufferFromAzureBlobStorage : public WriteBufferFromFileBase std::shared_ptr settings_, const String & container_for_logging_ = {}, BlobStorageLogWriterPtr blob_log_ = {}, - ThreadPoolCallbackRunnerUnsafe schedule_ = {}); + ThreadPoolCallbackRunnerUnsafe schedule_ = {}, + AzureBlobStorage::ContainerClientRefreshCallback credentials_refresh_callback_ = {}); ~WriteBufferFromAzureBlobStorage() override; @@ -60,9 +63,17 @@ class WriteBufferFromAzureBlobStorage : public WriteBufferFromFileBase void setFakeBufferWhenPreFinalized(); void finalizeImpl() override; - void execWithRetry(std::function func, size_t num_tries, size_t cost = 0); + /// `func` gets the client for its attempt: a credentials refresh replaces it, so it must not be captured outside. + void execWithRetry(std::function func, size_t num_tries, size_t cost = 0); void uploadBlock(const char * data, size_t size); + AzureClientPtr getClient() const; + + /// On an auth failure, swap in a client rebuilt with refreshed credentials; returns true to retry. + /// `used_client` is the client of the failed attempt, so that an attempt racing with a refresh + /// done by another part upload retries with the fresh client instead of giving up. + bool tryRefreshCredentials(const Azure::Core::RequestFailedException & e, const AzureClientPtr & used_client); + /// Returns true if not a single byte was written to the buffer bool isEmpty() const { return total_size == 0 && count() == 0 && hidden_size == 0 && offset() == 0; } @@ -81,7 +92,13 @@ class WriteBufferFromAzureBlobStorage : public WriteBufferFromFileBase /// Track that prefinalize() is called only once bool is_prefinalized = false; - AzureClientPtr blob_container_client; + /// Part uploads run in parallel, so a refresh may replace the client while they are in flight. + mutable std::mutex client_mutex; + AzureClientPtr blob_container_client TSA_GUARDED_BY(client_mutex); + bool credentials_refreshed TSA_GUARDED_BY(client_mutex) = false; + + const AzureBlobStorage::ContainerClientRefreshCallback credentials_refresh_callback; + std::vector block_ids; using MemoryBufferPtr = std::unique_ptr>; diff --git a/src/Disks/IO/WriteBufferFromAzureDataLakeStorage.cpp b/src/Disks/IO/WriteBufferFromAzureDataLakeStorage.cpp index f165db0f843a..f0927b596d2a 100644 --- a/src/Disks/IO/WriteBufferFromAzureDataLakeStorage.cpp +++ b/src/Disks/IO/WriteBufferFromAzureDataLakeStorage.cpp @@ -109,13 +109,15 @@ WriteBufferFromAzureDataLakeStorage::WriteBufferFromAzureDataLakeStorage( const WriteSettings & write_settings_, std::shared_ptr settings_, const String & container_for_logging_, - BlobStorageLogWriterPtr blob_log_) + BlobStorageLogWriterPtr blob_log_, + FileClientRefreshCallback credentials_refresh_callback_) : WriteBufferFromFileBase(buf_size_, nullptr, 0) , log(getLogger("WriteBufferFromAzureDataLakeStorage")) , file_client(makeAdlsGen2FileClient(endpoint_, auth_method_, blob_client_options_, blob_path_)) , blob_path(blob_path_) , write_settings(write_settings_) , max_unexpected_write_error_retries(settings_->max_unexpected_write_error_retries) + , credentials_refresh_callback(std::move(credentials_refresh_callback_)) , container_for_logging(container_for_logging_) , blob_log(std::move(blob_log_)) { @@ -133,6 +135,21 @@ WriteBufferFromAzureDataLakeStorage::~WriteBufferFromAzureDataLakeStorage() } } +bool WriteBufferFromAzureDataLakeStorage::tryRefreshCredentials(const Azure::Core::RequestFailedException & e) +{ + if (credentials_refreshed || !credentials_refresh_callback || !isAzureAccessTokenExpiredError(e)) + return false; + + auto new_file_client = credentials_refresh_callback(); + if (!new_file_client) + return false; + + file_client = std::move(*new_file_client); + credentials_refreshed = true; + LOG_DEBUG(log, "Refreshed Azure credentials for `{}` after an authentication failure", blob_path); + return true; +} + void WriteBufferFromAzureDataLakeStorage::runWithRetries( const std::function & op, const char * what, @@ -167,6 +184,13 @@ void WriteBufferFromAzureDataLakeStorage::runWithRetries( } catch (const Azure::Core::RequestFailedException & e) { + /// Credentials are refreshed at most once, so the retry after it costs no attempt. + if (tryRefreshCredentials(e)) + { + --attempt; + continue; + } + const bool retryable = isRetryableAzureException(e, write_settings.is_initial_access_check); if (!retryable || attempt >= max_unexpected_write_error_retries) { diff --git a/src/Disks/IO/WriteBufferFromAzureDataLakeStorage.h b/src/Disks/IO/WriteBufferFromAzureDataLakeStorage.h index 6e1844b1d6b9..f5f86a896da6 100644 --- a/src/Disks/IO/WriteBufferFromAzureDataLakeStorage.h +++ b/src/Disks/IO/WriteBufferFromAzureDataLakeStorage.h @@ -5,6 +5,7 @@ #if USE_AZURE_BLOB_STORAGE #include +#include #include #include @@ -21,6 +22,9 @@ namespace DB class WriteBufferFromAzureDataLakeStorage : public WriteBufferFromFileBase { public: + /// Returns a file client built with refreshed credentials, or an empty value if none are available. + using FileClientRefreshCallback = std::function()>; + WriteBufferFromAzureDataLakeStorage( const AzureBlobStorage::Endpoint & endpoint_, const AzureBlobStorage::AuthMethod & auth_method_, @@ -30,7 +34,8 @@ class WriteBufferFromAzureDataLakeStorage : public WriteBufferFromFileBase const WriteSettings & write_settings_, std::shared_ptr settings_, const String & container_for_logging_ = {}, - BlobStorageLogWriterPtr blob_log_ = {}); + BlobStorageLogWriterPtr blob_log_ = {}, + FileClientRefreshCallback credentials_refresh_callback_ = {}); ~WriteBufferFromAzureDataLakeStorage() override; @@ -50,6 +55,9 @@ class WriteBufferFromAzureDataLakeStorage : public WriteBufferFromFileBase BlobStorageLogElement::EventType event_type, size_t data_size); + /// On an auth failure, swap in a file client rebuilt with refreshed credentials; returns true to retry. + bool tryRefreshCredentials(const Azure::Core::RequestFailedException & e); + LoggerPtr log; Azure::Storage::Files::DataLake::DataLakeFileClient file_client; @@ -57,6 +65,9 @@ class WriteBufferFromAzureDataLakeStorage : public WriteBufferFromFileBase const WriteSettings write_settings; const size_t max_unexpected_write_error_retries; + const FileClientRefreshCallback credentials_refresh_callback; + bool credentials_refreshed = false; + bool file_created = false; bool is_prefinalized = false; int64_t bytes_appended = 0; diff --git a/src/IO/AzureBlobStorage/isRetryableAzureException.cpp b/src/IO/AzureBlobStorage/isRetryableAzureException.cpp index 83d20065e284..dd023ae0ae59 100644 --- a/src/IO/AzureBlobStorage/isRetryableAzureException.cpp +++ b/src/IO/AzureBlobStorage/isRetryableAzureException.cpp @@ -21,6 +21,12 @@ bool isRetryableAzureException(const Azure::Core::RequestFailedException & e, bo return e.StatusCode >= Azure::Core::Http::HttpStatusCode::InternalServerError; } +bool isAzureAccessTokenExpiredError(const Azure::Core::RequestFailedException & e) +{ + return e.StatusCode == Azure::Core::Http::HttpStatusCode::Unauthorized + || e.StatusCode == Azure::Core::Http::HttpStatusCode::Forbidden; +} + } #endif diff --git a/src/IO/AzureBlobStorage/isRetryableAzureException.h b/src/IO/AzureBlobStorage/isRetryableAzureException.h index b1bea9331169..7967a95e8d3b 100644 --- a/src/IO/AzureBlobStorage/isRetryableAzureException.h +++ b/src/IO/AzureBlobStorage/isRetryableAzureException.h @@ -9,6 +9,9 @@ namespace DB bool isRetryableAzureException(const Azure::Core::RequestFailedException & e, bool may_be_provisioning_access = false); +/// True for HTTP 401/403, which indicate the SAS token (or other credentials) is rejected or expired. +bool isAzureAccessTokenExpiredError(const Azure::Core::RequestFailedException & e); + } #endif diff --git a/src/Storages/ObjectStorage/Azure/Configuration.cpp b/src/Storages/ObjectStorage/Azure/Configuration.cpp index a5930cc25015..3fb98712856c 100644 --- a/src/Storages/ObjectStorage/Azure/Configuration.cpp +++ b/src/Storages/ObjectStorage/Azure/Configuration.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -89,13 +90,38 @@ StorageObjectStorageQuerySettings StorageAzureConfiguration::getQuerySettings(co }; } -ObjectStoragePtr StorageAzureConfiguration::createObjectStorage(ContextPtr context, bool is_readonly, CredentialsConfigurationCallback /*refresh_credentials_callback*/) /// NOLINT +ObjectStoragePtr StorageAzureConfiguration::createObjectStorage(ContextPtr context, bool is_readonly, CredentialsConfigurationCallback refresh_credentials_callback) /// NOLINT { assertInitialized(); auto settings = AzureBlobStorage::getRequestSettings(context->getSettingsRef()); auto client = AzureBlobStorage::getContainerClient(connection_params, is_readonly); + /// For catalogs that vend refreshable SAS tokens, rebuild the clients with a fresh token on auth failure. + AzureObjectStorage::AzureCredentialsRefreshCallback credentials_refresher; + if (refresh_credentials_callback) + { + credentials_refresher = [refresh_credentials_callback, params = connection_params]() -> std::optional + { + auto new_credentials = (*refresh_credentials_callback)(); + if (!new_credentials) + return {}; + + auto azure_credentials = std::dynamic_pointer_cast(new_credentials); + if (!azure_credentials) + throw Exception(ErrorCodes::BAD_ARGUMENTS, "Unexpected credentials type for Azure storage"); + if (azure_credentials->isEmpty()) + return {}; + + auto new_params = params; + std::string sas = azure_credentials->getToken(); + if (!sas.empty() && sas.front() == '?') + sas.erase(0, 1); + new_params.endpoint.sas_auth = std::move(sas); + return new_params; + }; + } + return std::make_unique( "AzureBlobStorage", connection_params.auth_method, @@ -104,7 +130,8 @@ ObjectStoragePtr StorageAzureConfiguration::createObjectStorage(ContextPtr conte connection_params, connection_params.getContainer(), connection_params.getConnectionURL(), - /*common_key_prefix*/ ""); + /*common_key_prefix*/ "", + std::move(credentials_refresher)); } AzureBlobStorage::ConnectionParams getAzureConnectionParams( diff --git a/tests/integration/compose/docker_compose_iceberg_lakekeeper_catalog.yml b/tests/integration/compose/docker_compose_iceberg_lakekeeper_catalog.yml index 834dfcf071c9..85fa53adeea4 100644 --- a/tests/integration/compose/docker_compose_iceberg_lakekeeper_catalog.yml +++ b/tests/integration/compose/docker_compose_iceberg_lakekeeper_catalog.yml @@ -1,6 +1,6 @@ services: lakekeeper: - image: vakamo/lakekeeper:v0.9.4 + image: vakamo/lakekeeper:v0.13.1 environment: - LAKEKEEPER__PG_ENCRYPTION_KEY=This-is-NOT-Secure! - LAKEKEEPER__PG_DATABASE_URL_READ=postgresql://postgres:postgres@db:5432/postgres @@ -23,7 +23,7 @@ services: cpus: 3 migrate: - image: vakamo/lakekeeper:v0.9.4 + image: vakamo/lakekeeper:v0.13.1 environment: - LAKEKEEPER__PG_ENCRYPTION_KEY=This-is-NOT-Secure! - LAKEKEEPER__PG_DATABASE_URL_READ=postgresql://postgres:postgres@db:5432/postgres diff --git a/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py b/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py index 01b25fd75f86..47b57a014bed 100644 --- a/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py +++ b/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py @@ -9,6 +9,7 @@ from pyiceberg.schema import Schema from pyiceberg.types import ( DoubleType, + IntegerType, NestedField, StringType, ) @@ -399,3 +400,166 @@ def test_invalid_auth_header_format(started_cluster): ) assert "Invalid auth header format" in str(err.value) + +def get_credentials_profile_events(node, query_id): + node.query("SYSTEM FLUSH LOGS") + vended = int(node.query( + f"SELECT ProfileEvents['DataLakeRestCatalogCredentialsVended'] " + f"FROM system.query_log WHERE query_id = '{query_id}' AND type = 'QueryFinish'" + )) + hits = int(node.query( + f"SELECT ProfileEvents['DataLakeRestCatalogCredentialsCacheHits'] " + f"FROM system.query_log WHERE query_id = '{query_id}' AND type = 'QueryFinish'" + )) + return vended, hits + + +def create_int_table(catalog, namespace, table_name, rows=1): + if namespace not in catalog.list_namespaces(): + catalog.create_namespace(namespace) + schema = Schema( + NestedField(field_id=1, name="id", field_type=IntegerType(), required=False), + NestedField(field_id=2, name="data", field_type=StringType(), required=False), + ) + table = catalog.create_table( + namespace + (table_name,), + schema=schema, + properties={"write.metadata.compression-codec": "none"}, + ) + table.append( + pa.Table.from_pandas( + pd.DataFrame({"id": list(range(rows)), "data": ["x"] * rows}).astype( + {"id": "int32"} + ) + ) + ) + + +def test_vended_credentials_cache(started_cluster): + node = started_cluster.instances["node1"] + catalog = load_catalog_impl(started_cluster) + + test_ref = f"test_vended_credentials_cache_{uuid.uuid4().hex[:8]}" + namespace = (f"{test_ref}_namespace",) + table_name = f"{test_ref}_table" + db_name = f"{test_ref}_database" + + create_int_table(catalog, namespace, table_name) + + query = f"SELECT count() FROM {db_name}.`{namespace[0]}.{table_name}`" + + # Caching enabled (default TTL): the second query reuses cached credentials + # and does not ask the catalog to vend them again. + create_clickhouse_iceberg_database(started_cluster, node, db_name) + + qid = f"{test_ref}-cache-1-{uuid.uuid4()}" + node.query(query, query_id=qid) + vended, _ = get_credentials_profile_events(node, qid) + assert vended >= 1 + + qid = f"{test_ref}-cache-2-{uuid.uuid4()}" + node.query(query, query_id=qid) + vended, hits = get_credentials_profile_events(node, qid) + assert vended == 0 and hits >= 1 + + # Caching disabled (TTL = 0): every query asks the catalog to vend credentials. + create_clickhouse_iceberg_database( + started_cluster, node, db_name, + additional_settings={"vended_credentials_cache_ttl": 0}, + ) + + qid = f"{test_ref}-nocache-1-{uuid.uuid4()}" + node.query(query, query_id=qid) + vended, hits = get_credentials_profile_events(node, qid) + assert vended >= 1 and hits == 0 + + qid = f"{test_ref}-nocache-2-{uuid.uuid4()}" + node.query(query, query_id=qid) + vended, hits = get_credentials_profile_events(node, qid) + assert vended >= 1 and hits == 0 + + +def test_vended_credentials_cache_invalidated_on_table_replace(started_cluster): + node = started_cluster.instances["node1"] + catalog = load_catalog_impl(started_cluster) + + test_ref = f"test_vended_credentials_cache_replace_{uuid.uuid4().hex[:8]}" + namespace = (f"{test_ref}_namespace",) + table_name = f"{test_ref}_table" + db_name = f"{test_ref}_database" + + create_int_table(catalog, namespace, table_name) + create_clickhouse_iceberg_database(started_cluster, node, db_name) + query = f"SELECT count() FROM {db_name}.`{namespace[0]}.{table_name}`" + + # Populate the cache, then confirm the next query reuses it. + node.query(query, query_id=f"{test_ref}-1-{uuid.uuid4()}") + qid = f"{test_ref}-2-{uuid.uuid4()}" + node.query(query, query_id=qid) + vended, hits = get_credentials_profile_events(node, qid) + assert vended == 0 and hits >= 1 + + # Replace the table (new UUID and location) while the cache entry is still valid. + catalog.drop_table(namespace + (table_name,)) + create_int_table(catalog, namespace, table_name, rows=2) + + # The stale entry must be detected, so credentials are re-vended and the new table is read. + qid = f"{test_ref}-3-{uuid.uuid4()}" + assert node.query(query, query_id=qid).strip() == "2" + vended, _ = get_credentials_profile_events(node, qid) + assert vended >= 1 + + # The re-vended credentials are cached under the new identity, so the next query reuses them. + qid = f"{test_ref}-4-{uuid.uuid4()}" + node.query(query, query_id=qid) + vended, hits = get_credentials_profile_events(node, qid) + assert vended == 0 and hits >= 1 + + +def test_vended_credentials_cache_cleared_on_auth_change(started_cluster): + node = started_cluster.instances["node1"] + catalog = load_catalog_impl(started_cluster) + + test_ref = f"test_vended_credentials_cache_auth_{uuid.uuid4().hex[:8]}" + namespace = (f"{test_ref}_namespace",) + table_name = f"{test_ref}_table" + db_name = f"{test_ref}_database" + + create_int_table(catalog, namespace, table_name) + + # Header-mode database: `auth_header` is the only auth setting that can be altered. + node.query(f"DROP DATABASE IF EXISTS {db_name}") + node.query( + f""" + CREATE DATABASE {db_name} + ENGINE = DataLakeCatalog('{BASE_URL}') + SETTINGS + catalog_type = 'rest', + warehouse = 'demo', + storage_endpoint = 'http://minio1:9001/warehouse-rest', + auth_header = 'Authorization: Bearer initial_dummy' + """, + settings={"allow_experimental_database_iceberg": 1}, + ) + + query = f"SELECT count() FROM {db_name}.`{namespace[0]}.{table_name}`" + + # Populate the cache and confirm the next query reuses it. + node.query(query, query_id=f"{test_ref}-1-{uuid.uuid4()}") + qid = f"{test_ref}-2-{uuid.uuid4()}" + node.query(query, query_id=qid) + vended, hits = get_credentials_profile_events(node, qid) + assert vended == 0 and hits >= 1 + + # Committing new auth settings must drop the cached credentials. + node.query( + f"ALTER DATABASE {db_name} MODIFY SETTING auth_header = 'Authorization: Bearer altered_dummy'" + ) + + qid = f"{test_ref}-3-{uuid.uuid4()}" + node.query(query, query_id=qid) + vended, _ = get_credentials_profile_events(node, qid) + assert vended >= 1 + + node.query(f"DROP DATABASE IF EXISTS {db_name}") + From 9bd22cad5988909dcc12239570fa7a23c219092a Mon Sep 17 00:00:00 2001 From: Andrey Zvonov <32552679+zvonand@users.noreply.github.com> Date: Mon, 3 Aug 2026 19:45:23 +0200 Subject: [PATCH 2/2] Resolve conflicts in cherry-pick of #107960 Kept only the source PR's own additions in the conflicted regions; the surrounding "theirs" context (stateless-worker ProfileEvents, the `initializeOrLeaveUnavailable` catalog loading path, and the `ALTER DATABASE ... MODIFY SETTING` / `CatalogState` machinery) belongs to other upstream PRs that are not on antalya-26.6. Adapted: DatabaseDataLake::getCatalog reads vended_credentials_cache_ttl from the plain `settings` member, because antalya-26.6 has no MultiVersion database_settings Dropped: catalog auth-state invalidation of the vended credentials cache (cacheCredentials state-snapshot check and the commitSettingsChanges cache clear) - depends on the RestCatalog CatalogState/prepareSettingsChanges machinery not yet on antalya-26.6 --- src/Common/ProfileEvents.cpp | 17 -- src/Databases/DataLake/DatabaseDataLake.cpp | 25 +- src/Databases/DataLake/ICatalog.h | 25 -- src/Databases/DataLake/RestCatalog.cpp | 250 +----------------- src/Databases/DataLake/RestCatalog.h | 55 +--- .../test.py | 11 + 6 files changed, 18 insertions(+), 365 deletions(-) diff --git a/src/Common/ProfileEvents.cpp b/src/Common/ProfileEvents.cpp index 59b79ae83811..ddb67cf8a800 100644 --- a/src/Common/ProfileEvents.cpp +++ b/src/Common/ProfileEvents.cpp @@ -1503,26 +1503,9 @@ The server successfully detected this situation and will download merged part fr M(AIRowsProcessed, "Number of rows that received an AI result.", ValueType::Number) \ M(AIRowsSkipped, "Number of rows that received a default value due to quota or error.", ValueType::Number) \ \ -<<<<<<< HEAD -======= M(DataLakeRestCatalogCredentialsVended, "Number of table metadata requests to REST catalog asking to vend storage credentials.", ValueType::Number) \ M(DataLakeRestCatalogCredentialsCacheHits, "Number of table metadata requests to REST catalog reusing cached storage credentials.", ValueType::Number) \ \ - M(StatelessWorkerRequested, "Number of stateless workers requested by queries for distributed query execution.", ValueType::Number) \ - M(StatelessWorkerProvided, "Number of stateless workers provided to queries for distributed query execution.", ValueType::Number) \ - M(StatelessWorkerProvisioningMicroseconds, "Total time queries spent waiting for stateless workers to be provisioned.", ValueType::Microseconds) \ - M(StatelessWorkerProvisioningWaits, "Number of times a query had to wait for the stateless worker discovery service to provide workers.", ValueType::Number) \ - M(StatelessWorkerCreateLeaseRequests, "Number of create_lease requests sent to the stateless worker discovery service to acquire workers.", ValueType::Number) \ - M(StatelessWorkerCreateLeaseMicroseconds, "Total time spent in create_lease requests to the stateless worker discovery service.", ValueType::Microseconds) \ - M(StatelessWorkerCreateLeaseErrors, "Number of failed create_lease requests to the stateless worker discovery service (any error or timeout).", ValueType::Number) \ - M(StatelessWorkerGetLeaseRequests, "Number of get_lease requests sent to the stateless worker discovery service to refresh the list of leased workers.", ValueType::Number) \ - M(StatelessWorkerGetLeaseMicroseconds, "Total time spent in get_lease requests to the stateless worker discovery service.", ValueType::Microseconds) \ - M(StatelessWorkerGetLeaseErrors, "Number of failed get_lease requests to the stateless worker discovery service (any error or timeout).", ValueType::Number) \ - M(StatelessWorkerUpdateLeaseRequests, "Number of update_lease requests sent to the stateless worker discovery service to renew worker leases.", ValueType::Number) \ - M(StatelessWorkerUpdateLeaseMicroseconds, "Total time spent in update_lease requests to the stateless worker discovery service.", ValueType::Microseconds) \ - M(StatelessWorkerUpdateLeaseErrors, "Number of failed update_lease requests to the stateless worker discovery service (any error or timeout).", ValueType::Number) \ - \ ->>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) #ifdef APPLY_FOR_EXTERNAL_EVENTS #define APPLY_FOR_EVENTS(M) APPLY_FOR_BUILTIN_EVENTS(M) APPLY_FOR_EXTERNAL_EVENTS(M) diff --git a/src/Databases/DataLake/DatabaseDataLake.cpp b/src/Databases/DataLake/DatabaseDataLake.cpp index f31cbf5dc22a..a44d4f97cc14 100644 --- a/src/Databases/DataLake/DatabaseDataLake.cpp +++ b/src/Databases/DataLake/DatabaseDataLake.cpp @@ -320,35 +320,14 @@ std::shared_ptr DatabaseDataLake::getCatalog() const std::lock_guard lock(catalog_mutex); /// Lazily build the catalog on first access for databases attached at startup (see ctor). if (!catalog_impl) -<<<<<<< HEAD -======= - { - initializeOrLeaveUnavailable(); - if (!catalog_impl) - throw Exception( - ErrorCodes::ACCESS_DENIED, - "DataLakeCatalog database is inaccessible: its catalog uses server-managed credentials that are " - "restricted for user queries and could not be resolved when the database was loaded from metadata. " - "Provide explicit credentials, or enable `s3_allow_server_credentials_in_user_queries`. Reason: {}", - catalog_unavailable_reason); - } + initialize(); - const auto settings_version = database_settings.get(); catalog_impl->setVendedCredentialsCacheTTL( - std::chrono::seconds((*settings_version)[DatabaseDataLakeSetting::vended_credentials_cache_ttl].value)); + std::chrono::seconds(settings[DatabaseDataLakeSetting::vended_credentials_cache_ttl].value)); return catalog_impl; } -void DatabaseDataLake::initializeOrLeaveUnavailable() const -{ - try - { ->>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) - initialize(); - return catalog_impl; -} - std::shared_ptr DatabaseDataLake::getConfiguration( DatabaseDataLakeStorageType type, DataLakeStorageSettingsPtr storage_settings) const diff --git a/src/Databases/DataLake/ICatalog.h b/src/Databases/DataLake/ICatalog.h index 499c62dd331b..e0dec0115a71 100644 --- a/src/Databases/DataLake/ICatalog.h +++ b/src/Databases/DataLake/ICatalog.h @@ -213,33 +213,8 @@ class ICatalog return std::nullopt; } -<<<<<<< HEAD -======= virtual void setVendedCredentialsCacheTTL(std::chrono::seconds /*ttl*/) {} - /// Result of `prepareSettingsChanges`: the new catalog state built off to the side, - /// ready to be published by `commitSettingsChanges`. - struct PreparedSettingsChanges - { - virtual ~PreparedSettingsChanges() = default; - }; - using PreparedSettingsChangesPtr = std::unique_ptr; - - /// Validate `ALTER DATABASE ... MODIFY SETTING` changes and build the new catalog - /// state without publishing anything (may throw, may do network I/O). The state - /// becomes visible only after `commitSettingsChanges`, so the caller can persist - /// the changes in between and abandon the prepared state on failure. - virtual PreparedSettingsChangesPtr prepareSettingsChanges(const DB::SettingsChanges & changes); - - /// Publish the state built by `prepareSettingsChanges`. Must not fail. - virtual void commitSettingsChanges(PreparedSettingsChangesPtr prepared); - - void applySettingsChanges(const DB::SettingsChanges & changes) - { - commitSettingsChanges(prepareSettingsChanges(changes)); - } - ->>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) protected: /// Name of the warehouse, /// which is sometimes also called "catalog name". diff --git a/src/Databases/DataLake/RestCatalog.cpp b/src/Databases/DataLake/RestCatalog.cpp index 321b3aec8056..d48cf5e35bdb 100644 --- a/src/Databases/DataLake/RestCatalog.cpp +++ b/src/Databases/DataLake/RestCatalog.cpp @@ -49,15 +49,11 @@ #include #include #include -<<<<<<< HEAD -======= #include #include #include #include #include -#include ->>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) namespace DB::ErrorCodes @@ -78,25 +74,12 @@ namespace DB::FailPoints extern const char check_database_datalake_negative[]; } -<<<<<<< HEAD -======= namespace ProfileEvents { extern const Event DataLakeRestCatalogCredentialsVended; extern const Event DataLakeRestCatalogCredentialsCacheHits; } -namespace DB::DatabaseDataLakeSetting -{ - extern const DatabaseDataLakeSettingsString catalog_credential; - extern const DatabaseDataLakeSettingsString auth_header; - extern const DatabaseDataLakeSettingsString onelake_tenant_id; - extern const DatabaseDataLakeSettingsString onelake_bearer_token; - extern const DatabaseDataLakeSettingsString onelake_client_id; - extern const DatabaseDataLakeSettingsString onelake_client_secret; -} - ->>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) namespace DataLake { @@ -348,229 +331,7 @@ OneLakeCatalog::OneLakeCatalog( config = loadConfig(); } -<<<<<<< HEAD AccessToken RestCatalog::retrieveAccessToken() const -======= -void RestCatalog::validateSettingsChangesImpl( - const DB::SettingsChanges & changes, - const std::unordered_set & alterable_settings, - const std::string & auth_mode_description) -{ - for (const auto & change : changes) - { - if (alterable_settings.empty()) - throw DB::Exception( - DB::ErrorCodes::BAD_ARGUMENTS, - "Setting `{}` cannot be altered for a {}: the database was created without authentication settings", - change.name, - auth_mode_description); - - if (!alterable_settings.contains(change.name)) - throw DB::Exception( - DB::ErrorCodes::BAD_ARGUMENTS, - "Setting `{}` cannot be altered for a {} " - "(the authentication mode is fixed when the database is created; " - "alterable settings are: {})", - change.name, - auth_mode_description, - fmt::join(alterable_settings, ", ")); - - if (change.value.getType() != DB::Field::Types::String) - throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Setting `{}` must be a string", change.name); - - if (change.value.safeGet().empty()) - throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Setting `{}` cannot be set to an empty value", change.name); - } -} - -void RestCatalog::validateSettingsChanges(const DB::SettingsChanges & changes, bool credential_mode, bool header_mode) -{ - static const std::unordered_set credential_mode_settings = { - DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::catalog_credential)}; - static const std::unordered_set header_mode_settings = { - DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::auth_header)}; - static const std::unordered_set no_auth_settings = {}; - - if (credential_mode) - validateSettingsChangesImpl(changes, credential_mode_settings, "REST catalog with catalog credential authentication"); - else if (header_mode) - validateSettingsChangesImpl(changes, header_mode_settings, "REST catalog with auth header authentication"); - else - validateSettingsChangesImpl(changes, no_auth_settings, "REST catalog"); -} - -struct RestCatalog::PreparedAuthChanges : ICatalog::PreparedSettingsChanges -{ - std::unique_ptr new_state; - /// Set only when the OAuth credentials changed. - std::unique_ptr new_access_token; -}; - -ICatalog::PreparedSettingsChangesPtr RestCatalog::prepareSettingsChanges(const DB::SettingsChanges & changes) -{ - const auto old_state = state.get(); - CatalogState new_state = *old_state; - - auto prepared = std::make_unique(); - std::optional new_auth_headers; - applySettingsChangesToState(changes, *old_state, new_state, new_auth_headers, prepared->new_access_token); - - /// The config was loaded with the old credentials; the new ones may resolve the - /// warehouse to a different prefix or base location, so reload it before publishing. - new_state.config = loadConfig(new_state, new_auth_headers); - prepared->new_state = std::make_unique(std::move(new_state)); - return prepared; -} - -void RestCatalog::commitSettingsChanges(ICatalog::PreparedSettingsChangesPtr prepared) -{ - auto * prepared_auth = dynamic_cast(prepared.get()); - if (!prepared_auth || !prepared_auth->new_state) - throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Settings changes to commit were not prepared by this catalog"); - - state.set(std::move(prepared_auth->new_state)); - if (prepared_auth->new_access_token) - access_token.set(std::move(prepared_auth->new_access_token)); - - /// Cached storage credentials were vended under the old auth identity; do not reuse them. - { - std::lock_guard lock(credentials_cache_mutex); - credentials_cache.clear(); - } -} - -void RestCatalog::applySettingsChangesToState( - const DB::SettingsChanges & changes, - const CatalogState & old_state, - CatalogState & new_state, - std::optional & new_auth_headers, - std::unique_ptr & new_access_token) -{ - const bool credential_mode = !old_state.client_id.empty(); - const bool header_mode = old_state.auth_header.has_value(); - - validateSettingsChanges(changes, credential_mode, header_mode); - - for (const auto & change : changes) - { - if (change.name == DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::catalog_credential)) - { - std::tie(new_state.client_id, new_state.client_secret) = parseCatalogCredential(change.value.safeGet()); - } - else if (change.name == DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::auth_header)) - { - new_state.auth_header = parseAuthHeader(change.value.safeGet()); - validateAuthHeaders(new_state.auth_header.value()); - } - else - throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Unexpected setting `{}` after validation", change.name); - } - - if (credential_mode && (new_state.client_id != old_state.client_id || new_state.client_secret != old_state.client_secret)) - { - /// Eagerly fetch a token with the not-yet-published credentials: wrong credentials - /// fail the ALTER right here, and the config reload authenticates with that token - /// instead of the cached one. - new_access_token = std::make_unique(retrieveAccessToken(new_state.client_id, new_state.client_secret)); - new_auth_headers = DB::HTTPHeaderEntries{{"Authorization", "Bearer " + new_access_token->token}}; - } -} - -DB::HTTPHeaderEntries OneLakeCatalog::getAuthHeaders(const CatalogState & catalog_state, bool update_token) const -{ - auto headers = RestCatalog::getAuthHeaders(catalog_state, update_token); - headers.emplace_back("User-Agent", fmt::format("ClickHouse/{}{} OneLake-Catalog", VERSION_STRING, VERSION_OFFICIAL)); - return headers; -} - -void OneLakeCatalog::validateSettingsChanges(const DB::SettingsChanges & changes, bool bearer_mode) -{ - static const std::unordered_set bearer_mode_settings = { - DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_tenant_id), - DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_bearer_token)}; - static const std::unordered_set client_mode_settings = { - DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_tenant_id), - DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_client_id), - DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_client_secret)}; - - RestCatalog::validateSettingsChangesImpl( - changes, - bearer_mode ? bearer_mode_settings : client_mode_settings, - bearer_mode ? "OneLake catalog with bearer token authentication" : "OneLake catalog with client credentials authentication"); -} - -void OneLakeCatalog::applySettingsChangesToState( - const DB::SettingsChanges & changes, - const CatalogState & old_state, - CatalogState & new_state, - std::optional & new_auth_headers, - std::unique_ptr & new_access_token) -{ - const bool bearer_mode = !old_state.bearer_token.empty(); - - validateSettingsChanges(changes, bearer_mode); - - for (const auto & change : changes) - { - if (change.name == DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_tenant_id)) - new_state.tenant_id = change.value.safeGet(); - else if (change.name == DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_bearer_token)) - new_state.bearer_token = change.value.safeGet(); - else if (change.name == DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_client_id)) - new_state.client_id = change.value.safeGet(); - else if (change.name == DB::DatabaseDataLakeSettings::getSettingName(DB::DatabaseDataLakeSetting::onelake_client_secret)) - new_state.client_secret = change.value.safeGet(); - else - throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Unexpected setting `{}` after validation", change.name); - } - - if (bearer_mode) - { - new_state.auth_header = DB::HTTPHeaderEntry("Authorization", "Bearer " + new_state.bearer_token); - validateAuthHeaders(new_state.auth_header.value()); - } - else if (new_state.client_id != old_state.client_id || new_state.client_secret != old_state.client_secret) - { - /// Eagerly fetch a token with the not-yet-published credentials: wrong credentials - /// fail the ALTER right here, and the config reload authenticates with that token - /// instead of the cached one. - new_access_token = std::make_unique(retrieveAccessToken(new_state.client_id, new_state.client_secret)); - new_auth_headers = DB::HTTPHeaderEntries{{"Authorization", "Bearer " + new_access_token->token}}; - } -} - -namespace -{ - -[[maybe_unused]] const bool rest_settings_alter_validator_registered = [] -{ - CatalogSettingsAlterValidatorFactory::instance().registerValidator( - DB::DatabaseDataLakeCatalogType::ICEBERG_REST, - [](const DB::DatabaseDataLakeSettings & current_settings, const DB::SettingsChanges & changes) - { - const bool credential_mode = !current_settings[DB::DatabaseDataLakeSetting::catalog_credential].value.empty(); - const bool header_mode = !current_settings[DB::DatabaseDataLakeSetting::auth_header].value.empty(); - RestCatalog::validateSettingsChanges(changes, credential_mode, header_mode); - }); - return true; -}(); - -[[maybe_unused]] const bool onelake_settings_alter_validator_registered = [] -{ - CatalogSettingsAlterValidatorFactory::instance().registerValidator( - DB::DatabaseDataLakeCatalogType::ICEBERG_ONELAKE, - [](const DB::DatabaseDataLakeSettings & current_settings, const DB::SettingsChanges & changes) - { - const bool bearer_mode = !current_settings[DB::DatabaseDataLakeSetting::onelake_bearer_token].value.empty(); - OneLakeCatalog::validateSettingsChanges(changes, bearer_mode); - }); - return true; -}(); - -} - -AccessToken RestCatalog::retrieveAccessToken(const std::string & client_id, const std::string & client_secret) const ->>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) { static constexpr auto oauth_tokens_endpoint = "oauth/tokens"; @@ -1474,7 +1235,7 @@ bool RestCatalog::getTableMetadataImpl( if (parsed.credentials) { result.setStorageCredentials(parsed.credentials); - cacheCredentials(namespace_name, table_name, parsed, state_snapshot); + cacheCredentials(namespace_name, table_name, parsed); } if (!parsed.endpoint.empty()) result.setEndpoint(parsed.endpoint); @@ -1909,8 +1670,7 @@ std::optional RestCatalog::tryGetCachedCredentials( void RestCatalog::cacheCredentials( const std::string & namespace_name, const std::string & table_name, - const VendedStorageCredentials & parsed, - const CatalogStateVersion & state_snapshot) const + const VendedStorageCredentials & parsed) const { const auto ttl = vended_credentials_cache_ttl.load(std::memory_order_relaxed); if (ttl <= std::chrono::seconds::zero()) @@ -1934,10 +1694,6 @@ void RestCatalog::cacheCredentials( std::lock_guard lock(credentials_cache_mutex); - /// Do not cache credentials vended under an outdated auth state. - if (state.get() != state_snapshot) - return; - if (credentials_cache.size() >= credentials_cache_cleanup_threshold) std::erase_if(credentials_cache, [&now](const auto & entry) { return now >= entry.second.expires_at.value(); }); credentials_cache[{namespace_name, table_name}] @@ -1996,7 +1752,7 @@ ICatalog::CredentialsRefreshCallback RestCatalog::getCredentialsConfigurationCal if (metadata_object) parsed.table_uuid = parseTableUuid(metadata_object); /// Refresh the per-table cache so subsequent queries reuse these freshly vended credentials. - cacheCredentials(namespace_name, table_name, parsed, state_snapshot); + cacheCredentials(namespace_name, table_name, parsed); return parsed.credentials; }; } diff --git a/src/Databases/DataLake/RestCatalog.h b/src/Databases/DataLake/RestCatalog.h index 25035c61fba0..d673427e4341 100644 --- a/src/Databases/DataLake/RestCatalog.h +++ b/src/Databases/DataLake/RestCatalog.h @@ -12,13 +12,9 @@ #include #include #include -<<<<<<< HEAD -======= #include #include #include -#include ->>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) #include namespace DB @@ -105,54 +101,10 @@ class RestCatalog : public ICatalog, public DB::WithContext ICatalog::CredentialsRefreshCallback getCredentialsConfigurationCallback(const DB::StorageID & storage_id) override; -<<<<<<< HEAD String getClientId() const { return client_id; } String getClientSecret() const { return client_secret; } -======= - void setVendedCredentialsCacheTTL(std::chrono::seconds ttl) override { vended_credentials_cache_ttl.store(ttl, std::memory_order_relaxed); } - struct Config - { - /// Prefix is a path of the catalog endpoint, - /// e.g. /v1/{prefix}/namespaces/{namespace}/tables/{table} - std::filesystem::path prefix; - /// Base location is location of data in storage - /// (in filesystem or object storage). - std::string default_base_location; - - std::string toString() const; - }; - - /// Credentials together with the catalog configuration they resolve to - /// (the /v1/config response depends on the credentials), published as one - /// atomic snapshot so readers never see a torn combination of them. - struct CatalogState - { - std::optional auth_header; - std::string client_id; - std::string client_secret; - std::string tenant_id; - std::string bearer_token; - Config config; - }; - using CatalogStateVersion = MultiVersion::Version; - - CatalogStateVersion getStateSnapshot() const { return state.get(); } - - ICatalog::PreparedSettingsChangesPtr prepareSettingsChanges(const DB::SettingsChanges & changes) override; - - void commitSettingsChanges(ICatalog::PreparedSettingsChangesPtr prepared) override; - - /// Check that we actually support these settings alter - static void validateSettingsChangesImpl( - const DB::SettingsChanges & changes, - const std::unordered_set & alterable_settings, - const std::string & auth_mode_description); - - /// `credential_mode` means the catalog authenticates with `catalog_credential`, - /// `header_mode` with `auth_header`. The mode is fixed when the database is created. - static void validateSettingsChanges(const DB::SettingsChanges & changes, bool credential_mode, bool header_mode); ->>>>>>> 8e1a1d4d0b0 (Cache vended credentials from REST data lake catalogs) + void setVendedCredentialsCacheTTL(std::chrono::seconds ttl) override { vended_credentials_cache_ttl.store(ttl, std::memory_order_relaxed); } protected: RestCatalog( @@ -254,13 +206,10 @@ class RestCatalog : public ICatalog, public DB::WithContext std::optional tryGetCachedCredentials( const std::string & namespace_name, const std::string & table_name) const; - /// `state_snapshot` is the state the credentials were vended under: if the auth settings - /// changed since, the entry is not inserted. void cacheCredentials( const std::string & namespace_name, const std::string & table_name, - const VendedStorageCredentials & parsed, - const CatalogStateVersion & state_snapshot) const; + const VendedStorageCredentials & parsed) const; AccessToken retrieveAccessToken() const; }; diff --git a/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py b/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py index 47b57a014bed..7a7343c2f738 100644 --- a/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py +++ b/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py @@ -516,6 +516,17 @@ def test_vended_credentials_cache_invalidated_on_table_replace(started_cluster): assert vended == 0 and hits >= 1 +@pytest.mark.skip( + reason="ALTER DATABASE ... MODIFY SETTING is not supported for the DataLakeCatalog " + "engine on this branch: DatabaseDataLake does not override " + "IDatabase::applySettingsChanges, and the RestCatalog CatalogState / " + "prepareSettingsChanges / commitSettingsChanges machinery this test relies on is " + "upstream code that is not part of this PR and has not been backported here, so " + "the ALTER fails with NOT_IMPLEMENTED before the cache behaviour is ever exercised. " + "The credentials cache itself is covered by test_vended_credentials_cache and " + "test_vended_credentials_cache_invalidated_on_table_replace. Re-enable this test " + "together with the ALTER DATABASE ... MODIFY SETTING backport." +) def test_vended_credentials_cache_cleared_on_auth_change(started_cluster): node = started_cluster.instances["node1"] catalog = load_catalog_impl(started_cluster)