diff --git a/src/Common/ProfileEvents.cpp b/src/Common/ProfileEvents.cpp index c100f5999c0e..d23fbb1137cc 100644 --- a/src/Common/ProfileEvents.cpp +++ b/src/Common/ProfileEvents.cpp @@ -1724,6 +1724,56 @@ The server successfully detected this situation and will download merged part fr M(StatelessWorkerDiscoveryHeartbeatsRejected, "Number of heartbeats the stateless worker discovery service rejected because the worker had already been evicted.", ValueType::Number) \ M(StatelessWorkerDiscoveryKeeperTransactionRetries, "Number of write transactions the stateless worker discovery service retried because its coordination store (Keeper) state was modified concurrently.", ValueType::Number) \ \ + M(DataLakeRestCatalogLoadConfig, "Number of 'load config' requests to Iceberg REST catalog.", ValueType::Number) \ + M(DataLakeRestCatalogLoadConfigMicroseconds, "Total time of 'load config' requests to Iceberg REST catalog.", ValueType::Microseconds) \ + M(DataLakeRestCatalogGetNamespaces, "Number of 'get namespaces' requests to Iceberg REST catalog.", ValueType::Number) \ + M(DataLakeRestCatalogGetNamespacesMicroseconds, "Total time of 'get namespaces' requests to Iceberg REST catalog.", ValueType::Microseconds) \ + M(DataLakeRestCatalogGetTables, "Number of 'get tables' requests to Iceberg REST catalog.", ValueType::Number) \ + M(DataLakeRestCatalogGetTablesMicroseconds, "Total time of 'get tables' requests to Iceberg REST catalog.", ValueType::Microseconds) \ + M(DataLakeRestCatalogGetTableMetadata, "Number of 'get table metadata' requests to Iceberg REST catalog.", ValueType::Number) \ + M(DataLakeRestCatalogGetTableMetadataMicroseconds, "Total time of 'get table metadata' requests to Iceberg REST catalog.", ValueType::Microseconds) \ + M(DataLakeRestCatalogGetCredentials, "Number of 'get credentials' requests to Iceberg REST catalog.", ValueType::Number) \ + M(DataLakeRestCatalogGetCredentialsMicroseconds, "Total time of 'get credentials' requests to Iceberg REST catalog.", ValueType::Microseconds) \ + M(DataLakeRestCatalogAuthTokenCachedValid, "Number of requests to Iceberg REST catalog that reused a cached access token and did not fetch a new one.", ValueType::Number) \ + M(DataLakeRestCatalogAuthTokenRetrieve, "Number of new access tokens fetched for Iceberg REST catalog (OAuth client-credentials or GCP metadata/ADC).", ValueType::Number) \ + M(DataLakeRestCatalogAuthTokenRefreshedMicroseconds, "Total time spent fetching access tokens for Iceberg REST catalog.", ValueType::Microseconds) \ + M(DataLakeRestCatalogUnauthorized, "Number of Iceberg REST catalog HTTP requests retried with a new access token after HTTP 401 or 403.", ValueType::Number) \ + M(DataLakeRestCatalogCreateNamespace, "Number of 'create namespace' requests to Iceberg REST catalog.", ValueType::Number) \ + M(DataLakeRestCatalogCreateNamespaceMicroseconds, "Total time of 'create namespace' requests to Iceberg REST catalog.", ValueType::Microseconds) \ + M(DataLakeRestCatalogCreateTable, "Number of 'create table' requests to Iceberg REST catalog.", ValueType::Number) \ + M(DataLakeRestCatalogCreateTableMicroseconds, "Total time of 'create table' requests to Iceberg REST catalog.", ValueType::Microseconds) \ + M(DataLakeRestCatalogUpdateTable, "Number of 'update table' requests to Iceberg REST catalog.", ValueType::Number) \ + M(DataLakeRestCatalogUpdateTableMicroseconds, "Total time of 'update table' requests to Iceberg REST catalog.", ValueType::Microseconds) \ + M(DataLakeRestCatalogDropTable, "Number of 'drop table' requests to Iceberg REST catalog.", ValueType::Number) \ + M(DataLakeRestCatalogDropTableMicroseconds, "Total time of 'drop table' requests to Iceberg REST catalog.", ValueType::Microseconds) \ + \ + M(DataLakeGlueCatalogGetDatabases, "Number of 'get databases' requests to Iceberg Glue catalog.", ValueType::Number) \ + M(DataLakeGlueCatalogGetDatabasesMicroseconds, "Total time of 'get databases' requests to Iceberg Glue catalog.", ValueType::Microseconds) \ + M(DataLakeGlueCatalogGetTables, "Number of 'get tables' requests to Iceberg Glue catalog.", ValueType::Number) \ + M(DataLakeGlueCatalogGetTablesMicroseconds, "Total time of 'get tables' requests to Iceberg Glue catalog.", ValueType::Microseconds) \ + M(DataLakeGlueCatalogGetTable, "Number of 'get table' requests to Iceberg Glue catalog.", ValueType::Number) \ + M(DataLakeGlueCatalogGetTableMicroseconds, "Total time of 'get table' requests to Iceberg Glue catalog.", ValueType::Microseconds) \ + M(DataLakeGlueCatalogCreateDatabase, "Number of 'create database' requests to Iceberg Glue catalog.", ValueType::Number) \ + M(DataLakeGlueCatalogCreateDatabaseMicroseconds, "Total time of 'create database' requests to Iceberg Glue catalog.", ValueType::Microseconds) \ + M(DataLakeGlueCatalogCreateTable, "Number of 'create table' requests to Iceberg Glue catalog.", ValueType::Number) \ + M(DataLakeGlueCatalogCreateTableMicroseconds, "Total time of 'create table' requests to Iceberg Glue catalog.", ValueType::Microseconds) \ + M(DataLakeGlueCatalogUpdateTable, "Number of 'update table' requests to Iceberg Glue catalog.", ValueType::Number) \ + M(DataLakeGlueCatalogUpdateTableMicroseconds, "Total time of 'update table' requests to Iceberg Glue catalog.", ValueType::Microseconds) \ + M(DataLakeGlueCatalogDropTable, "Number of 'drop table' requests to Iceberg Glue catalog.", ValueType::Number) \ + M(DataLakeGlueCatalogDropTableMicroseconds, "Total time of 'drop table' requests to Iceberg Glue catalog.", ValueType::Microseconds) \ + \ + M(DataLakeUnityCatalogGetTables, "Number of 'get tables' requests to Iceberg Unity catalog.", ValueType::Number) \ + M(DataLakeUnityCatalogGetTablesMicroseconds, "Total time of 'get tables' requests to Iceberg Unity catalog.", ValueType::Microseconds) \ + M(DataLakeUnityCatalogGetTable, "Number of 'get table' requests to Iceberg Unity catalog.", ValueType::Number) \ + M(DataLakeUnityCatalogGetTableMicroseconds, "Total time of 'get table' requests to Iceberg Unity catalog.", ValueType::Microseconds) \ + M(DataLakeUnityCatalogGetTableMetadata, "Number of 'get table metadata' requests to Iceberg Unity catalog.", ValueType::Number) \ + M(DataLakeUnityCatalogGetTableMetadataMicroseconds, "Total time of 'get table metadata' requests to Iceberg Unity catalog.", ValueType::Microseconds) \ + M(DataLakeUnityCatalogGetSchemas, "Number of 'get schemas' requests to Iceberg Unity catalog.", ValueType::Number) \ + M(DataLakeUnityCatalogGetSchemasMicroseconds, "Total time of 'get schemas' requests to Iceberg Unity catalog.", ValueType::Microseconds) \ + M(DataLakeUnityCatalogGetCredentials, "Number of 'get credentials' requests to Iceberg Unity catalog.", ValueType::Number) \ + M(DataLakeUnityCatalogGetCredentialsMicroseconds, "Total time of 'get credentials' requests to Iceberg Unity catalog.", ValueType::Microseconds) \ + \ + #ifdef APPLY_FOR_EXTERNAL_EVENTS #define APPLY_FOR_EVENTS(M) APPLY_FOR_BUILTIN_EVENTS(M) APPLY_FOR_EXTERNAL_EVENTS(M) diff --git a/src/Core/Settings.cpp b/src/Core/Settings.cpp index dc978ae678fc..a314ef978360 100644 --- a/src/Core/Settings.cpp +++ b/src/Core/Settings.cpp @@ -9151,6 +9151,9 @@ Multiple algorithms can be specified as a comma-separated list, e.g. `dphyp,gree )", EXPERIMENTAL) \ DECLARE(Bool, allow_experimental_database_paimon_rest_catalog, false, R"( Allow experimental database engine DataLakeCatalog with catalog_type = 'paimon_rest' +)", EXPERIMENTAL) \ + DECLARE(Bool, allow_experimental_database_s3_tables, false, R"( +Allow experimental database engine DataLakeCatalog with catalog_type = 's3tables' (Amazon S3 Tables Iceberg REST with SigV4) )", EXPERIMENTAL) \ DECLARE(UInt64, webassembly_udf_max_fuel, 100'000, R"( Fuel limit per WebAssembly UDF instance execution. Each WebAssembly instruction consumes some amount of fuel. The value is scaled by 1024 before being passed to the runtime, so `webassembly_udf_max_fuel = 1` corresponds to approximately 1024 fuel units. Set to 0 for no finite limit. Applies only to functions whose per-function setting `webassembly_udf_enable_fuel` is true, which is the default. diff --git a/src/Core/SettingsChangesHistory.cpp b/src/Core/SettingsChangesHistory.cpp index 0a3ddb973f7b..f870c1826133 100644 --- a/src/Core/SettingsChangesHistory.cpp +++ b/src/Core/SettingsChangesHistory.cpp @@ -265,6 +265,7 @@ const VersionToSettingsChangesMap & getSettingsChangesHistory() {"allow_experimental_query_deduplication", false, false, "The setting is obsolete, the feature has been removed."}, {"query_plan_min_columns_for_join_lazy_indexing", 0, 3, "Control the minimum number of payload columns from the left side required for enabling lazy indexing optimization in JOIN"}, {"query_plan_max_limit_for_join_lazy_indexing", 1000, 1000, "Added new setting to control maximum limit value that allows to use query plan for lazy join indexing optimization. If zero, there is no limit"}, + {"allow_experimental_database_s3_tables", false, false, "New setting to enable experimental database S3 tables (AWS Iceberg REST catalog)."}, {"statistics_max_set_size_for_exact_selectivity_estimation", 10000, 10000, "The bound on the cost of estimating the selectivity of `IN` with a large set is kept under `compatibility` with an earlier version: the previous value is deliberately equal to the new one, so that the uncapped estimation, which could add hundreds of milliseconds to the planning of a single query, is not restored."}, }); diff --git a/src/Databases/DataLake/DatabaseDataLake.cpp b/src/Databases/DataLake/DatabaseDataLake.cpp index e812f591c57a..2295ffe70173 100644 --- a/src/Databases/DataLake/DatabaseDataLake.cpp +++ b/src/Databases/DataLake/DatabaseDataLake.cpp @@ -104,6 +104,7 @@ namespace Setting extern const SettingsBool allow_experimental_database_glue_catalog; extern const SettingsBool allow_experimental_database_hms_catalog; extern const SettingsBool allow_experimental_database_paimon_rest_catalog; + extern const SettingsBool allow_experimental_database_s3_tables; extern const SettingsBool use_hive_partitioning; extern const SettingsBool log_queries; extern const SettingsBool parallel_replicas_for_cluster_engines; @@ -1664,11 +1665,11 @@ void registerDatabaseDataLake(DatabaseFactory & factory) case DatabaseDataLakeCatalogType::S3_TABLES: { if (!args.create_query.attach - && !args.context->getSettingsRef()[Setting::allow_experimental_database_iceberg]) + && !args.context->getSettingsRef()[Setting::allow_experimental_database_s3_tables]) { throw Exception(ErrorCodes::SUPPORT_IS_DISABLED, - "DatabaseDataLake with S3 Tables catalog (Iceberg REST) is beta. " - "To allow its usage, enable setting allow_database_iceberg"); + "DatabaseDataLake with S3 Tables catalog is experimental. " + "To allow its usage, enable setting allow_experimental_database_s3_tables"); } engine_func->name = "Iceberg"; diff --git a/src/Databases/DataLake/GlueCatalog.cpp b/src/Databases/DataLake/GlueCatalog.cpp index a2817627e33d..859a73c0fc0a 100644 --- a/src/Databases/DataLake/GlueCatalog.cpp +++ b/src/Databases/DataLake/GlueCatalog.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include @@ -83,6 +84,24 @@ namespace DB::ServerSetting extern const ServerSettingsUInt64 s3_retry_attempts; } +namespace ProfileEvents +{ + extern const Event DataLakeGlueCatalogGetDatabases; + extern const Event DataLakeGlueCatalogGetDatabasesMicroseconds; + extern const Event DataLakeGlueCatalogGetTables; + extern const Event DataLakeGlueCatalogGetTablesMicroseconds; + extern const Event DataLakeGlueCatalogGetTable; + extern const Event DataLakeGlueCatalogGetTableMicroseconds; + extern const Event DataLakeGlueCatalogCreateDatabase; + extern const Event DataLakeGlueCatalogCreateDatabaseMicroseconds; + extern const Event DataLakeGlueCatalogCreateTable; + extern const Event DataLakeGlueCatalogCreateTableMicroseconds; + extern const Event DataLakeGlueCatalogUpdateTable; + extern const Event DataLakeGlueCatalogUpdateTableMicroseconds; + extern const Event DataLakeGlueCatalogDropTable; + extern const Event DataLakeGlueCatalogDropTableMicroseconds; +} + namespace CurrentMetrics { extern const Metric MarkCacheBytes; @@ -212,7 +231,14 @@ DataLake::ICatalog::Namespaces GlueCatalog::getDatabases(const std::string & pre do { request.SetNextToken(next_token); - auto outcome = glue_client->GetDatabases(request); + + Aws::Glue::Model::GetDatabasesOutcome outcome; + { + ProfileEvents::increment(ProfileEvents::DataLakeGlueCatalogGetDatabases); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeGlueCatalogGetDatabasesMicroseconds); + outcome = glue_client->GetDatabases(request); + } + if (outcome.IsSuccess()) { const auto & databases_result = outcome.GetResult(); @@ -261,7 +287,12 @@ CatalogTables GlueCatalog::getTablesForDatabase(const std::string & db_name, siz do { request.SetNextToken(next_token); - auto outcome = glue_client->GetTables(request); + Aws::Glue::Model::GetTablesOutcome outcome; + { + ProfileEvents::increment(ProfileEvents::DataLakeGlueCatalogGetTables); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeGlueCatalogGetTablesMicroseconds); + outcome = glue_client->GetTables(request); + } if (outcome.IsSuccess()) { const auto & tables_result = outcome.GetResult(); @@ -339,7 +370,12 @@ bool GlueCatalog::tryGetTableMetadata( request.SetDatabaseName(database_name); request.SetName(table_name); - auto outcome = glue_client->GetTable(request); + Aws::Glue::Model::GetTableOutcome outcome; + { + ProfileEvents::increment(ProfileEvents::DataLakeGlueCatalogGetTable); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeGlueCatalogGetTableMicroseconds); + outcome = glue_client->GetTable(request); + } if (outcome.IsSuccess()) { const auto & table_outcome = outcome.GetResult().GetTable(); @@ -635,6 +671,8 @@ void GlueCatalog::createNamespaceIfNotExists(const String & namespace_name, cons db_input.SetName(namespace_name); create_request.SetDatabaseInput(db_input); + ProfileEvents::increment(ProfileEvents::DataLakeGlueCatalogCreateDatabase); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeGlueCatalogCreateDatabaseMicroseconds); auto outcome = glue_client->CreateDatabase(create_request); if (!outcome.IsSuccess() && outcome.GetError().GetErrorType() != Aws::Glue::GlueErrors::ALREADY_EXISTS) { @@ -672,7 +710,13 @@ void GlueCatalog::createTable(const String & namespace_name, const String & tabl request.SetTableInput(table_input); - auto response = glue_client->CreateTable(request); + Aws::Glue::Model::CreateTableOutcome response; + + { + ProfileEvents::increment(ProfileEvents::DataLakeGlueCatalogCreateTable); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeGlueCatalogCreateTableMicroseconds); + response = glue_client->CreateTable(request); + } if (!response.IsSuccess()) throw DB::Exception(DB::ErrorCodes::DATALAKE_DATABASE_ERROR, "Can not create metadata in glue catalog: {}", response.GetError().GetMessage()); @@ -707,7 +751,13 @@ bool GlueCatalog::updateMetadata(const String & namespace_name, const String & t request.SetTableInput(table_input); - auto response = glue_client->UpdateTable(request); + Aws::Glue::Model::UpdateTableOutcome response; + + { + ProfileEvents::increment(ProfileEvents::DataLakeGlueCatalogUpdateTable); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeGlueCatalogUpdateTableMicroseconds); + response = glue_client->UpdateTable(request); + } if (!response.IsSuccess()) throw DB::Exception(DB::ErrorCodes::DATALAKE_DATABASE_ERROR, "Can not update metadata in glue catalog {}", response.GetError().GetMessage()); @@ -731,7 +781,13 @@ void GlueCatalog::dropTable(const String & namespace_name, const String & table_ request.SetDatabaseName(namespace_name); request.SetName(table_name); - auto response = glue_client->DeleteTable(request); + Aws::Glue::Model::DeleteTableOutcome response; + + { + ProfileEvents::increment(ProfileEvents::DataLakeGlueCatalogDropTable); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeGlueCatalogDropTableMicroseconds); + response = glue_client->DeleteTable(request); + } if (!response.IsSuccess()) throw DB::Exception( diff --git a/src/Databases/DataLake/ICatalog.cpp b/src/Databases/DataLake/ICatalog.cpp index cdf297f0671e..e51c96377cdf 100644 --- a/src/Databases/DataLake/ICatalog.cpp +++ b/src/Databases/DataLake/ICatalog.cpp @@ -320,8 +320,22 @@ std::string TableMetadata::getMetadataLocation(const std::string & iceberg_metad metadata_location = metadata_location.substr(storage_type_str.size()); if (data_location.starts_with(storage_type_str)) data_location = data_location.substr(storage_type_str.size()); - else if (!endpoint.empty() && data_location.starts_with(endpoint)) - data_location = data_location.substr(endpoint.size()); + else if (!endpoint.empty()) + { + std::string normalized_endpoint = endpoint; + if (normalized_endpoint.ends_with('/')) + normalized_endpoint.pop_back(); + + if (data_location.starts_with(normalized_endpoint)) + { + data_location = data_location.substr(normalized_endpoint.size()); + /// `metadata_location` is relative to the bucket (the `s3://` prefix is stripped above), + /// while `data_location` still has the leading slash left over from the endpoint, + /// e.g. "/bucket/table-uuid/". Drop it so that the prefix comparison below works. + if (azure_account_with_suffix.empty() && data_location.starts_with('/')) + data_location = data_location.substr(1); + } + } if (metadata_location.starts_with(data_location)) { diff --git a/src/Databases/DataLake/RestCatalog.cpp b/src/Databases/DataLake/RestCatalog.cpp index 3b3b86785c5a..9c15ff5774b6 100644 --- a/src/Databases/DataLake/RestCatalog.cpp +++ b/src/Databases/DataLake/RestCatalog.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -81,6 +82,28 @@ namespace ProfileEvents extern const Event OneLakeAccessTokenRequestFailures; extern const Event OneLakeAccessTokenRequestMicroseconds; extern const Event OneLakeAccessTokenExpirations; + extern const Event DataLakeRestCatalogLoadConfig; + extern const Event DataLakeRestCatalogLoadConfigMicroseconds; + extern const Event DataLakeRestCatalogGetNamespaces; + extern const Event DataLakeRestCatalogGetNamespacesMicroseconds; + extern const Event DataLakeRestCatalogGetTables; + extern const Event DataLakeRestCatalogGetTablesMicroseconds; + extern const Event DataLakeRestCatalogGetTableMetadata; + extern const Event DataLakeRestCatalogGetTableMetadataMicroseconds; + extern const Event DataLakeRestCatalogGetCredentials; + extern const Event DataLakeRestCatalogGetCredentialsMicroseconds; + extern const Event DataLakeRestCatalogAuthTokenCachedValid; + extern const Event DataLakeRestCatalogAuthTokenRetrieve; + extern const Event DataLakeRestCatalogAuthTokenRefreshedMicroseconds; + extern const Event DataLakeRestCatalogUnauthorized; + extern const Event DataLakeRestCatalogCreateNamespace; + extern const Event DataLakeRestCatalogCreateNamespaceMicroseconds; + extern const Event DataLakeRestCatalogCreateTable; + extern const Event DataLakeRestCatalogCreateTableMicroseconds; + extern const Event DataLakeRestCatalogUpdateTable; + extern const Event DataLakeRestCatalogUpdateTableMicroseconds; + extern const Event DataLakeRestCatalogDropTable; + extern const Event DataLakeRestCatalogDropTableMicroseconds; } namespace DB::DatabaseDataLakeSetting @@ -248,10 +271,15 @@ RestCatalog::RestCatalog( RestCatalog::Config RestCatalog::loadConfig(const CatalogState & catalog_state, const std::optional & auth_headers) { Poco::URI::QueryParameters params = {{"warehouse", warehouse}}; - auto buf = createReadBuffer(catalog_state, CONFIG_ENDPOINT, params, /* headers */ {}, auth_headers); std::string json_str; - readJSONObjectPossiblyInvalid(json_str, *buf); + + { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogLoadConfig); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeRestCatalogLoadConfigMicroseconds); + auto buf = createReadBuffer(catalog_state, CONFIG_ENDPOINT, params, /* headers */ {}, auth_headers); + readJSONObjectPossiblyInvalid(json_str, *buf); + } LOG_DEBUG(log, "Received catalog configuration settings: {}", json_str); @@ -295,13 +323,19 @@ void RestCatalog::validateAuthHeaders(const DB::HTTPHeaderEntry & header) const getContext()->getGlobalContext()->getHTTPHeaderFilter().checkAndNormalizeHeaders(header_to_check); } -DB::HTTPHeaderEntries RestCatalog::getAuthHeaders(const CatalogState & catalog_state, bool update_token) const +DB::HTTPHeaderEntries RestCatalog::getAuthHeaders( + const CatalogState & catalog_state, + bool update_token, + bool * used_cached_oauth_token) const { fiu_do_on(DB::FailPoints::check_database_datalake_negative, { throw DB::Exception(DB::ErrorCodes::FAULT_INJECTED, "Injecting fault when checking database"); }); + if (used_cached_oauth_token) + *used_cached_oauth_token = false; + /// Option 1: user specified auth header manually. /// Header has format: 'Authorization: '. if (catalog_state.auth_header.has_value()) @@ -319,11 +353,15 @@ DB::HTTPHeaderEntries RestCatalog::getAuthHeaders(const CatalogState & catalog_s /// request fails with 401/403 and is retried with `update_token = true`, fetching /// a token with the snapshot's credentials. auto current = access_token.get(); - if (!current || update_token) + if (!current || update_token || current->isExpired()) { access_token.set(std::make_unique(retrieveAccessToken(catalog_state.client_id, catalog_state.client_secret))); current = access_token.get(); } + else if (used_cached_oauth_token) + { + *used_cached_oauth_token = true; + } DB::HTTPHeaderEntries headers; headers.emplace_back("Authorization", "Bearer " + current->token); @@ -501,7 +539,10 @@ void RestCatalog::applySettingsChangesToState( } } -DB::HTTPHeaderEntries OneLakeCatalog::getAuthHeaders(const CatalogState & catalog_state, bool update_token) const +DB::HTTPHeaderEntries OneLakeCatalog::getAuthHeaders( + const CatalogState & catalog_state, + bool update_token, + bool * used_cached_oauth_token) const { DB::HTTPHeaderEntries headers; if (!catalog_state.refresh_token.empty()) @@ -511,7 +552,7 @@ DB::HTTPHeaderEntries OneLakeCatalog::getAuthHeaders(const CatalogState & catalo } else { - headers = RestCatalog::getAuthHeaders(catalog_state, update_token); + headers = RestCatalog::getAuthHeaders(catalog_state, update_token, used_cached_oauth_token); } headers.emplace_back("User-Agent", fmt::format("ClickHouse/{}{} OneLake-Catalog", VERSION_STRING, VERSION_OFFICIAL)); return headers; @@ -743,6 +784,9 @@ namespace AccessToken RestCatalog::retrieveAccessToken(const std::string & client_id, const std::string & client_secret) const { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogAuthTokenRetrieve); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeRestCatalogAuthTokenRefreshedMicroseconds); + static constexpr auto oauth_tokens_endpoint = "oauth/tokens"; /// TODO: @@ -997,7 +1041,10 @@ BigLakeCatalog::BigLakeCatalog( state.set(std::make_unique(std::move(initial_state))); } -DB::HTTPHeaderEntries BigLakeCatalog::getAuthHeaders(const CatalogState & catalog_state, bool update_token) const +DB::HTTPHeaderEntries BigLakeCatalog::getAuthHeaders( + const CatalogState & catalog_state, + bool update_token, + bool * used_cached_oauth_token) const { /// Google Cloud OAuth2 for BigLake. /// Uses GCP metadata service or Application Default Credentials to get access token. @@ -1005,12 +1052,19 @@ DB::HTTPHeaderEntries BigLakeCatalog::getAuthHeaders(const CatalogState & catalo /// https://developers.google.com/identity/protocols/oauth2 if (!google_project_id.empty() || !google_adc_client_id.empty()) { + if (used_cached_oauth_token) + *used_cached_oauth_token = false; + auto current = access_token.get(); if (!current || update_token || current->isExpired()) { access_token.set(std::make_unique(retrieveGoogleCloudAccessToken())); current = access_token.get(); } + else if (used_cached_oauth_token) + { + *used_cached_oauth_token = true; + } DB::HTTPHeaderEntries headers; headers.emplace_back("Authorization", "Bearer " + current->token); @@ -1029,7 +1083,7 @@ DB::HTTPHeaderEntries BigLakeCatalog::getAuthHeaders(const CatalogState & catalo return headers; } - return RestCatalog::getAuthHeaders(catalog_state, update_token); + return RestCatalog::getAuthHeaders(catalog_state, update_token, used_cached_oauth_token); } AccessToken BigLakeCatalog::retrieveGoogleCloudAccessTokenFromRefreshToken() const @@ -1051,6 +1105,9 @@ AccessToken BigLakeCatalog::retrieveGoogleCloudAccessTokenFromRefreshToken() con AccessToken BigLakeCatalog::retrieveGoogleCloudAccessToken() const { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogAuthTokenRetrieve); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeRestCatalogAuthTokenRefreshedMicroseconds); + const auto & context = getContext(); /// An explicit Application Default Credentials triple is a user-supplied credential, so it is honored. @@ -1182,9 +1239,9 @@ DB::ReadWriteBufferFromHTTPPtr RestCatalog::createReadBuffer( if (!params.empty()) url.setQueryParameters(params); - auto create_buffer = [&](bool update_token) + auto create_buffer = [&](bool update_token, bool & used_cached_oauth_token) { - auto result_headers = auth_headers ? *auth_headers : getAuthHeaders(catalog_state, update_token); + auto result_headers = auth_headers ? *auth_headers : getAuthHeaders(catalog_state, update_token, &used_cached_oauth_token); std::move(headers.begin(), headers.end(), std::back_inserter(result_headers)); return DB::BuilderRWBufferFromHTTP(url) @@ -1202,7 +1259,11 @@ DB::ReadWriteBufferFromHTTPPtr RestCatalog::createReadBuffer( try { - return create_buffer(false); + bool used_cached_oauth_token = false; + auto buf = create_buffer(false, used_cached_oauth_token); + if (used_cached_oauth_token) + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogAuthTokenCachedValid); + return buf; } catch (const DB::HTTPException & e) { @@ -1211,7 +1272,9 @@ DB::ReadWriteBufferFromHTTPPtr RestCatalog::createReadBuffer( (status == Poco::Net::HTTPResponse::HTTPStatus::HTTP_UNAUTHORIZED || status == Poco::Net::HTTPResponse::HTTPStatus::HTTP_FORBIDDEN)) { - return create_buffer(true); + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogUnauthorized); + bool used_cached_oauth_token_on_retry = false; + return create_buffer(true, used_cached_oauth_token_on_retry); } throw; } @@ -1380,6 +1443,8 @@ RestCatalog::Namespaces RestCatalog::listChildNamespaces(const std::string & bas if (!page_token.empty()) params.push_back({"pageToken", page_token}); + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogGetNamespaces); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeRestCatalogGetNamespacesMicroseconds); auto buf = createReadBuffer( *state_snapshot, state_snapshot->config.prefix / NAMESPACES_ENDPOINT, params, /* headers */ {}, /* auth_headers */ std::nullopt); String next_page_token; @@ -1531,6 +1596,8 @@ DB::Names RestCatalog::listTablesInNamespace(const std::string & base_namespace, if (!page_token.empty()) params.push_back({"pageToken", page_token}); + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogGetTables); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeRestCatalogGetTablesMicroseconds); auto buf = createReadBuffer( *state_snapshot, state_snapshot->config.prefix / endpoint, params, /* headers */ {}, /* auth_headers */ std::nullopt); @@ -1677,16 +1744,22 @@ bool RestCatalog::getTableMetadataImpl( const auto state_snapshot = state.get(); const std::string endpoint = std::filesystem::path(NAMESPACES_ENDPOINT) / encodeNamespaceForURI(namespace_name) / "tables" / table_name; - auto buf = createReadBuffer(*state_snapshot, state_snapshot->config.prefix / endpoint, /* params */ {}, headers, /* auth_headers */ std::nullopt); + String json_str; - if (buf->eof()) { - LOG_DEBUG(log, "Table doesn't exist (endpoint: {})", endpoint); - return false; - } + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogGetTableMetadata); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeRestCatalogGetTableMetadataMicroseconds); + auto buf = createReadBuffer( + *state_snapshot, state_snapshot->config.prefix / endpoint, /* params */ {}, headers, /* auth_headers */ std::nullopt); - String json_str; - readJSONObjectPossiblyInvalid(json_str, *buf); + if (buf->eof()) + { + LOG_DEBUG(log, "Table doesn't exist (endpoint: {})", endpoint); + return false; + } + + readJSONObjectPossiblyInvalid(json_str, *buf); + } #ifdef DEBUG_OR_SANITIZER_BUILD /// This log message might contain credentials, @@ -1763,10 +1836,6 @@ void RestCatalog::sendRequest(const CatalogState & catalog_state, const String & LOG_TEST(log, "REST catalog {} {} body ({} bytes): {}", method, endpoint, body_str.size(), body_str); - DB::HTTPHeaderEntries headers = getAuthHeaders(catalog_state, /* update_token = */ true); - headers.emplace_back("Content-Type", "application/json"); - headers.emplace_back("X-Iceberg-Access-Delegation", "vended-credentials"); - const auto & context = getContext(); DB::ReadWriteBufferFromHTTP::OutStreamCallback out_stream_callback; @@ -1781,25 +1850,62 @@ void RestCatalog::sendRequest(const CatalogState & catalog_state, const String & /// enable_url_encoding=false to allow using tables with encoded sequences in names like 'foo%2Fbar' Poco::URI url(endpoint, /* enable_url_encoding */ false); - auto wb = DB::BuilderRWBufferFromHTTP(url) - .withConnectionGroup(DB::HTTPConnectionGroupType::HTTP) - .withMethod(method) - .withSettings(context->getReadSettings()) - .withTimeouts(DB::ConnectionTimeouts::getHTTPTimeouts(context->getSettingsRef(), context->getServerSettings())) - .withHostFilter(&context->getRemoteHostFilter()) - .withHeaders(headers) - .withOutCallback(out_stream_callback) - /// Send the JSON body with an explicit Content-Length: Snowflake Horizon rejects - /// chunked transfer encoding on catalog commits with HTTP 500 and an empty body. - .withOutCallbackFixedContentLength(body_str.size()) - .withSkipNotFound(false) - .create(credentials); - - String response_str; - if (!ignore_result) - readJSONObjectPossiblyInvalid(response_str, *wb); - else - wb->ignoreAll(); + auto create_buffer = [&](bool update_token, bool & used_cached_oauth_token) + { + DB::HTTPHeaderEntries headers = getAuthHeaders(catalog_state, update_token, &used_cached_oauth_token); + headers.emplace_back("Content-Type", "application/json"); + headers.emplace_back("X-Iceberg-Access-Delegation", "vended-credentials"); + return DB::BuilderRWBufferFromHTTP(url) + .withConnectionGroup(DB::HTTPConnectionGroupType::HTTP) + .withMethod(method) + .withSettings(context->getReadSettings()) + .withTimeouts(DB::ConnectionTimeouts::getHTTPTimeouts(context->getSettingsRef(), context->getServerSettings())) + .withHostFilter(&context->getRemoteHostFilter()) + .withHeaders(headers) + .withOutCallback(out_stream_callback) + /// Send the JSON body with an explicit Content-Length: Snowflake Horizon rejects + /// chunked transfer encoding on catalog commits with HTTP 500 and an empty body. + .withOutCallbackFixedContentLength(body_str.size()) + .withSkipNotFound(false) + .create(credentials); + }; + + try + { + bool used_cached_oauth_token = false; + auto wb = create_buffer(false, used_cached_oauth_token); + + String response_str; + if (!ignore_result) + readJSONObjectPossiblyInvalid(response_str, *wb); + else + wb->ignoreAll(); + + if (used_cached_oauth_token) + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogAuthTokenCachedValid); + } + catch (const DB::HTTPException & e) + { + const auto status = e.getHTTPStatus(); + if (update_token_if_expired && + (status == Poco::Net::HTTPResponse::HTTPStatus::HTTP_UNAUTHORIZED + || status == Poco::Net::HTTPResponse::HTTPStatus::HTTP_FORBIDDEN)) + { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogUnauthorized); + bool used_cached_oauth_token_on_retry = false; + auto wb = create_buffer(true, used_cached_oauth_token_on_retry); + + String response_str; + if (!ignore_result) + readJSONObjectPossiblyInvalid(response_str, *wb); + else + wb->ignoreAll(); + } + else + { + throw; + } + } } void RestCatalog::createNamespaceIfNotExists(const String & namespace_name, const String & location) const @@ -1837,6 +1943,8 @@ void RestCatalog::createNamespaceIfNotExists(const String & namespace_name, cons try { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogCreateNamespace); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeRestCatalogCreateNamespaceMicroseconds); sendRequest(*state_snapshot, endpoint, request_body, Poco::Net::HTTPRequest::HTTP_POST, /* ignore_result */ false); } catch (const DB::HTTPException & e) @@ -1881,6 +1989,8 @@ void RestCatalog::createTable(const String & namespace_name, const String & tabl try { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogCreateTable); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeRestCatalogCreateTableMicroseconds); sendRequest(*state_snapshot, endpoint, request_body, Poco::Net::HTTPRequest::HTTP_POST, /* ignore_result */ false); } catch (const DB::HTTPException & ex) @@ -1948,6 +2058,8 @@ bool RestCatalog::updateMetadata(const String & namespace_name, const String & t try { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogUpdateTable); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeRestCatalogUpdateTableMicroseconds); sendRequest(*state_snapshot, endpoint, request_body, Poco::Net::HTTPRequest::HTTP_POST, /* ignore_result */ false); } catch (const DB::HTTPException & ex) @@ -2040,6 +2152,8 @@ void RestCatalog::dropTable(const String & namespace_name, const String & table_ Poco::JSON::Object::Ptr request_body = nullptr; try { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogDropTable); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeRestCatalogDropTableMicroseconds); sendRequest(*state_snapshot, endpoint, request_body, Poco::Net::HTTPRequest::HTTP_DELETE, true); } catch (const DB::HTTPException & ex) @@ -2128,16 +2242,22 @@ ICatalog::CredentialsRefreshCallback RestCatalog::getCredentialsConfigurationCal const auto & table = storage_id.getTableName(); auto [namespace_name, table_name] = DataLake::parseTableName(table); const std::string endpoint = std::filesystem::path(NAMESPACES_ENDPOINT) / encodeNamespaceForURI(namespace_name) / "tables" / table_name; - auto buf = createReadBuffer(*state_snapshot, state_snapshot->config.prefix / endpoint, /* params */ {}, headers, /* auth_headers */ std::nullopt); + String json_str; - if (buf->eof()) { - LOG_DEBUG(log, "Table doesn't exist (endpoint: {})", endpoint); - return nullptr; - } + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogGetCredentials); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeRestCatalogGetCredentialsMicroseconds); + auto buf = createReadBuffer( + *state_snapshot, state_snapshot->config.prefix / endpoint, /* params */ {}, headers, /* auth_headers */ std::nullopt); - String json_str; - readJSONObjectPossiblyInvalid(json_str, *buf); + if (buf->eof()) + { + LOG_DEBUG(log, "Table doesn't exist (endpoint: {})", endpoint); + return nullptr; + } + + readJSONObjectPossiblyInvalid(json_str, *buf); + } Poco::JSON::Parser parser; Poco::Dynamic::Var json = parser.parse(json_str); diff --git a/src/Databases/DataLake/RestCatalog.h b/src/Databases/DataLake/RestCatalog.h index a21474fd5081..4444c3eed5f6 100644 --- a/src/Databases/DataLake/RestCatalog.h +++ b/src/Databases/DataLake/RestCatalog.h @@ -207,9 +207,13 @@ class RestCatalog : public ICatalog, public DB::WithContext /// Load catalog config (special http handler) utilizing information from catalog_state and auth_headers. Config loadConfig(const CatalogState & catalog_state, const std::optional & auth_headers = std::nullopt); - virtual DB::HTTPHeaderEntries getAuthHeaders(const CatalogState & catalog_state, bool update_token) const; + virtual DB::HTTPHeaderEntries getAuthHeaders( + const CatalogState & catalog_state, + bool update_token, + bool * used_cached_oauth_token) const; void validateAuthHeaders(const DB::HTTPHeaderEntry & header) const; + static void parseCatalogConfigurationSettings(const Poco::JSON::Object::Ptr & object, Config & result); virtual void sendRequest( @@ -278,7 +282,10 @@ class OneLakeCatalog : public RestCatalog return DB::DatabaseDataLakeCatalogType::ICEBERG_ONELAKE; } - DB::HTTPHeaderEntries getAuthHeaders(const CatalogState & catalog_state, bool update_token) const override; + DB::HTTPHeaderEntries getAuthHeaders( + const CatalogState & catalog_state, + bool update_token, + bool * used_cached_oauth_token) const override; static void validateSettingsChanges(const DB::SettingsChanges & changes, AuthMode auth_mode); @@ -325,7 +332,10 @@ class BigLakeCatalog : public RestCatalog return DB::DatabaseDataLakeCatalogType::ICEBERG_BIGLAKE; } - DB::HTTPHeaderEntries getAuthHeaders(const CatalogState & catalog_state, bool update_token) const override; + DB::HTTPHeaderEntries getAuthHeaders( + const CatalogState & catalog_state, + bool update_token, + bool * used_cached_oauth_token) const override; const std::string & getGoogleADCClientId() const { return google_adc_client_id; } const std::string & getGoogleADCClientSecret() const { return google_adc_client_secret; } diff --git a/src/Databases/DataLake/S3TablesCredentialRefresh.cpp b/src/Databases/DataLake/S3TablesCredentialRefresh.cpp new file mode 100644 index 000000000000..8de09e4d553e --- /dev/null +++ b/src/Databases/DataLake/S3TablesCredentialRefresh.cpp @@ -0,0 +1,43 @@ +#include "config.h" + +#if USE_AVRO && USE_SSL && USE_AWS_S3 + +#include + +namespace DataLake +{ + +namespace +{ + +std::shared_ptr getCatalogIAMCredentials(Aws::Auth::AWSCredentialsProvider & provider) +{ + auto aws_creds = provider.GetAWSCredentials(); + if (aws_creds.GetAWSAccessKeyId().empty() || aws_creds.GetAWSSecretKey().empty()) + return nullptr; + return std::make_shared( + aws_creds.GetAWSAccessKeyId(), aws_creds.GetAWSSecretKey(), aws_creds.GetSessionToken()); +} + +} + +std::shared_ptr resolveS3TablesRefreshCredentials( + const ICatalog::CredentialsRefreshCallback & base_callback, + Aws::Auth::AWSCredentialsProvider & credentials_provider) +{ + if (base_callback) + { + if (auto creds = (*base_callback)()) + { + auto s3_creds = std::dynamic_pointer_cast(creds); + if (s3_creds && !s3_creds->isEmpty()) + return creds; + } + } + + return getCatalogIAMCredentials(credentials_provider); +} + +} + +#endif diff --git a/src/Databases/DataLake/S3TablesCredentialRefresh.h b/src/Databases/DataLake/S3TablesCredentialRefresh.h new file mode 100644 index 000000000000..e5b959f88860 --- /dev/null +++ b/src/Databases/DataLake/S3TablesCredentialRefresh.h @@ -0,0 +1,23 @@ +#pragma once + +#include "config.h" + +#if USE_AVRO && USE_SSL && USE_AWS_S3 + +#include +#include + +#include + +#include + +namespace DataLake +{ + +std::shared_ptr resolveS3TablesRefreshCredentials( + const ICatalog::CredentialsRefreshCallback & base_callback, + Aws::Auth::AWSCredentialsProvider & credentials_provider); + +} + +#endif diff --git a/src/Databases/DataLake/UnityCatalog.cpp b/src/Databases/DataLake/UnityCatalog.cpp index 2d78df3ec7d8..da9ece058891 100644 --- a/src/Databases/DataLake/UnityCatalog.cpp +++ b/src/Databases/DataLake/UnityCatalog.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -15,6 +16,20 @@ #include #include +namespace ProfileEvents +{ + extern const Event DataLakeUnityCatalogGetTables; + extern const Event DataLakeUnityCatalogGetTablesMicroseconds; + extern const Event DataLakeUnityCatalogGetTable; + extern const Event DataLakeUnityCatalogGetTableMicroseconds; + extern const Event DataLakeUnityCatalogGetTableMetadata; + extern const Event DataLakeUnityCatalogGetTableMetadataMicroseconds; + extern const Event DataLakeUnityCatalogGetCredentials; + extern const Event DataLakeUnityCatalogGetCredentialsMicroseconds; + extern const Event DataLakeUnityCatalogGetSchemas; + extern const Event DataLakeUnityCatalogGetSchemasMicroseconds; +} + namespace DB::ErrorCodes { extern const int DATALAKE_DATABASE_ERROR; @@ -139,7 +154,14 @@ Poco::JSON::Object::Ptr UnityCatalog::requestReadCredentials(const String & tabl request_body.set("operation", "READ"); auto callback = [&request_body] (std::ostream & os) { request_body.stringify(os); }; - auto [json, _] = postJSONRequest(TEMPORARY_CREDENTIALS_ENDPOINT, callback); + + Poco::Dynamic::Var json; + { + ProfileEvents::increment(ProfileEvents::DataLakeUnityCatalogGetCredentials); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeUnityCatalogGetCredentialsMicroseconds); + std::string _; + std::tie(json, _) = postJSONRequest(TEMPORARY_CREDENTIALS_ENDPOINT, callback); + } return json.extract(); } @@ -200,7 +222,11 @@ bool UnityCatalog::tryGetTableMetadata( std::string json_str; try { - std::tie(json, json_str) = getJSONRequest(std::filesystem::path{TABLES_ENDPOINT} / full_table_name); + { + ProfileEvents::increment(ProfileEvents::DataLakeUnityCatalogGetTableMetadata); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeUnityCatalogGetTableMetadataMicroseconds); + std::tie(json, json_str) = getJSONRequest(std::filesystem::path{TABLES_ENDPOINT} / full_table_name); + } const Poco::JSON::Object::Ptr & object = json.extract(); if (hasValueAndItsNotNone("name", object) && object->get("name").extract() == table_name) { @@ -321,7 +347,11 @@ bool UnityCatalog::existsTable(const std::string & schema_name, const std::strin Poco::Dynamic::Var json; try { - std::tie(json, json_str) = getJSONRequest(std::filesystem::path{TABLES_ENDPOINT} / (warehouse + "." + schema_name + "." + table_name)); + { + ProfileEvents::increment(ProfileEvents::DataLakeUnityCatalogGetTable); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeUnityCatalogGetTableMicroseconds); + std::tie(json, json_str) = getJSONRequest(std::filesystem::path{TABLES_ENDPOINT} / (warehouse + "." + schema_name + "." + table_name)); + } const Poco::JSON::Object::Ptr & object = json.extract(); if (hasValueAndItsNotNone("name", object) && object->get("name").extract() == table_name) return true; @@ -349,7 +379,11 @@ CatalogTables UnityCatalog::getTablesForSchema(const std::string & schema, size_ try { - std::tie(json, json_str) = getJSONRequest(TABLES_ENDPOINT, params); + { + ProfileEvents::increment(ProfileEvents::DataLakeUnityCatalogGetTables); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeUnityCatalogGetTablesMicroseconds); + std::tie(json, json_str) = getJSONRequest(TABLES_ENDPOINT, params); + } const Poco::JSON::Object::Ptr & object = json.extract(); if (!hasValueAndItsNotNone("tables", object)) @@ -416,7 +450,11 @@ DataLake::ICatalog::Namespaces UnityCatalog::getSchemas(const std::string & base try { - std::tie(json, json_str) = getJSONRequest(SCHEMAS_ENDPOINT, params); + { + ProfileEvents::increment(ProfileEvents::DataLakeUnityCatalogGetSchemas); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeUnityCatalogGetSchemasMicroseconds); + std::tie(json, json_str) = getJSONRequest(SCHEMAS_ENDPOINT, params); + } const Poco::JSON::Object::Ptr & object = json.extract(); auto schemas_object = object->get("schemas").extract(); diff --git a/src/Databases/DataLake/tests/gtest_azure_abfss_parsing.cpp b/src/Databases/DataLake/tests/gtest_azure_abfss_parsing.cpp index cc80f91b7382..58f4925b1bb8 100644 --- a/src/Databases/DataLake/tests/gtest_azure_abfss_parsing.cpp +++ b/src/Databases/DataLake/tests/gtest_azure_abfss_parsing.cpp @@ -284,4 +284,31 @@ TEST_F(AzureAbfssParsingTest, TableMetadataS3EndpointAlreadyEndsWithBucketForceA EXPECT_EQ(metadata.getLocation(), "http://minio:9000/warehouse-rest/warehouse-rest/data/testns/testtable/"); } +TEST_F(AzureAbfssParsingTest, TableMetadataGetMetadataLocationS3TablesWithAwsEndpoint) +{ + TableMetadata metadata; + metadata.withLocation(); + metadata.setLocation("s3://bucket/table-uuid/"); + metadata.setEndpoint("https://s3.us-east-2.amazonaws.com"); + + EXPECT_EQ(metadata.getLocation(), "https://s3.us-east-2.amazonaws.com/bucket/table-uuid/"); + + const std::string metadata_file = + "s3://bucket/table-uuid/metadata/v1.metadata.json"; + EXPECT_EQ(metadata.getMetadataLocation(metadata_file), "metadata/v1.metadata.json"); +} + +TEST_F(AzureAbfssParsingTest, TableMetadataGetMetadataLocationS3TablesEmptyPathWithAwsEndpoint) +{ + TableMetadata metadata; + metadata.withLocation(); + metadata.setLocation("s3://bucket"); + metadata.setEndpoint("https://s3.us-east-2.amazonaws.com"); + + EXPECT_EQ(metadata.getLocation(), "https://s3.us-east-2.amazonaws.com/bucket/"); + + const std::string metadata_file = "s3://bucket/metadata/v1.metadata.json"; + EXPECT_EQ(metadata.getMetadataLocation(metadata_file), "metadata/v1.metadata.json"); +} + } diff --git a/src/Databases/DataLake/tests/gtest_s3tables_credential_refresh.cpp b/src/Databases/DataLake/tests/gtest_s3tables_credential_refresh.cpp new file mode 100644 index 000000000000..6c5bc705fea5 --- /dev/null +++ b/src/Databases/DataLake/tests/gtest_s3tables_credential_refresh.cpp @@ -0,0 +1,112 @@ +#include "config.h" + +#if USE_AVRO && USE_SSL && USE_AWS_S3 + +#include + +#include +#include + +#include +#include + +#include + +namespace +{ + +class RotatingAWSCredentialsProvider : public Aws::Auth::AWSCredentialsProvider +{ +public: + explicit RotatingAWSCredentialsProvider(std::vector credentials_sets_) + : credentials_sets(std::move(credentials_sets_)) + { + } + + Aws::Auth::AWSCredentials GetAWSCredentials() override + { + std::lock_guard lock(mutex); + const size_t index = call_count++; + if (index >= credentials_sets.size()) + return credentials_sets.back(); + return credentials_sets[index]; + } + +private: + std::vector credentials_sets; + std::mutex mutex; + size_t call_count = 0; +}; + +} + +TEST(S3TablesCredentialRefresh, FallsBackToCatalogIAMWhenVendedCredentialsMissing) +{ + RotatingAWSCredentialsProvider provider({ + Aws::Auth::AWSCredentials("access_key_1", "secret_key_1", "session_token_1"), + Aws::Auth::AWSCredentials("access_key_2", "secret_key_2", "session_token_2"), + }); + + DataLake::ICatalog::CredentialsRefreshCallback base_callback = []() -> std::shared_ptr + { + return nullptr; + }; + + auto first = DataLake::resolveS3TablesRefreshCredentials(base_callback, provider); + ASSERT_NE(first, nullptr); + auto first_s3 = std::dynamic_pointer_cast(first); + ASSERT_NE(first_s3, nullptr); + EXPECT_EQ(first_s3->getAccessKeyId(), "access_key_1"); + EXPECT_EQ(first_s3->getSecretAccessKey(), "secret_key_1"); + EXPECT_EQ(first_s3->getSessionToken(), "session_token_1"); + + auto second = DataLake::resolveS3TablesRefreshCredentials(base_callback, provider); + ASSERT_NE(second, nullptr); + auto second_s3 = std::dynamic_pointer_cast(second); + ASSERT_NE(second_s3, nullptr); + EXPECT_EQ(second_s3->getAccessKeyId(), "access_key_2"); + EXPECT_EQ(second_s3->getSecretAccessKey(), "secret_key_2"); + EXPECT_EQ(second_s3->getSessionToken(), "session_token_2"); +} + +TEST(S3TablesCredentialRefresh, PrefersVendedCredentialsWhenPresent) +{ + RotatingAWSCredentialsProvider provider({ + Aws::Auth::AWSCredentials("catalog_access", "catalog_secret", "catalog_token"), + }); + + DataLake::ICatalog::CredentialsRefreshCallback base_callback = []() -> std::shared_ptr + { + return std::make_shared("vended_access", "vended_secret", "vended_token"); + }; + + auto creds = DataLake::resolveS3TablesRefreshCredentials(base_callback, provider); + ASSERT_NE(creds, nullptr); + auto s3_creds = std::dynamic_pointer_cast(creds); + ASSERT_NE(s3_creds, nullptr); + EXPECT_EQ(s3_creds->getAccessKeyId(), "vended_access"); + EXPECT_EQ(s3_creds->getSecretAccessKey(), "vended_secret"); + EXPECT_EQ(s3_creds->getSessionToken(), "vended_token"); +} + +TEST(S3TablesCredentialRefresh, FallsBackWhenVendedCredentialsEmpty) +{ + RotatingAWSCredentialsProvider provider({ + Aws::Auth::AWSCredentials("catalog_access", "catalog_secret", "catalog_token"), + }); + + DataLake::ICatalog::CredentialsRefreshCallback base_callback = []() -> std::shared_ptr + { + return std::make_shared("", "", ""); + }; + + auto creds = DataLake::resolveS3TablesRefreshCredentials(base_callback, provider); + ASSERT_NE(creds, nullptr); + auto s3_creds = std::dynamic_pointer_cast(creds); + ASSERT_NE(s3_creds, nullptr); + EXPECT_EQ(s3_creds->getAccessKeyId(), "catalog_access"); + EXPECT_EQ(s3_creds->getSecretAccessKey(), "catalog_secret"); + EXPECT_EQ(s3_creds->getSessionToken(), "catalog_token"); +} + +#endif diff --git a/src/Databases/enableAllExperimentalSettings.cpp b/src/Databases/enableAllExperimentalSettings.cpp index 6b542bf1e641..a7ca1b28bda6 100644 --- a/src/Databases/enableAllExperimentalSettings.cpp +++ b/src/Databases/enableAllExperimentalSettings.cpp @@ -71,6 +71,7 @@ void enableAllExperimentalSettings(ContextMutablePtr context) context->setSetting("allow_dynamic_type_in_join_keys", 1); context->setSetting("allow_experimental_alias_table_engine", 1); context->setSetting("allow_experimental_database_paimon_rest_catalog", 1); + context->setSetting("allow_experimental_database_s3_tables", 1); context->setSetting("allow_experimental_object_storage_queue_hive_partitioning", 1); context->setSetting("allow_experimental_json_lazy_type_hints", 1); context->setSetting("allow_experimental_url_wildcard_from_index_pages", 1); diff --git a/src/IO/S3/tests/gtest_s3_uri.cpp b/src/IO/S3/tests/gtest_s3_uri.cpp index a429c645fca3..92556b442884 100644 --- a/src/IO/S3/tests/gtest_s3_uri.cpp +++ b/src/IO/S3/tests/gtest_s3_uri.cpp @@ -6,6 +6,8 @@ #if USE_AWS_S3 +#include + TEST(IOTestS3URI, PathStyleNoKey) { using namespace DB; @@ -38,4 +40,29 @@ TEST(IOTestS3URI, PathStyleWithKey) ASSERT_EQ(uri_with_no_key_and_with_slash.key, "key/key/key/key"); } +TEST(IOTestS3URI, ResolveS3Endpoint) +{ + using namespace DB; + + /// expandRegionToAmazonPath() goes through the SDK's endpoint provider, which requires + /// the AWS SDK (and its CRT allocator) to be initialized. In the server this happens + /// implicitly because every S3 client is created through ClientFactory; here the test may + /// be the only thing running, so initialize it explicitly. + S3::ClientFactory::instance(); + + ASSERT_EQ(S3::expandRegionToAmazonPath("us-east-1"), + "https://s3.us-east-1.amazonaws.com"); + ASSERT_EQ(S3::expandRegionToAmazonPath("eu-west-1"), + "https://s3.eu-west-1.amazonaws.com"); + + auto cn_north = S3::expandRegionToAmazonPath("cn-north-1"); + ASSERT_TRUE(cn_north.ends_with(".amazonaws.com.cn")) + << "China region should resolve to .amazonaws.com.cn suffix, got: " << cn_north; + ASSERT_TRUE(cn_north.find("cn-north-1") != std::string::npos) + << "Got: " << cn_north; + + ASSERT_EQ(S3::expandRegionToAmazonPath("us-gov-west-1"), + "https://s3.us-gov-west-1.amazonaws.com"); +} + #endif diff --git a/tests/integration/compose/docker_compose_iceberg_lakekeeper_catalog.yml b/tests/integration/compose/docker_compose_iceberg_lakekeeper_catalog.yml index 834dfcf071c9..a2f68b606c77 100644 --- a/tests/integration/compose/docker_compose_iceberg_lakekeeper_catalog.yml +++ b/tests/integration/compose/docker_compose_iceberg_lakekeeper_catalog.yml @@ -70,3 +70,27 @@ services: retries: 5 start_period: 10s cpus: 3 + + mock-oauth: + image: python:3.12-alpine + command: + - python + - -c + - | + from http.server import HTTPServer, BaseHTTPRequestHandler + import json + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_token() + def do_POST(self): + self.send_token() + def send_token(self): + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + payload = {"access_token": "test-token", "expires_in": 3600, "token_type": "Bearer"} + self.wfile.write(json.dumps(payload).encode()) + def log_message(self, format, *args): + pass + HTTPServer(("0.0.0.0", 9999), Handler).serve_forever() + cpus: 1 diff --git a/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py b/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py index c11966d21895..d8ac831a78d8 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, ) @@ -19,6 +20,7 @@ from helpers.test_tools import TSV, csv_compare BASE_URL = "http://lakekeeper:8181/catalog" +MOCK_OAUTH_URL = "http://mock-oauth:9999/token" CATALOG_NAME = "demo" WAREHOUSE_NAME = "demo" @@ -474,3 +476,74 @@ def test_invalid_auth_header_format(started_cluster): ) assert "Invalid auth header format" in str(err.value) + +def get_auth_token_profile_events(node, query_id): + node.query("SYSTEM FLUSH LOGS") + refreshed = int(node.query( + f"SELECT ProfileEvents['DataLakeRestCatalogAuthTokenRetrieve'] " + f"FROM system.query_log WHERE query_id = '{query_id}' AND type = 'QueryFinish'" + )) + cache_hits = int(node.query( + f"SELECT ProfileEvents['DataLakeRestCatalogAuthTokenCachedValid'] " + f"FROM system.query_log WHERE query_id = '{query_id}' AND type = 'QueryFinish'" + )) + return refreshed, cache_hits + + +def test_auth_token_profile_events(started_cluster): + node = started_cluster.instances["node1"] + + test_ref = f"test_auth_token_profile_events_{uuid.uuid4().hex[:8]}" + db_name = f"{test_ref}_database" + namespace = (f"{test_ref}_namespace",) + table_name = f"{test_ref}_table" + + catalog = load_catalog_impl(started_cluster) + 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), + ) + catalog.create_table( + namespace + (table_name,), + schema=schema, + properties={"write.metadata.compression-codec": "none"}, + ) + + # The catalog client is built eagerly by CREATE DATABASE (only ATTACH defers it), + # so the first access token is fetched by that query. OAuth credentials must use + # client_id:client_secret format; oauth_server_uri points to a mock token endpoint + # in docker compose. + node.query(f"DROP DATABASE IF EXISTS {db_name}") + qid_create = f"{test_ref}-create-{uuid.uuid4()}" + node.query( + f""" + CREATE DATABASE {db_name} + ENGINE = DataLakeCatalog('{BASE_URL}', 'minio', '{minio_secret_key}') + SETTINGS + catalog_type = 'rest', + warehouse = 'demo', + storage_endpoint = 'http://minio1:9001/warehouse-rest', + catalog_credential = 'test:secret', + oauth_server_uri = '{MOCK_OAUTH_URL}' + """, + query_id=qid_create, + settings={"allow_experimental_database_iceberg": 1}, + ) + refreshed, _ = get_auth_token_profile_events(node, qid_create) + assert refreshed >= 1 + + # Every later catalog request reuses the cached token instead of fetching a new one. + qid1 = f"{test_ref}-show-1-{uuid.uuid4()}" + node.query(f"SHOW TABLES FROM {db_name}", query_id=qid1) + assert table_name in node.query(f"SHOW TABLES FROM {db_name}") + refreshed, cache_hits = get_auth_token_profile_events(node, qid1) + assert refreshed == 0 and cache_hits >= 1 + + qid2 = f"{test_ref}-show-2-{uuid.uuid4()}" + node.query(f"SHOW TABLES FROM {db_name}", query_id=qid2) + refreshed, cache_hits = get_auth_token_profile_events(node, qid2) + assert refreshed == 0 and cache_hits >= 1 +