Skip to main content

valence_backend_sql/
postgres_ops.rs

1//! Postgres-specific SQL document operations.
2
3use serde_json::{Map, Value};
4use sqlx::postgres::PgPool;
5use sqlx::Row;
6use valence_core::error::{Error, Result};
7use valence_core::record_id::RecordId;
8
9use crate::query::prepare_compiled_postgres;
10use crate::sqlite_ops::{assert_safe_table, expand_body_rows, storage_id};
11use crate::{json_merge, row_from_body, upsert_body_fields};
12
13pub fn ensure_table_ddl_postgres(table: &str) -> String {
14    format!(
15        "CREATE TABLE IF NOT EXISTS {table} (\
16         id TEXT PRIMARY KEY NOT NULL, \
17         body JSONB NOT NULL DEFAULT '{{}}'::jsonb)"
18    )
19}
20
21pub async fn ensure_edges_postgres(pool: &PgPool) -> Result<()> {
22    sqlx::query(&crate::document::ensure_edges_table_ddl())
23        .execute(pool)
24        .await
25        .map_err(|e| Error::database(e.to_string()))?;
26    Ok(())
27}
28
29pub async fn ensure_table_postgres(pool: &PgPool, table: &str) -> Result<()> {
30    assert_safe_table(table)?;
31    sqlx::query(&ensure_table_ddl_postgres(table))
32        .execute(pool)
33        .await
34        .map_err(|e| Error::database(e.to_string()))?;
35    Ok(())
36}
37
38fn bind_pg<'q>(
39    query: sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>,
40    value: &Value,
41) -> sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments> {
42    match value {
43        Value::Null => query.bind(None::<String>),
44        Value::Bool(b) => query.bind(*b),
45        Value::Number(n) => {
46            if let Some(i) = n.as_i64() {
47                query.bind(i)
48            } else if let Some(f) = n.as_f64() {
49                query.bind(f)
50            } else {
51                query.bind(n.to_string())
52            }
53        }
54        Value::String(s) => query.bind(s.clone()),
55        other => query.bind(other.to_string()),
56    }
57}
58
59pub async fn execute_select_postgres(
60    pool: &PgPool,
61    compiled: &valence_core::compiled_query::CompiledQuery,
62    default_table: &str,
63) -> Result<Vec<Value>> {
64    let (sql, params) = prepare_compiled_postgres(compiled)?;
65    let mut q = sqlx::query(&sql);
66    for p in &params {
67        q = bind_pg(q, p);
68    }
69    let rows = match q.fetch_all(pool).await {
70        Ok(rows) => rows,
71        // 42P01 undefined_table: read-only compiled queries on a missing table
72        // return empty, matching the sqlite adapter's "no such table" behavior.
73        Err(e)
74            if e.as_database_error()
75                .and_then(|d| d.code())
76                .is_some_and(|code| code == "42P01") =>
77        {
78            return Ok(vec![]);
79        }
80        Err(e) => return Err(Error::database(e.to_string())),
81    };
82
83    if sql.contains("COUNT(") {
84        let count = rows
85            .first()
86            .and_then(|r| r.try_get::<i64, _>(0).ok())
87            .unwrap_or(0);
88        return Ok(vec![Value::Number(count.into())]);
89    }
90
91    if sql.contains("SELECT id") && !sql.contains("body") {
92        // Return object rows so callers can deserialize as `IdOnlyRecord` / model types.
93        return Ok(rows
94            .iter()
95            .filter_map(|r| r.try_get::<String, _>(0).ok())
96            .map(|id| serde_json::json!({ "id": id }))
97            .collect());
98    }
99
100    let table = compiled
101        .query_string
102        .split("FROM")
103        .nth(1)
104        .and_then(|s| s.split_whitespace().next())
105        .unwrap_or(default_table);
106
107    let pairs: Vec<(String, String)> = rows
108        .iter()
109        .map(|r| {
110            let id: String = r.try_get(0).unwrap_or_default();
111            let body: Value = r.try_get(1).unwrap_or_else(|_| Value::Object(Map::new()));
112            (id, body.to_string())
113        })
114        .collect();
115    expand_body_rows(table, pairs)
116}
117
118pub async fn get_record_postgres(pool: &PgPool, table: &str, id: &str) -> Result<Option<Value>> {
119    ensure_table_postgres(pool, table).await?;
120    let q = format!("SELECT id, body FROM {table} WHERE id = $1");
121    let row = sqlx::query(&q)
122        .bind(id)
123        .fetch_optional(pool)
124        .await
125        .map_err(|e| Error::database(e.to_string()))?;
126    Ok(row.map(|r| {
127        let id: String = r.get(0);
128        let body: Value = r.get(1);
129        row_from_body(table, &id, body)
130    }))
131}
132
133pub async fn create_record_postgres(pool: &PgPool, table: &str, content: Value) -> Result<Value> {
134    ensure_table_postgres(pool, table).await?;
135    let id = storage_id(&content).unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
136    let mut body = upsert_body_fields(content);
137    body.remove("id");
138    let body_val = Value::Object(body);
139    let q = format!("INSERT INTO {table} (id, body) VALUES ($1, $2)");
140    sqlx::query(&q)
141        .bind(&id)
142        .bind(&body_val)
143        .execute(pool)
144        .await
145        .map_err(|e| Error::database(e.to_string()))?;
146    Ok(row_from_body(table, &id, body_val))
147}
148
149pub async fn update_record_postgres(
150    pool: &PgPool,
151    table: &str,
152    id: &str,
153    content: Value,
154) -> Result<Value> {
155    if get_record_postgres(pool, table, id).await?.is_none() {
156        return Err(Error::NotFound(format!("{table}:{id}")));
157    }
158    let mut body = upsert_body_fields(content);
159    body.remove("id");
160    let body_val = Value::Object(body);
161    let q = format!("UPDATE {table} SET body = $1 WHERE id = $2");
162    sqlx::query(&q)
163        .bind(&body_val)
164        .bind(id)
165        .execute(pool)
166        .await
167        .map_err(|e| Error::database(e.to_string()))?;
168    Ok(row_from_body(table, id, body_val))
169}
170
171pub async fn merge_record_postgres(
172    pool: &PgPool,
173    table: &str,
174    id: &str,
175    patch: Value,
176) -> Result<Value> {
177    let existing = get_record_postgres(pool, table, id)
178        .await?
179        .ok_or_else(|| Error::NotFound(format!("{table}:{id}")))?;
180    let mut base = existing.as_object().cloned().unwrap_or_default();
181    base.remove("id");
182    if let Some(patch_obj) = patch.as_object() {
183        json_merge(&mut base, patch_obj);
184    }
185    update_record_postgres(pool, table, id, Value::Object(base)).await
186}
187
188pub async fn delete_record_postgres(pool: &PgPool, table: &str, id: &str) -> Result<()> {
189    let q = format!("DELETE FROM {table} WHERE id = $1");
190    sqlx::query(&q)
191        .bind(id)
192        .execute(pool)
193        .await
194        .map_err(|e| Error::database(e.to_string()))?;
195    Ok(())
196}
197
198pub async fn relate_edge_postgres(
199    pool: &PgPool,
200    from: &RecordId,
201    edge_table: &str,
202    to: &RecordId,
203) -> Result<()> {
204    ensure_edges_postgres(pool).await?;
205    sqlx::query(
206        "INSERT INTO valence_edges (from_table, from_id, edge_type, to_table, to_id) \
207         VALUES ($1, $2, $3, $4, $5) ON CONFLICT DO NOTHING",
208    )
209    .bind(from.table())
210    .bind(from.id())
211    .bind(edge_table)
212    .bind(to.table())
213    .bind(to.id())
214    .execute(pool)
215    .await
216    .map_err(|e| Error::database(e.to_string()))?;
217    Ok(())
218}
219
220pub async fn unrelate_edge_postgres(
221    pool: &PgPool,
222    from: &RecordId,
223    edge_table: &str,
224    to: &RecordId,
225) -> Result<()> {
226    sqlx::query(
227        "DELETE FROM valence_edges WHERE from_table = $1 AND from_id = $2 AND edge_type = $3 \
228         AND to_table = $4 AND to_id = $5",
229    )
230    .bind(from.table())
231    .bind(from.id())
232    .bind(edge_table)
233    .bind(to.table())
234    .bind(to.id())
235    .execute(pool)
236    .await
237    .map_err(|e| Error::database(e.to_string()))?;
238    Ok(())
239}
240
241pub async fn get_edge_targets_postgres(
242    pool: &PgPool,
243    from: &RecordId,
244    edge_table: &str,
245) -> Result<Vec<RecordId>> {
246    let rows = sqlx::query(
247        "SELECT to_table, to_id FROM valence_edges \
248         WHERE from_table = $1 AND from_id = $2 AND edge_type = $3",
249    )
250    .bind(from.table())
251    .bind(from.id())
252    .bind(edge_table)
253    .fetch_all(pool)
254    .await
255    .map_err(|e| Error::database(e.to_string()))?;
256    Ok(rows
257        .iter()
258        .map(|r| RecordId::new(r.get::<String, _>(0), r.get::<String, _>(1)))
259        .collect())
260}
261
262pub async fn define_unique_index_postgres(pool: &PgPool, table: &str, field: &str) -> Result<()> {
263    assert_safe_table(table)?;
264    valence_core::safe_ident::assert_safe_ident(field)?;
265    ensure_table_postgres(pool, table).await?;
266    let idx = format!("valence_unique_{table}_{field}");
267    let q = format!("CREATE UNIQUE INDEX IF NOT EXISTS {idx} ON {table} ((body->>'{field}'))");
268    sqlx::query(&q)
269        .execute(pool)
270        .await
271        .map_err(|e| Error::database(e.to_string()))?;
272    Ok(())
273}