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..ddb67cf8a800 100644 --- a/src/Common/ProfileEvents.cpp +++ b/src/Common/ProfileEvents.cpp @@ -1503,6 +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) \ \ + 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) \ + \ #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..a44d4f97cc14 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; @@ -319,6 +321,10 @@ std::shared_ptr DatabaseDataLake::getCatalog() const /// Lazily build the catalog on first access for databases attached at startup (see ctor). if (!catalog_impl) initialize(); + + catalog_impl->setVendedCredentialsCacheTTL( + std::chrono::seconds(settings[DatabaseDataLakeSetting::vended_credentials_cache_ttl].value)); + return catalog_impl; } @@ -1208,6 +1214,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..e0dec0115a71 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,8 @@ class ICatalog return std::nullopt; } + virtual void setVendedCredentialsCacheTTL(std::chrono::seconds /*ttl*/) {} + 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 9f85e9d80d6d..a46f57fbffb1 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,11 @@ #include #include #include +#include +#include +#include +#include +#include namespace DB::ErrorCodes @@ -68,6 +74,12 @@ namespace DB::FailPoints extern const char check_database_datalake_negative[]; } +namespace ProfileEvents +{ + extern const Event DataLakeRestCatalogCredentialsVended; + extern const Event DataLakeRestCatalogCredentialsCacheHits; +} + namespace DataLake { @@ -77,6 +89,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 ":" @@ -1031,23 +1050,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; @@ -1076,6 +1190,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()) { @@ -1101,16 +1217,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); + } + if (!parsed.endpoint.empty()) + result.setEndpoint(parsed.endpoint); + } } if (result.requiresDataLakeSpecificProperties()) @@ -1122,8 +1258,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; } @@ -1381,7 +1517,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) @@ -1389,22 +1587,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)) @@ -1413,9 +1618,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: { @@ -1437,15 +1646,65 @@ 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 +{ + 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); + + 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) @@ -1475,32 +1734,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); + return parsed.credentials; }; } diff --git a/src/Databases/DataLake/RestCatalog.h b/src/Databases/DataLake/RestCatalog.h index 982475ee2c96..d673427e4341 100644 --- a/src/Databases/DataLake/RestCatalog.h +++ b/src/Databases/DataLake/RestCatalog.h @@ -8,7 +8,13 @@ #include #include #include +#include +#include +#include #include +#include +#include +#include #include namespace DB @@ -32,6 +38,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: @@ -90,6 +104,8 @@ class RestCatalog : public ICatalog, public DB::WithContext 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); } + protected: RestCatalog( const std::string & warehouse_, @@ -131,6 +147,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 +188,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 +201,15 @@ 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; + + void cacheCredentials( + const std::string & namespace_name, + const std::string & table_name, + const VendedStorageCredentials & parsed) 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..7a7343c2f738 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,177 @@ 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 + + +@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) + + 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}") +