1 files changed, 67 insertions, 1 deletions

diff --git a/src/utils/abort_signal.rs b/src/utils/abort_signal.rs
index a13e5bb..47534ae 100644
--- a/src/utils/abort_signal.rs
+++ b/src/utils/abort_signal.rs
@@ -2,7 +2,7 @@ use anyhow::Result;
 use crossterm::event::{self, Event, KeyCode, KeyModifiers};
 use std::{
     sync::{
-        atomic::{AtomicBool, Ordering},
+        atomic::{AtomicBool, AtomicU64, Ordering},
         Arc,
     },
     time::Duration,
@@ -10,6 +10,72 @@ use std::{
 
 pub type AbortSignal = Arc<AbortSignalInner>;
 
+#[derive(Debug, Default)]
+pub struct QueryIdGen {
+    counter: AtomicU64,
+    parents: parking_lot::RwLock<std::collections::HashMap<String, u64>>,
+}
+
+impl QueryIdGen {
+    pub fn root_id(&self) -> u64 {
+        self.counter.fetch_add(1, Ordering::SeqCst)
+    }
+
+    pub fn child_id(&self, parent_path: &[u64]) -> u64 {
+        if parent_path.is_empty() {
+            return self.root_id();
+        }
+        let key = parent_path
+            .iter()
+            .map(|v| v.to_string())
+            .collect::<Vec<_>>()
+            .join("/");
+        let mut parents = self.parents.write();
+        let entry = parents.entry(key).or_insert(0);
+        let id = *entry;
+        *entry += 1;
+        id
+    }
+}
+
+pub fn format_query_path(path: &[u64]) -> String {
+    if path.is_empty() {
+        return String::new();
+    }
+    format!(
+        "[{}]",
+        path.iter()
+            .map(|v| v.to_string())
+            .collect::<Vec<_>>()
+            .join("/")
+    )
+}
+
+tokio::task_local! {
+    static QUERY_SCOPE: String;
+    static QUERY_ID_GEN: QueryIdGen;
+}
+
+pub fn current_query_scope() -> Option<String> {
+    QUERY_SCOPE
+        .try_with(|v| if v.is_empty() { None } else { Some(v.clone()) })
+        .ok()
+        .flatten()
+}
+
+pub async fn with_query_scope<F: std::future::Future>(scope: String, future: F) -> F::Output {
+    QUERY_SCOPE.scope(scope, future).await
+}
+
+pub fn next_query_id(parent_path: &[u64]) -> u64 {
+    if parent_path.is_empty() {
+        return QUERY_ID_GEN.try_with(|g| g.root_id()).unwrap_or(0);
+    }
+    QUERY_ID_GEN
+        .try_with(|g| g.child_id(parent_path))
+        .unwrap_or(0)
+}
+
 pub struct AbortSignalInner {
     ctrlc: AtomicBool,
     ctrld: AtomicBool,