locus_surreal_adapter/
client.rs

1use anyhow::{Context, Result};
2use async_trait::async_trait;
3use locus_core_rs::storage::surrealdb::QueryParams;
4use locus_core_rs::storage::{SurrealDbClient, SurrealDbRuntimeOptions};
5use serde_json::Value;
6use surrealdb::engine::any::{Any, connect};
7use surrealdb::opt::auth::Root;
8#[cfg(feature = "native")]
9use tracing::{debug, error};
10
11pub struct RuntimeSurrealDbClient {
12    db: surrealdb::Surreal<Any>,
13}
14
15impl RuntimeSurrealDbClient {
16    pub async fn connect(
17        runtime: &SurrealDbRuntimeOptions,
18        user: Option<&str>,
19        password: Option<&str>,
20    ) -> Result<Self> {
21        let use_remote = crate::endpoint::effective_use_remote(&runtime.endpoint, runtime.use_remote);
22
23        let db = connect(runtime.endpoint.as_str()).await.with_context(|| {
24            format!(
25                "failed to connect to SurrealDB endpoint '{}'",
26                runtime.endpoint
27            )
28        })?;
29
30        if use_remote {
31            let username = user.filter(|v| !v.trim().is_empty()).unwrap_or("root");
32            let password = password.filter(|v| !v.trim().is_empty()).unwrap_or("root");
33
34            db.signin(Root {
35                username: username.to_string(),
36                password: password.to_string(),
37            })
38            .await
39            .context("failed to authenticate against remote SurrealDB")?;
40        } else if let (Some(username), Some(password)) = (
41            user.filter(|v| !v.trim().is_empty()),
42            password.filter(|v| !v.trim().is_empty()),
43        ) {
44            let _ = db
45                .signin(Root {
46                    username: username.to_string(),
47                    password: password.to_string(),
48                })
49                .await;
50        }
51
52        db.use_ns(runtime.namespace.as_str())
53            .use_db(runtime.database.as_str())
54            .await
55            .with_context(|| {
56                format!(
57                    "failed to select namespace '{}' and database '{}'",
58                    runtime.namespace, runtime.database
59                )
60            })?;
61
62        Ok(Self { db })
63    }
64
65    fn is_read_query(query: &str) -> bool {
66        query
67            .trim_start()
68            .to_ascii_uppercase()
69            .starts_with("SELECT")
70    }
71}
72
73#[async_trait]
74impl SurrealDbClient for RuntimeSurrealDbClient {
75    async fn raw_query(&self, query: &str, parameters: QueryParams) -> Result<Vec<Value>> {
76        #[cfg(feature = "native")]
77        let operation = query
78            .split_whitespace()
79            .next()
80            .unwrap_or("UNKNOWN")
81            .to_ascii_uppercase();
82        let is_read_query = Self::is_read_query(query);
83        #[cfg(feature = "native")]
84        let has_parameters = !parameters.is_empty();
85
86        let response = if parameters.is_empty() {
87            self.db.query(query).await?
88        } else {
89            self.db.query(query).bind(parameters).await?
90        };
91
92        let mut response = match response.check() {
93            Ok(value) => value,
94            Err(err) => {
95                #[cfg(feature = "native")]
96                error!(
97                    operation = %operation,
98                    read_query = is_read_query,
99                    has_parameters,
100                    error = %err,
101                    "Surreal query failed"
102                );
103                return Err(err.into());
104            }
105        };
106
107        #[cfg(feature = "native")]
108        debug!(
109            operation = %operation,
110            read_query = is_read_query,
111            has_parameters,
112            "Surreal query succeeded"
113        );
114
115        if !is_read_query {
116            return Ok(Vec::new());
117        }
118
119        if let Ok(rows) = response.take::<Vec<Value>>(0) {
120            #[cfg(feature = "native")]
121            debug!(operation = %operation, row_count = rows.len(), "Surreal read query returned rows");
122            return Ok(rows);
123        }
124
125        if let Ok(Some(row)) = response.take::<Option<Value>>(0) {
126            #[cfg(feature = "native")]
127            debug!(operation = %operation, row_count = 1, "Surreal read query returned a single row");
128            return Ok(vec![row]);
129        }
130
131        #[cfg(feature = "native")]
132        debug!(operation = %operation, row_count = 0, "Surreal read query returned no rows");
133
134        Ok(Vec::new())
135    }
136}
locus_surreal_adapter/client.rs

locus_surreal_adapter/
client.rs
