Datalake catalog auth token profile events

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) \

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

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

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)