Skip to main content

valence_backend_sql/
sqlite_ops.rs

1//! Shared SQL document backend logic for SQLite and Postgres.
2
3use serde_json::{Map, Value};
4use sqlx::Row;
5use valence_core::compiled_query::CompiledQuery;
6use valence_core::error::{Error, Result};
7use valence_core::record_id::RecordId;
8
9use crate::{ensure_table, json_merge, prepare_compiled, row_from_body, upsert_body_fields};
10
11/// Dialect-specific SQL fragments.
12#[allow(dead_code)]
13pub trait SqlDialect: Send + Sync + 'static {
14    fn insert_or_ignore(&self) -> &'static str;
15    fn body_column_type(&self) -> &'static str;
16}
17
18pub fn assert_safe_table(table: &str) -> Result<()> {
19    if table.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
20        Ok(())
21    } else {
22        Err(Error::Validation(format!("unsafe table name: {table}")))
23    }
24}
25
26pub fn storage_id(content: &Value) -> Option<String> {
27    content.get("id").and_then(|v| {
28        v.get("id")
29            .and_then(|x| x.as_str())
30            .map(str::to_string)
31            .or_else(|| v.as_str().map(str::to_string))
32    })
33}
34
35pub fn expand_body_rows(table: &str, rows: Vec<(String, String)>) -> Result<Vec<Value>> {
36    rows.into_iter()
37        .map(|(id, body_text)| {
38            let body: Value =
39                serde_json::from_str(&body_text).unwrap_or_else(|_| Value::Object(Map::new()));
40            Ok(row_from_body(table, &id, body))
41        })
42        .collect()
43}
44
45pub fn bind_json_value<'q>(
46    query: sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'q>>,
47    value: &Value,
48) -> sqlx::query::Query<'q, sqlx::Sqlite, sqlx::sqlite::SqliteArguments<'q>> {
49    match value {
50        Value::Null => query.bind(None::<String>),
51        Value::Bool(b) => query.bind(*b),
52        Value::Number(n) => {
53            if let Some(i) = n.as_i64() {
54                query.bind(i)
55            } else if let Some(f) = n.as_f64() {
56                query.bind(f)
57            } else {
58                query.bind(n.to_string())
59            }
60        }
61        Value::String(s) => query.bind(s.clone()),
62        other => query.bind(other.to_string()),
63    }
64}
65
66pub async fn execute_select_sqlite(
67    pool: &sqlx::SqlitePool,
68    compiled: &CompiledQuery,
69    default_table: &str,
70) -> Result<Vec<Value>> {
71    let (sql, params) = prepare_compiled(compiled)?;
72    let mut q = sqlx::query(&sql);
73    for p in &params {
74        q = bind_json_value(q, p);
75    }
76    let rows = match q.fetch_all(pool).await {
77        Ok(rows) => rows,
78        Err(e) if e.to_string().to_lowercase().contains("no such table") => return Ok(vec![]),
79        Err(e) => return Err(Error::Database(e.to_string())),
80    };
81
82    if sql.contains("COUNT(") {
83        let count = rows
84            .first()
85            .and_then(|r| r.try_get::<i64, _>(0).ok())
86            .unwrap_or(0);
87        return Ok(vec![Value::Number(count.into())]);
88    }
89
90    if sql.contains("SELECT id") && !sql.contains("body") {
91        // Return object rows so callers can deserialize as `IdOnlyRecord` / model types.
92        return Ok(rows
93            .iter()
94            .filter_map(|r| r.try_get::<String, _>(0).ok())
95            .map(|id| serde_json::json!({ "id": id }))
96            .collect());
97    }
98
99    let table = compiled
100        .query_string
101        .split("FROM")
102        .nth(1)
103        .and_then(|s| s.split_whitespace().next())
104        .unwrap_or(default_table);
105
106    let pairs: Vec<(String, String)> = rows
107        .iter()
108        .map(|r| {
109            let id: String = r.try_get(0).unwrap_or_default();
110            let body: String = r.try_get(1).unwrap_or_else(|_| "{}".into());
111            (id, body)
112        })
113        .collect();
114    expand_body_rows(table, pairs)
115}
116
117pub async fn ensure_edges_sqlite(pool: &sqlx::SqlitePool) -> Result<()> {
118    sqlx::query(&crate::document::ensure_edges_table_ddl())
119        .execute(pool)
120        .await
121        .map_err(|e| Error::Database(e.to_string()))?;
122    Ok(())
123}
124
125pub async fn ensure_table_sqlite(pool: &sqlx::SqlitePool, table: &str) -> Result<()> {
126    assert_safe_table(table)?;
127    sqlx::query(&ensure_table(table))
128        .execute(pool)
129        .await
130        .map_err(|e| Error::Database(e.to_string()))?;
131    Ok(())
132}
133
134pub async fn get_record_sqlite(
135    pool: &sqlx::SqlitePool,
136    table: &str,
137    id: &str,
138) -> Result<Option<Value>> {
139    ensure_table_sqlite(pool, table).await?;
140    let q = format!("SELECT id, body FROM {table} WHERE id = ?");
141    let row = sqlx::query(&q)
142        .bind(id)
143        .fetch_optional(pool)
144        .await
145        .map_err(|e| Error::Database(e.to_string()))?;
146    Ok(row.map(|r| {
147        let id: String = r.get(0);
148        let body: String = r.get(1);
149        row_from_body(
150            table,
151            &id,
152            serde_json::from_str(&body).unwrap_or(Value::Object(Map::new())),
153        )
154    }))
155}
156
157pub async fn create_record_sqlite(
158    pool: &sqlx::SqlitePool,
159    table: &str,
160    content: Value,
161) -> Result<Value> {
162    ensure_table_sqlite(pool, table).await?;
163    let id = storage_id(&content).unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
164    let mut body = upsert_body_fields(content);
165    body.remove("id");
166    let body_text = serde_json::to_string(&Value::Object(body))
167        .map_err(|e| Error::Serialization(e.to_string()))?;
168    let q = format!("INSERT INTO {table} (id, body) VALUES (?, ?)");
169    sqlx::query(&q)
170        .bind(&id)
171        .bind(&body_text)
172        .execute(pool)
173        .await
174        .map_err(|e| Error::Database(e.to_string()))?;
175    Ok(row_from_body(
176        table,
177        &id,
178        serde_json::from_str(&body_text).unwrap_or(Value::Object(Map::new())),
179    ))
180}
181
182pub async fn update_record_sqlite(
183    pool: &sqlx::SqlitePool,
184    table: &str,
185    id: &str,
186    content: Value,
187) -> Result<Value> {
188    if get_record_sqlite(pool, table, id).await?.is_none() {
189        return Err(Error::NotFound(format!("{table}:{id}")));
190    }
191    let mut body = upsert_body_fields(content);
192    body.remove("id");
193    let body_text = serde_json::to_string(&Value::Object(body))
194        .map_err(|e| Error::Serialization(e.to_string()))?;
195    let q = format!("UPDATE {table} SET body = ? WHERE id = ?");
196    sqlx::query(&q)
197        .bind(&body_text)
198        .bind(id)
199        .execute(pool)
200        .await
201        .map_err(|e| Error::Database(e.to_string()))?;
202    Ok(row_from_body(
203        table,
204        id,
205        serde_json::from_str(&body_text).unwrap_or(Value::Object(Map::new())),
206    ))
207}
208
209pub async fn merge_record_sqlite(
210    pool: &sqlx::SqlitePool,
211    table: &str,
212    id: &str,
213    patch: Value,
214) -> Result<Value> {
215    let existing = get_record_sqlite(pool, table, id)
216        .await?
217        .ok_or_else(|| Error::NotFound(format!("{table}:{id}")))?;
218    let mut base = existing.as_object().cloned().unwrap_or_default();
219    base.remove("id");
220    if let Some(patch_obj) = patch.as_object() {
221        json_merge(&mut base, patch_obj);
222    }
223    update_record_sqlite(pool, table, id, Value::Object(base)).await
224}
225
226pub async fn delete_record_sqlite(pool: &sqlx::SqlitePool, table: &str, id: &str) -> Result<()> {
227    let q = format!("DELETE FROM {table} WHERE id = ?");
228    sqlx::query(&q)
229        .bind(id)
230        .execute(pool)
231        .await
232        .map_err(|e| Error::Database(e.to_string()))?;
233    Ok(())
234}
235
236pub async fn relate_edge_sqlite(
237    pool: &sqlx::SqlitePool,
238    from: &RecordId,
239    edge_table: &str,
240    to: &RecordId,
241) -> Result<()> {
242    ensure_edges_sqlite(pool).await?;
243    sqlx::query(
244        "INSERT OR IGNORE INTO valence_edges (from_table, from_id, edge_type, to_table, to_id) \
245         VALUES (?, ?, ?, ?, ?)",
246    )
247    .bind(from.table())
248    .bind(from.id())
249    .bind(edge_table)
250    .bind(to.table())
251    .bind(to.id())
252    .execute(pool)
253    .await
254    .map_err(|e| Error::Database(e.to_string()))?;
255    Ok(())
256}
257
258pub async fn unrelate_edge_sqlite(
259    pool: &sqlx::SqlitePool,
260    from: &RecordId,
261    edge_table: &str,
262    to: &RecordId,
263) -> Result<()> {
264    sqlx::query(
265        "DELETE FROM valence_edges WHERE from_table = ? AND from_id = ? AND edge_type = ? \
266         AND to_table = ? AND to_id = ?",
267    )
268    .bind(from.table())
269    .bind(from.id())
270    .bind(edge_table)
271    .bind(to.table())
272    .bind(to.id())
273    .execute(pool)
274    .await
275    .map_err(|e| Error::Database(e.to_string()))?;
276    Ok(())
277}
278
279pub async fn get_edge_targets_sqlite(
280    pool: &sqlx::SqlitePool,
281    from: &RecordId,
282    edge_table: &str,
283) -> Result<Vec<RecordId>> {
284    let rows = sqlx::query(
285        "SELECT to_table, to_id FROM valence_edges \
286         WHERE from_table = ? AND from_id = ? AND edge_type = ?",
287    )
288    .bind(from.table())
289    .bind(from.id())
290    .bind(edge_table)
291    .fetch_all(pool)
292    .await
293    .map_err(|e| Error::Database(e.to_string()))?;
294    Ok(rows
295        .iter()
296        .map(|r| RecordId::new(r.get::<String, _>(0), r.get::<String, _>(1)))
297        .collect())
298}
299
300pub async fn define_unique_index_sqlite(
301    pool: &sqlx::SqlitePool,
302    table: &str,
303    field: &str,
304) -> Result<()> {
305    assert_safe_table(table)?;
306    if !field.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
307        return Err(Error::Validation(format!("unsafe field: {field}")));
308    }
309    ensure_table_sqlite(pool, table).await?;
310    let idx = format!("valence_unique_{table}_{field}");
311    let q = format!(
312        "CREATE UNIQUE INDEX IF NOT EXISTS {idx} ON {table} (json_extract(body, '$.{field}'))"
313    );
314    sqlx::query(&q)
315        .execute(pool)
316        .await
317        .map_err(|e| Error::Database(e.to_string()))?;
318    Ok(())
319}
320
321pub fn ttl_deferred() -> valence_core::ttl::BackendTtlCapability {
322    valence_core::ttl::BackendTtlCapability::Deferred
323}
324
325#[allow(dead_code, clippy::unused_async)]
326pub async fn apply_ttl_noop(
327    _table: &str,
328    _policy: &valence_core::ttl::SchemaTtlPolicy,
329) -> Result<()> {
330    Ok(())
331}
332
333pub const fn sql_capabilities(label: &'static str) -> valence_core::BackendCapabilities {
334    valence_core::BackendCapabilities {
335        supports_merge: true,
336        supports_graph_edges: true,
337        telemetry_label: label,
338    }
339}