Antalya-26.6: Datalake catalog auth token profile events

src/Common/ProfileEvents.cpp

@@ -1549,6 +1549,10 @@ The server successfully detected this situation and will download merged part fr
     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(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) \
@@ -1589,6 +1593,7 @@ The server successfully detected this situation and will download merged part fr
     M(ObjectStorageClusterProcessedTasks, "Number of processed tasks in ObjectStorageCluster request.", ValueType::Number) \
     M(ObjectStorageClusterWaitingMicroseconds, "Time of waiting for tasks in ObjectStorageCluster request.", ValueType::Microseconds) \
 
+
 #ifdef APPLY_FOR_EXTERNAL_EVENTS
     #define APPLY_FOR_EVENTS(M) APPLY_FOR_BUILTIN_EVENTS(M) APPLY_FOR_EXTERNAL_EVENTS(M)
 #else

src/Databases/DataLake/RestCatalog.cpp

@@ -90,6 +90,10 @@ namespace ProfileEvents
     extern const Event DataLakeRestCatalogGetTableMetadataMicroseconds;
     extern const Event DataLakeRestCatalogGetCredentials;
     extern const Event DataLakeRestCatalogGetCredentialsMicroseconds;
+    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;
@@ -313,13 +317,17 @@ DB::HTTPHeaderEntries RestCatalog::getAuthHeaders(
     const String & /*method*/,
     const Poco::URI & /*url*/,
     const DB::HTTPHeaderEntries & /*extra_headers*/,
-    const String & /*body*/) const
+    const String & /*body*/,
+    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: <scheme> <token>'.
     if (auth_header.has_value())
@@ -333,11 +341,15 @@ DB::HTTPHeaderEntries RestCatalog::getAuthHeaders(
     if (!client_id.empty())
     {
         auto current = access_token.get();
-        if (!current || update_token)
+        if (!current || update_token || current->isExpired())
         {
             access_token.set(std::make_unique<AccessToken>(retrieveAccessToken()));
             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);
@@ -390,6 +402,9 @@ String OneLakeCatalog::getBearerToken() const
 
 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:
@@ -507,23 +522,31 @@ BigLakeCatalog::BigLakeCatalog(
 
 DB::HTTPHeaderEntries BigLakeCatalog::getAuthHeaders(
     bool update_token,
-    const String & /*method*/,
-    const Poco::URI & /*url*/,
-    const DB::HTTPHeaderEntries & /*extra_headers*/,
-    const String & /*body*/) const
+    const String & method,
+    const Poco::URI & url,
+    const DB::HTTPHeaderEntries & extra_headers,
+    const String & body,
+    bool * used_cached_oauth_token) const
 {
     /// Google Cloud OAuth2 for BigLake.
     /// Uses GCP metadata service or Application Default Credentials to get access token.
     /// Only use Google OAuth if explicitly configured (google_project_id or google_adc_client_id).
     /// 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<AccessToken>(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);
@@ -542,7 +565,7 @@ DB::HTTPHeaderEntries BigLakeCatalog::getAuthHeaders(
         return headers;
     }
 
-    return RestCatalog::getAuthHeaders(update_token);
+    return RestCatalog::getAuthHeaders(update_token, method, url, extra_headers, body, used_cached_oauth_token);
 }
 
 AccessToken BigLakeCatalog::retrieveGoogleCloudAccessTokenFromRefreshToken() const
@@ -564,6 +587,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
@@ -682,9 +708,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 = getAuthHeaders(update_token, Poco::Net::HTTPRequest::HTTP_GET, url, headers, {});
+        auto result_headers = getAuthHeaders(update_token, Poco::Net::HTTPRequest::HTTP_GET, url, headers, {}, &used_cached_oauth_token);
         std::move(headers.begin(), headers.end(), std::back_inserter(result_headers));
 
         return DB::BuilderRWBufferFromHTTP(url)
@@ -702,7 +728,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::DataLakeRestCatalogAuthTokenCacheHits);
+        return buf;
     }
     catch (const DB::HTTPException & e)
     {
@@ -711,7 +741,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::DataLakeRestCatalogAuthTokenRefreshedOnUnauthorized);
+            bool used_cached_oauth_token_on_retry = false;
+            return create_buffer(true, used_cached_oauth_token_on_retry);
         }
         throw;
     }
@@ -1390,24 +1422,58 @@ 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, bool & used_cached_oauth_token)
+    {
+        DB::HTTPHeaderEntries headers = getAuthHeaders(update_token, method, url, extra_headers, body_str, &used_cached_oauth_token);
+        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
+    {
+        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::DataLakeRestCatalogAuthTokenCacheHits);
+    }
+    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);
+            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

src/Databases/DataLake/RestCatalog.h

@@ -222,7 +222,8 @@ class RestCatalog : public ICatalog, public DB::WithContext
         const String & method = {},
         const Poco::URI & url = {},
         const DB::HTTPHeaderEntries & extra_headers = {},
-        const String & body = {}) const;
+        const String & body = {},
+        bool * used_cached_oauth_token = nullptr) const;
 
     void validateAuthHeaders(const DB::HTTPHeaderEntry & header) const;
 
@@ -305,7 +306,8 @@ class BigLakeCatalog : public RestCatalog
         const String & method = {},
         const Poco::URI & url = {},
         const DB::HTTPHeaderEntries & extra_headers = {},
-        const String & body = {}) const override;
+        const String & body = {},
+        bool * used_cached_oauth_token = nullptr) const override;
 
     const std::string & getGoogleADCClientId() const { return google_adc_client_id; }
     const std::string & getGoogleADCClientSecret() const { return google_adc_client_secret; }

src/Databases/DataLake/S3TablesCatalog.cpp

@@ -241,7 +241,8 @@ DB::HTTPHeaderEntries S3TablesCatalog::getAuthHeaders(
     const String & method,
     const Poco::URI & url,
     const DB::HTTPHeaderEntries & extra_headers,
-    const String & body) const
+    const String & body,
+    bool * /*used_cached_oauth_token*/) const
 {
     DB::HTTPHeaderEntries all_signed;
     signRequestWithAWSV4(method, url, extra_headers, body, *signer, region, "s3tables", all_signed);

src/Databases/DataLake/S3TablesCatalog.h

@@ -51,7 +51,8 @@ class S3TablesCatalog final : public RestCatalog
         const String & method = {},
         const Poco::URI & url = {},
         const DB::HTTPHeaderEntries & extra_headers = {},
-        const String & body = {}) const override;
+        const String & body = {},
+        bool * used_cached_oauth_token = nullptr) const override;
 
 private:
     const String region;

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

tests/integration/test_database_iceberg_lakekeeper_catalog/test.py

@@ -20,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"
 
@@ -414,6 +415,76 @@ 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)
+
+    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"
+
+
 def create_int_table(catalog, namespace, table_name, rows=1):
     if namespace not in catalog.list_namespaces():
         catalog.create_namespace(namespace)