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..bcb6ec26fc95 100644 --- a/src/Common/ProfileEvents.cpp +++ b/src/Common/ProfileEvents.cpp @@ -1503,6 +1503,8 @@ 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 Iceberg REST catalog that asked the catalog to vend storage credentials (i.e. cache miss).", ValueType::Number) \ + M(DataLakeRestCatalogCredentialsCacheHits, "Number of table metadata requests to Iceberg REST catalog that reused cached storage credentials and did not ask the catalog to vend new ones.", 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..2b92f627102d 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,11 @@ std::shared_ptr DatabaseDataLake::getCatalog() const /// Lazily build the catalog on first access for databases attached at startup (see ctor). if (!catalog_impl) initialize(); + + if (catalog_impl) + catalog_impl->setVendedCredentialsCacheTTL( + std::chrono::seconds(settings[DatabaseDataLakeSetting::vended_credentials_cache_ttl].value)); + return catalog_impl; } 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 28c1195082e4..8b4fcb9a8f9a 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 { @@ -1032,15 +1044,28 @@ bool RestCatalog::getTableMetadataImpl( 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) { + 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::DataLakeRestCatalogCredentialsCacheHits); + else + { + 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; @@ -1094,16 +1119,28 @@ 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) + { + result.setStorageCredentials(cached_credentials->credentials); + if (!cached_credentials->endpoint.empty()) + result.setEndpoint(cached_credentials->endpoint); + } + else if (object->has("config")) + { + auto config_object = object->get("config").extract(); + if (!config_object) + throw DB::Exception(DB::ErrorCodes::LOGICAL_ERROR, "Cannot parse config result"); + auto parsed = getCredentialsAndEndpoint(config_object, location); + if (parsed.credentials) + { + result.setStorageCredentials(parsed.credentials); + cacheCredentials(namespace_name, table_name, parsed); + } + if (!parsed.endpoint.empty()) + result.setEndpoint(parsed.endpoint); + } } if (result.requiresDataLakeSpecificProperties()) @@ -1374,7 +1411,44 @@ 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 +{ +std::optional 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 (...) + { + return std::nullopt; + } + return std::nullopt; + } + return std::nullopt; +} +} + +VendedStorageCredentials RestCatalog::getCredentialsAndEndpoint(Poco::JSON::Object::Ptr object, const String & location) const { auto storage_type = parseStorageTypeFromLocation(location); switch (storage_type) @@ -1386,18 +1460,20 @@ std::pair, String> RestCatalog::getCredenti 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)) { 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), ""}; + return {std::make_shared(gcs_token), "", std::nullopt}; } 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 +1482,26 @@ 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(); + if (object->has(session_token_expires_at_ms_str)) + { + 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(session_token_expires_at_ms_str).convert(); + if (expires_at_ms <= 0) + expires_at = std::chrono::system_clock::time_point{}; /// Already invalid: do not cache. + else if (expires_at_ms / 1000 < max_representable_sec) + expires_at = std::chrono::system_clock::from_time_t(static_cast(expires_at_ms / 1000)); + } + catch (...) + { + LOG_DEBUG(log, "Failed to parse '{}' from vended credentials config", session_token_expires_at_ms_str); + } + } 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 +1523,59 @@ 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 && parsed.expires_at.value() < refresh_after) + refresh_after = parsed.expires_at.value(); + 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}; } ICatalog::CredentialsRefreshCallback RestCatalog::getCredentialsConfigurationCallback(const DB::StorageID & storage_id) @@ -1492,8 +1629,10 @@ ICatalog::CredentialsRefreshCallback RestCatalog::getCredentialsConfigurationCal throw DB::Exception(DB::ErrorCodes::BAD_ARGUMENTS, "Cannot read table {}, because no 'metadata-location' in response", table_name); } - auto [new_credentials, _] = getCredentialsAndEndpoint(config_object, location); - return new_credentials; + auto parsed = getCredentialsAndEndpoint(config_object, location); + /// 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..c76e7fb86c94 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,13 @@ struct AccessToken } }; +struct VendedStorageCredentials +{ + std::shared_ptr credentials; + std::string endpoint; + std::optional expires_at; +}; + class RestCatalog : public ICatalog, public DB::WithContext { public: @@ -87,6 +100,8 @@ class RestCatalog : public ICatalog, public DB::WithContext ICatalog::CredentialsRefreshCallback getCredentialsConfigurationCallback(const DB::StorageID & storage_id) override; + void setVendedCredentialsCacheTTL(std::chrono::seconds ttl) override { vended_credentials_cache_ttl.store(ttl, std::memory_order_relaxed); } + String getClientId() const { return client_id; } String getClientSecret() const { return client_secret; } @@ -131,6 +146,16 @@ 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; + 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( @@ -172,7 +197,13 @@ 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..43e71002f2aa 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) @@ -89,6 +94,8 @@ class GCSCredentials final : public IStorageCredentials const std::string & getToken() const { return oauth_token; } + bool isEmpty() const override { return oauth_token.empty(); } + private: std::string oauth_token; }; @@ -109,6 +116,8 @@ class AzureCredentials final : public IStorageCredentials engine_args.push_back(DB::make_intrusive(sas_token)); } + bool isEmpty() const override { return sas_token.empty(); } + private: std::string sas_token; }; diff --git a/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py b/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py index 01b25fd75f86..5ff56aabfa37 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,76 @@ 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 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" + + 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": [1], "data": ["x"]}).astype({"id": "int32"}) + ) + ) + + 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 +