From d8228d391df1aee196234f4eadd710051f45a4e0 Mon Sep 17 00:00:00 2001
From: discord9 <discord9@163.com>
Date: Thu, 16 Apr 2026 17:10:25 +0800
Subject: [PATCH] feat: query id

Signed-off-by: discord9 <discord9@163.com>
---
 src/session/src/context.rs  | 18 ++++++++-
 src/session/src/lib.rs      |  1 +
 src/session/src/query_id.rs | 76 +++++++++++++++++++++++++++++++++++++
 3 files changed, 93 insertions(+), 2 deletions(-)
 create mode 100644 src/session/src/query_id.rs

diff --git a/src/session/src/context.rs b/src/session/src/context.rs
index 7c6350dee1..756899b59e 100644
--- a/src/session/src/context.rs
+++ b/src/session/src/context.rs
@@ -31,9 +31,9 @@ use common_time::timezone::parse_timezone;
 use datafusion_common::config::ConfigOptions;
 use derive_builder::Builder;
 use sql::dialect::{Dialect, GenericDialect, GreptimeDbDialect, MySqlDialect, PostgreSqlDialect};
-use uuid::Uuid;
 
 use crate::protocol_ctx::ProtocolCtx;
+use crate::query_id::QueryId;
 use crate::session_config::{PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle};
 use crate::{MutableInner, ReadPreference};
 
@@ -44,7 +44,11 @@ const CURSOR_COUNT_WARNING_LIMIT: usize = 10;
 pub const REMOTE_QUERY_ID_EXTENSION_KEY: &str = "remote_query_id";
 
 pub fn generate_remote_query_id() -> String {
-    Uuid::now_v7().to_string()
+    generate_remote_query_id_value().to_string()
+}
+
+pub fn generate_remote_query_id_value() -> QueryId {
+    QueryId::new()
 }
 
 #[derive(Debug, Builder, Clone)]
@@ -354,6 +358,11 @@ impl QueryContext {
         self.extension(REMOTE_QUERY_ID_EXTENSION_KEY)
     }
 
+    pub fn remote_query_id_value(&self) -> Option<QueryId> {
+        self.remote_query_id()
+            .and_then(|query_id| query_id.parse().ok())
+    }
+
     pub fn extensions(&self) -> HashMap<String, String> {
         self.extensions.clone()
     }
@@ -798,9 +807,14 @@ mod test {
             .build();
 
         assert_eq!(ctx.remote_query_id(), Some(query_id));
+        assert_eq!(ctx.remote_query_id_value().unwrap().to_string(), query_id);
 
         let proto: api::v1::QueryContext = (&ctx).into();
         let restored = QueryContext::from(proto);
         assert_eq!(restored.remote_query_id(), Some(query_id));
+        assert_eq!(
+            restored.remote_query_id_value().unwrap().to_string(),
+            query_id
+        );
     }
 }
diff --git a/src/session/src/lib.rs b/src/session/src/lib.rs
index 1294a3368f..cba78a060e 100644
--- a/src/session/src/lib.rs
+++ b/src/session/src/lib.rs
@@ -15,6 +15,7 @@
 pub mod context;
 pub mod hints;
 pub mod protocol_ctx;
+pub mod query_id;
 pub mod session_config;
 pub mod table_name;
 
diff --git a/src/session/src/query_id.rs b/src/session/src/query_id.rs
new file mode 100644
index 0000000000..220e52577e
--- /dev/null
+++ b/src/session/src/query_id.rs
@@ -0,0 +1,76 @@
+// Copyright 2023 Greptime Team
+//
+// 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::fmt::{Display, Formatter};
+use std::str::FromStr;
+
+use uuid::Uuid;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct QueryId(Uuid);
+
+impl QueryId {
+    pub fn new() -> Self {
+        Self(Uuid::now_v7())
+    }
+
+    pub fn as_uuid(&self) -> &Uuid {
+        &self.0
+    }
+}
+
+impl Default for QueryId {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Display for QueryId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl FromStr for QueryId {
+    type Err = uuid::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Ok(Self(Uuid::parse_str(s)?))
+    }
+}
+
+impl From<Uuid> for QueryId {
+    fn from(value: Uuid) -> Self {
+        Self(value)
+    }
+}
+
+impl From<QueryId> for Uuid {
+    fn from(value: QueryId) -> Self {
+        value.0
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn query_id_round_trips_through_string() {
+        let query_id = QueryId::new();
+        let encoded = query_id.to_string();
+
+        assert_eq!(encoded.parse::<QueryId>().unwrap(), query_id);
+    }
+}