feat(query): add user agent and metrics for jwks request (#18535)

everpcpc · web-flow · commit 8b5ac3c94520 · 2025-08-15T13:10:01.000+08:00
* feat(query): add user agent for jwks request

* z

* z

* z

* z

* z

* z

* z

* z

* z

* z

* z

* z

* z
diff --git a/.typos.toml b/.typos.toml
@@ -15,6 +15,8 @@
 "Tke" = "Tke"
 "typ" = "typ"
 "dows" = "dows"
+# created_ons & updated_ons
+"ons" = "ons"
 
 [files]
 extend-exclude = [
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/common/metrics/src/lib.rs b/src/common/metrics/src/lib.rs
@@ -22,6 +22,7 @@ mod metrics;
 
 pub type VecLabels = Vec<(&'static str, String)>;
 
+pub use crate::metrics::auth;
 pub use crate::metrics::cache;
 pub use crate::metrics::cluster;
 pub use crate::metrics::external_server;
diff --git a/src/common/metrics/src/metrics/auth.rs b/src/common/metrics/src/metrics/auth.rs
@@ -0,0 +1,44 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::LazyLock;
+use std::time::Duration;
+
+use databend_common_base::runtime::metrics::register_counter_family;
+use databend_common_base::runtime::metrics::register_histogram_family_in_milliseconds;
+use databend_common_base::runtime::metrics::FamilyCounter;
+use databend_common_base::runtime::metrics::FamilyHistogram;
+
+use crate::VecLabels;
+
+static AUTH_JWKS_REQUESTS_COUNT: LazyLock<FamilyCounter<VecLabels>> =
+    LazyLock::new(|| register_counter_family("auth_jwks_requests_count"));
+static AUTH_JWKS_REFRESH_DURATION: LazyLock<FamilyHistogram<VecLabels>> =
+    LazyLock::new(|| register_histogram_family_in_milliseconds("auth_jwks_refresh_duration_ms"));
+
+pub fn metrics_incr_auth_jwks_requests_count(url: String, reason: String, status: u16) {
+    let labels = vec![
+        ("url", url),
+        ("reason", reason),
+        ("status", status.to_string()),
+    ];
+    AUTH_JWKS_REQUESTS_COUNT.get_or_create(&labels).inc();
+}
+
+pub fn metrics_observe_auth_jwks_refresh_duration(url: String, duration: Duration) {
+    let labels = vec![("url", url)];
+    AUTH_JWKS_REFRESH_DURATION
+        .get_or_create(&labels)
+        .observe(duration.as_millis() as f64);
+}
diff --git a/src/common/metrics/src/metrics/mod.rs b/src/common/metrics/src/metrics/mod.rs
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+pub mod auth;
 pub mod cache;
 pub mod cluster;
 pub mod external_server;
diff --git a/src/query/config/src/config.rs b/src/query/config/src/config.rs
@@ -1778,7 +1778,7 @@ pub struct QueryConfig {
     pub jwt_key_file: String,
 
     /// Interval in seconds to refresh jwks
-    #[clap(long, value_name = "VALUE", default_value = "600")]
+    #[clap(long, value_name = "VALUE", default_value = "86400")]
     pub jwks_refresh_interval: u64,
 
     /// Timeout in seconds to refresh jwks
diff --git a/src/query/config/src/inner.rs b/src/query/config/src/inner.rs
@@ -305,7 +305,7 @@ impl Default for QueryConfig {
             max_storage_io_requests: None,
             jwt_key_file: "".to_string(),
             jwt_key_files: Vec::new(),
-            jwks_refresh_interval: 600,
+            jwks_refresh_interval: 86400,
             jwks_refresh_timeout: 10,
             default_storage_format: "auto".to_string(),
             default_compression: "auto".to_string(),
diff --git a/src/query/config/src/lib.rs b/src/query/config/src/lib.rs
@@ -38,7 +38,6 @@ pub use builtin::*;
 pub use config::CacheStorageTypeConfig;
 pub use config::Commands;
 pub use config::Config;
-pub use config::QueryConfig;
 pub use config::StorageConfig;
 pub use config::StorageNetworkConfig;
 pub use global::GlobalConfig;
@@ -50,5 +49,6 @@ pub use inner::DiskCacheConfig as DiskCacheInnerConfig;
 pub use inner::DiskCacheKeyReloadPolicy;
 pub use inner::InnerConfig;
 pub use inner::MetaConfig;
+pub use inner::QueryConfig;
 pub use inner::SpillConfig;
 pub use inner::ThriftProtocol;
diff --git a/src/query/service/src/auth.rs b/src/query/service/src/auth.rs
@@ -83,12 +83,7 @@ impl AuthMgr {
 
     fn create(cfg: &InnerConfig) -> Arc<AuthMgr> {
         Arc::new(AuthMgr {
-            jwt_auth: JwtAuthenticator::create(
-                cfg.query.jwt_key_file.clone(),
-                cfg.query.jwt_key_files.clone(),
-                cfg.query.jwks_refresh_interval,
-                cfg.query.jwks_refresh_timeout,
-            ),
+            jwt_auth: JwtAuthenticator::create(&cfg.query),
         })
     }
 
diff --git a/src/query/service/tests/it/storages/testdata/configs_table_basic.txt b/src/query/service/tests/it/storages/testdata/configs_table_basic.txt
@@ -182,7 +182,7 @@ DB.Table: 'system'.'configs', Table: configs-table_id:1, ver:0, Engine: SystemCo
 | 'query'   | 'http_session_timeout_secs'                          | '14400'                                                                                                                                                                                                   | ''       |
 | 'query'   | 'internal_enable_sandbox_tenant'                     | 'false'                                                                                                                                                                                                   | ''       |
 | 'query'   | 'internal_merge_on_read_mutation'                    | 'false'                                                                                                                                                                                                   | ''       |
-| 'query'   | 'jwks_refresh_interval'                              | '600'                                                                                                                                                                                                     | ''       |
+| 'query'   | 'jwks_refresh_interval'                              | '86400'                                                                                                                                                                                                   | ''       |
 | 'query'   | 'jwks_refresh_timeout'                               | '10'                                                                                                                                                                                                      | ''       |
 | 'query'   | 'jwt_key_file'                                       | ''                                                                                                                                                                                                        | ''       |
 | 'query'   | 'jwt_key_files'                                      | ''                                                                                                                                                                                                        | ''       |
diff --git a/src/query/users/Cargo.toml b/src/query/users/Cargo.toml
@@ -26,6 +26,8 @@ databend-common-meta-cache = { workspace = true }
 databend-common-meta-kvapi = { workspace = true }
 databend-common-meta-store = { workspace = true }
 databend-common-meta-types = { workspace = true }
+databend-common-metrics = { workspace = true }
+databend-common-version = { workspace = true }
 enumflags2 = { workspace = true }
 itertools = { workspace = true }
 jwt-simple = { workspace = true }
diff --git a/src/query/users/src/jwt/authenticator.rs b/src/query/users/src/jwt/authenticator.rs
@@ -12,8 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use databend_common_config::QueryConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
+use databend_common_version::DATABEND_SEMVER;
 use jwt_simple::algorithms::ECDSAP256PublicKeyLike;
 use jwt_simple::algorithms::ES256PublicKey;
 use jwt_simple::algorithms::RS256PublicKey;
@@ -78,29 +80,33 @@ impl CustomClaims {
 }
 
 impl JwtAuthenticator {
-    pub fn create(
-        jwt_key_file: String,
-        jwt_key_files: Vec<String>,
-        jwks_refresh_interval: u64,
-        jwks_refresh_timeout: u64,
-    ) -> Option<Self> {
-        if jwt_key_file.is_empty() && jwt_key_files.is_empty() {
+    pub fn create(cfg: &QueryConfig) -> Option<Self> {
+        if cfg.jwt_key_file.is_empty() && cfg.jwt_key_files.is_empty() {
             return None;
         }
+        let user_agent = format!(
+            "Databend/{}/{}/{}",
+            *DATABEND_SEMVER,
+            cfg.tenant_id.tenant_name(),
+            cfg.cluster_id
+        );
+
         // init a vec of key store
         let mut key_stores = vec![];
-        if !jwt_key_file.is_empty() {
+        if !cfg.jwt_key_file.is_empty() {
             key_stores.push(
-                jwk::JwkKeyStore::new(jwt_key_file)
-                    .with_refresh_interval(jwks_refresh_interval)
-                    .with_refresh_timeout(jwks_refresh_timeout),
+                jwk::JwkKeyStore::new(cfg.jwt_key_file.clone())
+                    .with_user_agent(&user_agent)
+                    .with_refresh_interval(cfg.jwks_refresh_interval)
+                    .with_refresh_timeout(cfg.jwks_refresh_timeout),
             );
         }
-        for u in jwt_key_files {
+        for u in &cfg.jwt_key_files {
             key_stores.push(
-                jwk::JwkKeyStore::new(u)
-                    .with_refresh_interval(jwks_refresh_interval)
-                    .with_refresh_timeout(jwks_refresh_timeout),
+                jwk::JwkKeyStore::new(u.clone())
+                    .with_user_agent(&user_agent)
+                    .with_refresh_interval(cfg.jwks_refresh_interval)
+                    .with_refresh_timeout(cfg.jwks_refresh_timeout),
             );
         }
         Some(JwtAuthenticator { key_stores })
diff --git a/src/query/users/src/jwt/jwk.rs b/src/query/users/src/jwt/jwk.rs
@@ -22,6 +22,9 @@ use base64::engine::general_purpose;
 use base64::prelude::*;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
+use databend_common_metrics::auth::metrics_incr_auth_jwks_requests_count;
+use databend_common_metrics::auth::metrics_observe_auth_jwks_refresh_duration;
+use databend_common_version::DATABEND_SEMVER;
 use jwt_simple::prelude::ES256PublicKey;
 use jwt_simple::prelude::RS256PublicKey;
 use log::info;
@@ -35,7 +38,24 @@ use serde::Serialize;
 use super::PubKey;
 
 const JWKS_REFRESH_TIMEOUT: u64 = 10;
-const JWKS_REFRESH_INTERVAL: u64 = 600;
+const JWKS_REFRESH_INTERVAL: u64 = 86400;
+
+#[derive(Clone)]
+enum LoadKeyReason {
+    Initial,
+    Refresh,
+    Force,
+}
+
+impl LoadKeyReason {
+    fn as_str(&self) -> &str {
+        match self {
+            LoadKeyReason::Initial => "initial",
+            LoadKeyReason::Refresh => "refresh",
+            LoadKeyReason::Force => "force",
+        }
+    }
+}
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct JwkKey {
@@ -103,6 +123,7 @@ pub struct JwkKeys {
 /// error, as the key is not found in the cache. We'll try to refresh the keys and try again.
 pub struct JwkKeyStore {
     url: String,
+    user_agent: String,
     recent_cached_maps: Arc<RwLock<VecDeque<HashMap<String, PubKey>>>>,
     last_refreshed_time: RwLock<Option<Instant>>,
     last_retry_time: RwLock<Option<Instant>>,
@@ -117,6 +138,7 @@ impl JwkKeyStore {
     pub fn new(url: String) -> Self {
         Self {
             url,
+            user_agent: format!("Databend/{}", *DATABEND_SEMVER),
             recent_cached_maps: Arc::new(RwLock::new(VecDeque::new())),
             max_recent_cached_maps: 2,
             refresh_interval: Duration::from_secs(JWKS_REFRESH_INTERVAL),
@@ -137,6 +159,11 @@ impl JwkKeyStore {
         self
     }
 
+    pub fn with_user_agent(mut self, user_agent: &str) -> Self {
+        self.user_agent = user_agent.to_string();
+        self
+    }
+
     pub fn with_refresh_interval(mut self, interval: u64) -> Self {
         self.refresh_interval = Duration::from_secs(interval);
         self
@@ -164,26 +191,36 @@ impl JwkKeyStore {
 
 impl JwkKeyStore {
     #[async_backtrace::framed]
-    async fn load_keys(&self) -> Result<HashMap<String, PubKey>> {
+    async fn load_keys(&self, reason: &LoadKeyReason) -> Result<HashMap<String, PubKey>> {
         if let Some(load_keys_func) = &self.load_keys_func {
             return Ok(load_keys_func());
         }
 
         let client = reqwest::Client::builder()
+            .user_agent(&self.user_agent)
             .timeout(self.refresh_timeout)
             .build()
             .map_err(|e| {
                 ErrorCode::InvalidConfig(format!("Failed to create jwks client: {}", e))
             })?;
+        let start_time = Instant::now();
         let response = client
             .get(&self.url)
             .send()
             .await
             .map_err(|e| ErrorCode::Internal(format!("Could not download JWKS: {}", e)))?;
+        let status = response.status();
         let jwk_keys: JwkKeys = response
             .json()
             .await
             .map_err(|e| ErrorCode::InvalidConfig(format!("Failed to parse JWKS: {}", e)))?;
+        let duration = start_time.elapsed();
+        metrics_incr_auth_jwks_requests_count(
+            self.url.clone(),
+            reason.as_str().to_string(),
+            status.as_u16(),
+        );
+        metrics_observe_auth_jwks_refresh_duration(self.url.clone(), duration);
         let mut new_keys: HashMap<String, PubKey> = HashMap::new();
         for k in &jwk_keys.keys {
             new_keys.insert(k.kid.to_string(), k.get_public_key()?);
@@ -193,11 +230,30 @@ impl JwkKeyStore {
 
     #[async_backtrace::framed]
     async fn maybe_refresh_cached_keys(&self, force: bool) -> Result<()> {
-        let need_reload = force
-            || match *self.last_refreshed_time.read() {
-                None => true,
-                Some(last_refreshed_at) => last_refreshed_at.elapsed() > self.refresh_interval,
-            };
+        let mut need_reload = false;
+        let mut reason = LoadKeyReason::Refresh;
+
+        if force {
+            need_reload = true;
+            reason = LoadKeyReason::Force;
+        } else {
+            match *self.last_refreshed_time.read() {
+                None => {
+                    need_reload = true;
+                    reason = LoadKeyReason::Initial;
+                }
+                Some(last_refreshed_at) => {
+                    if last_refreshed_at.elapsed() > self.refresh_interval {
+                        need_reload = true;
+                        reason = LoadKeyReason::Refresh;
+                    }
+                }
+            }
+        }
+
+        if !need_reload {
+            return Ok(());
+        }
 
         let old_keys = self
             .recent_cached_maps
@@ -206,15 +262,17 @@ impl JwkKeyStore {
             .last()
             .cloned()
             .unwrap_or(HashMap::new());
-        if !need_reload {
-            return Ok(());
-        }
 
         // if got network issues on loading JWKS, fallback to the cached keys if available
-        let new_keys = match self.load_keys().await {
+        let new_keys = match self.load_keys(&reason).await {
             Ok(new_keys) => new_keys,
             Err(err) => {
-                warn!("Failed to load JWKS: {}", err);
+                metrics_incr_auth_jwks_requests_count(
+                    self.url.clone(),
+                    reason.as_str().to_string(),
+                    err.code(),
+                );
+                warn!("Failed to load JWKS from {}: {}", self.url, err);
                 if !old_keys.is_empty() {
                     return Ok(());
                 }
@@ -232,7 +290,7 @@ impl JwkKeyStore {
         if new_keys.keys().eq(old_keys.keys()) {
             return Ok(());
         }
-        info!("JWKS keys changed.");
+        info!("JWKS keys from {} changed.", self.url);
 
         // append the new keys to the end of recent_cached_maps
         {
diff --git a/src/query/users/tests/it/jwt/authenticator.rs b/src/query/users/tests/it/jwt/authenticator.rs
