diff --git a/src/Common/ProfileEvents.cpp b/src/Common/ProfileEvents.cpp index 0828ee9dc92f..4fd608806698 100644 --- a/src/Common/ProfileEvents.cpp +++ b/src/Common/ProfileEvents.cpp @@ -1431,6 +1431,10 @@ The server successfully detected this situation and will download merged part fr M(DataLakeRestCatalogGetCredentialsMicroseconds, "Total time of 'get credentials' requests to Iceberg REST catalog.", ValueType::Microseconds) \ 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) \ + M(DataLakeRestCatalogAuthTokenCacheHits, "Number of requests to Iceberg REST catalog that reused a cached access token and did not fetch a new one.", ValueType::Number) \ + M(DataLakeRestCatalogAuthTokenRefreshed, "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(DataLakeRestCatalogAuthTokenRefreshedOnUnauthorized, "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) \ diff --git a/src/Databases/DataLake/RestCatalog.cpp b/src/Databases/DataLake/RestCatalog.cpp index e030c873a152..728c062ea28a 100644 --- a/src/Databases/DataLake/RestCatalog.cpp +++ b/src/Databases/DataLake/RestCatalog.cpp @@ -73,6 +73,10 @@ namespace ProfileEvents extern const Event DataLakeRestCatalogGetCredentialsMicroseconds; extern const Event DataLakeRestCatalogCredentialsVended; extern const Event DataLakeRestCatalogCredentialsCacheHits; + extern const Event DataLakeRestCatalogAuthTokenCacheHits; + extern const Event DataLakeRestCatalogAuthTokenRefreshed; + extern const Event DataLakeRestCatalogAuthTokenRefreshedMicroseconds; + extern const Event DataLakeRestCatalogAuthTokenRefreshedOnUnauthorized; extern const Event DataLakeRestCatalogCreateNamespace; extern const Event DataLakeRestCatalogCreateNamespaceMicroseconds; extern const Event DataLakeRestCatalogCreateTable; @@ -274,10 +278,14 @@ DB::HTTPHeaderEntries RestCatalog::getAuthHeaders( /// https://github.com/apache/iceberg/blob/3badfe0c1fcf0c0adfc7aa4a10f0b50365c48cf9/open-api/rest-catalog-open-api.yaml#L3498C5-L3498C34 if (!client_id.empty()) { - if (!access_token.has_value() || update_token) + if (!access_token.has_value() || update_token || access_token->isExpired()) { access_token = retrieveAccessToken(); } + else + { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogAuthTokenCacheHits); + } DB::HTTPHeaderEntries headers; headers.emplace_back("Authorization", "Bearer " + access_token.value().token); @@ -313,6 +321,9 @@ OneLakeCatalog::OneLakeCatalog( AccessToken RestCatalog::retrieveAccessToken() const { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogAuthTokenRefreshed); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeRestCatalogAuthTokenRefreshedMicroseconds); + static constexpr auto oauth_tokens_endpoint = "oauth/tokens"; /// TODO: @@ -444,6 +455,10 @@ DB::HTTPHeaderEntries BigLakeCatalog::getAuthHeaders( { access_token = retrieveGoogleCloudAccessToken(); } + else + { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogAuthTokenCacheHits); + } DB::HTTPHeaderEntries headers; headers.emplace_back("Authorization", "Bearer " + access_token->token); @@ -484,6 +499,9 @@ AccessToken BigLakeCatalog::retrieveGoogleCloudAccessTokenFromRefreshToken() con AccessToken BigLakeCatalog::retrieveGoogleCloudAccessToken() const { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogAuthTokenRefreshed); + auto timer = DB::CurrentThread::getProfileEvents().timer(ProfileEvents::DataLakeRestCatalogAuthTokenRefreshedMicroseconds); + if (!google_adc_client_id.empty() && !google_adc_client_secret.empty() && !google_adc_refresh_token.empty()) { try @@ -614,6 +632,7 @@ DB::ReadWriteBufferFromHTTPPtr RestCatalog::createReadBuffer( (status == Poco::Net::HTTPResponse::HTTPStatus::HTTP_UNAUTHORIZED || status == Poco::Net::HTTPResponse::HTTPStatus::HTTP_FORBIDDEN)) { + ProfileEvents::increment(ProfileEvents::DataLakeRestCatalogAuthTokenRefreshedOnUnauthorized); return create_buffer(true); } throw; @@ -1077,24 +1096,53 @@ void RestCatalog::sendRequest(const String & endpoint, Poco::JSON::Object::Ptr r DB::HTTPHeaderEntries extra_headers; extra_headers.emplace_back("Content-Type", "application/json"); - DB::HTTPHeaderEntries headers = getAuthHeaders(/* update_token = */ true, method, url, extra_headers, body_str); - headers.emplace_back("Content-Type", "application/json"); - 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) - .withSkipNotFound(false) - .create(credentials); - - String response_str; - if (!ignore_result) - readJSONObjectPossiblyInvalid(response_str, *wb); - else - wb->ignoreAll(); + auto create_buffer = [&](bool update_token) + { + DB::HTTPHeaderEntries headers = getAuthHeaders(update_token, method, url, extra_headers, body_str); + headers.emplace_back("Content-Type", "application/json"); + 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) + .withSkipNotFound(false) + .create(credentials); + }; + + try + { + auto wb = create_buffer(false); + + String response_str; + if (!ignore_result) + readJSONObjectPossiblyInvalid(response_str, *wb); + else + wb->ignoreAll(); + } + 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::DataLakeRestCatalogAuthTokenRefreshedOnUnauthorized); + auto wb = create_buffer(true); + + 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 diff --git a/tests/integration/compose/docker_compose_iceberg_lakekeeper_catalog.yml b/tests/integration/compose/docker_compose_iceberg_lakekeeper_catalog.yml index 8adf5b43b4b9..a3cc6eaa7850 100644 --- a/tests/integration/compose/docker_compose_iceberg_lakekeeper_catalog.yml +++ b/tests/integration/compose/docker_compose_iceberg_lakekeeper_catalog.yml @@ -104,3 +104,27 @@ services: /usr/bin/mc policy set public minio/warehouse-rest; " # FIX 3: Removed 'tail -f /dev/null' to make this a one-shot setup task. 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 90a71b7ced26..e640cf79b6f3 100644 --- a/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py +++ b/tests/integration/test_database_iceberg_lakekeeper_catalog/test.py @@ -23,6 +23,7 @@ BASE_URL_LOCAL = "http://localhost:8181/catalog" BASE_URL = "http://lakekeeper:8181/catalog" +MOCK_OAUTH_URL = "http://mock-oauth:9999/token" CATALOG_NAME = "demo" WAREHOUSE_NAME = "demo" @@ -399,6 +400,66 @@ def get_credentials_profile_events(node, query_id): return vended, hits +def get_auth_token_profile_events(node, query_id): + node.query("SYSTEM FLUSH LOGS") + refreshed = int(node.query( + f"SELECT ProfileEvents['DataLakeRestCatalogAuthTokenRefreshed'] " + f"FROM system.query_log WHERE query_id = '{query_id}' AND type = 'QueryFinish'" + )) + cache_hits = int(node.query( + f"SELECT ProfileEvents['DataLakeRestCatalogAuthTokenCacheHits'] " + 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 initialized lazily on the first database access, + # not during CREATE DATABASE. OAuth credentials must use client_id:client_secret + # format; oauth_server_uri points to a mock token endpoint in docker compose. + create_clickhouse_iceberg_database( + started_cluster, + node, + db_name, + additional_settings={ + "catalog_credential": "test:secret", + "oauth_server_uri": MOCK_OAUTH_URL, + }, + ) + + 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 >= 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 + + def test_vended_credentials_cache(started_cluster): node = started_cluster.instances["node1"] catalog = load_catalog_impl(started_cluster)