Skip to main content

uni_db/api/
session.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2024-2026 Dragonscale Team
3
4use crate::api::Uni;
5use std::collections::HashMap;
6use std::sync::Arc;
7use uni_common::Result;
8use uni_query::{ExecuteResult, QueryResult, Value};
9
10/// Builder for creating query sessions with scoped variables
11pub struct SessionBuilder<'a> {
12    db: &'a Uni,
13    variables: HashMap<String, Value>,
14}
15
16impl<'a> SessionBuilder<'a> {
17    pub fn new(db: &'a Uni) -> Self {
18        Self {
19            db,
20            variables: HashMap::new(),
21        }
22    }
23
24    /// Set a session variable
25    pub fn set<K: Into<String>, V: Into<Value>>(mut self, key: K, value: V) -> Self {
26        self.variables.insert(key.into(), value.into());
27        self
28    }
29
30    /// Build the session (variables become immutable)
31    pub fn build(self) -> Session<'a> {
32        Session {
33            db: self.db,
34            variables: Arc::new(self.variables),
35        }
36    }
37}
38
39/// A query session with scoped variables
40pub struct Session<'a> {
41    db: &'a Uni,
42    variables: Arc<HashMap<String, Value>>,
43}
44
45impl<'a> Session<'a> {
46    /// Execute a query with session variables available
47    pub async fn query(&self, cypher: &str) -> Result<QueryResult> {
48        self.query_with(cypher).execute().await
49    }
50
51    /// Execute a query with additional parameters
52    pub fn query_with(&self, cypher: &str) -> SessionQueryBuilder<'a, '_> {
53        SessionQueryBuilder {
54            session: self,
55            cypher: cypher.to_string(),
56            params: HashMap::new(),
57        }
58    }
59
60    /// Execute a mutation
61    pub async fn execute(&self, cypher: &str) -> Result<ExecuteResult> {
62        self.query_with(cypher).execute_mutation().await
63    }
64
65    /// Execute a mutation with additional parameters using a builder.
66    ///
67    /// Alias for [`query_with`](Self::query_with) that clarifies intent for mutations.
68    pub fn execute_with(&self, cypher: &str) -> SessionQueryBuilder<'a, '_> {
69        self.query_with(cypher)
70    }
71
72    /// Get a session variable value
73    pub fn get(&self, key: &str) -> Option<&Value> {
74        self.variables.get(key)
75    }
76}
77
78pub struct SessionQueryBuilder<'a, 'b> {
79    session: &'b Session<'a>,
80    cypher: String,
81    params: HashMap<String, Value>,
82}
83
84impl<'a, 'b> SessionQueryBuilder<'a, 'b> {
85    pub fn param<K: Into<String>, V: Into<Value>>(mut self, key: K, value: V) -> Self {
86        self.params.insert(key.into(), value.into());
87        self
88    }
89
90    pub async fn execute(self) -> Result<QueryResult> {
91        let params = Self::merge_params_internal(self.params, &self.session.variables);
92        self.session.db.execute_internal(&self.cypher, params).await
93    }
94
95    pub async fn execute_mutation(self) -> Result<ExecuteResult> {
96        let params = Self::merge_params_internal(self.params, &self.session.variables);
97        let before = self.session.db.get_mutation_count().await;
98        let result = self
99            .session
100            .db
101            .execute_internal(&self.cypher, params)
102            .await?;
103        let affected_rows = if result.is_empty() {
104            self.session
105                .db
106                .get_mutation_count()
107                .await
108                .saturating_sub(before)
109        } else {
110            result.len()
111        };
112        Ok(ExecuteResult { affected_rows })
113    }
114
115    fn merge_params_internal(
116        mut params: HashMap<String, Value>,
117        session_vars: &HashMap<String, Value>,
118    ) -> HashMap<String, Value> {
119        let session_map: HashMap<String, Value> = session_vars.clone();
120
121        if let Some(Value::Map(existing)) = params.get_mut("session") {
122            for (k, v) in session_map {
123                existing.entry(k).or_insert(v);
124            }
125        } else {
126            params.insert("session".to_string(), Value::Map(session_map));
127        }
128        params
129    }
130}