1use 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 ¶ms {
67 q = bind_pg(q, p);
68 }
69 let rows = q
70 .fetch_all(pool)
71 .await
72 .map_err(|e| Error::Database(e.to_string()))?;
73
74 if sql.contains("COUNT(") {
75 let count = rows
76 .first()
77 .and_then(|r| r.try_get::<i64, _>(0).ok())
78 .unwrap_or(0);
79 return Ok(vec![Value::Number(count.into())]);
80 }
81
82 if sql.contains("SELECT id") && !sql.contains("body") {
83 return Ok(rows
85 .iter()
86 .filter_map(|r| r.try_get::<String, _>(0).ok())
87 .map(|id| serde_json::json!({ "id": id }))
88 .collect());
89 }
90
91 let table = compiled
92 .query_string
93 .split("FROM")
94 .nth(1)
95 .and_then(|s| s.split_whitespace().next())
96 .unwrap_or(default_table);
97
98 let pairs: Vec<(String, String)> = rows
99 .iter()
100 .map(|r| {
101 let id: String = r.try_get(0).unwrap_or_default();
102 let body: Value = r.try_get(1).unwrap_or(Value::Object(Map::new()));
103 (id, body.to_string())
104 })
105 .collect();
106 expand_body_rows(table, pairs)
107}
108
109pub async fn get_record_postgres(pool: &PgPool, table: &str, id: &str) -> Result<Option<Value>> {
110 ensure_table_postgres(pool, table).await?;
111 let q = format!("SELECT id, body FROM {table} WHERE id = $1");
112 let row = sqlx::query(&q)
113 .bind(id)
114 .fetch_optional(pool)
115 .await
116 .map_err(|e| Error::Database(e.to_string()))?;
117 Ok(row.map(|r| {
118 let id: String = r.get(0);
119 let body: Value = r.get(1);
120 row_from_body(table, &id, body)
121 }))
122}
123
124pub async fn create_record_postgres(pool: &PgPool, table: &str, content: Value) -> Result<Value> {
125 ensure_table_postgres(pool, table).await?;
126 let id = storage_id(&content).unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
127 let mut body = upsert_body_fields(content);
128 body.remove("id");
129 let body_val = Value::Object(body);
130 let q = format!("INSERT INTO {table} (id, body) VALUES ($1, $2)");
131 sqlx::query(&q)
132 .bind(&id)
133 .bind(&body_val)
134 .execute(pool)
135 .await
136 .map_err(|e| Error::Database(e.to_string()))?;
137 Ok(row_from_body(table, &id, body_val))
138}
139
140pub async fn update_record_postgres(
141 pool: &PgPool,
142 table: &str,
143 id: &str,
144 content: Value,
145) -> Result<Value> {
146 if get_record_postgres(pool, table, id).await?.is_none() {
147 return Err(Error::NotFound(format!("{table}:{id}")));
148 }
149 let mut body = upsert_body_fields(content);
150 body.remove("id");
151 let body_val = Value::Object(body);
152 let q = format!("UPDATE {table} SET body = $1 WHERE id = $2");
153 sqlx::query(&q)
154 .bind(&body_val)
155 .bind(id)
156 .execute(pool)
157 .await
158 .map_err(|e| Error::Database(e.to_string()))?;
159 Ok(row_from_body(table, id, body_val))
160}
161
162pub async fn merge_record_postgres(
163 pool: &PgPool,
164 table: &str,
165 id: &str,
166 patch: Value,
167) -> Result<Value> {
168 let existing = get_record_postgres(pool, table, id)
169 .await?
170 .ok_or_else(|| Error::NotFound(format!("{table}:{id}")))?;
171 let mut base = existing.as_object().cloned().unwrap_or_default();
172 base.remove("id");
173 if let Some(patch_obj) = patch.as_object() {
174 json_merge(&mut base, patch_obj);
175 }
176 update_record_postgres(pool, table, id, Value::Object(base)).await
177}
178
179pub async fn delete_record_postgres(pool: &PgPool, table: &str, id: &str) -> Result<()> {
180 let q = format!("DELETE FROM {table} WHERE id = $1");
181 sqlx::query(&q)
182 .bind(id)
183 .execute(pool)
184 .await
185 .map_err(|e| Error::Database(e.to_string()))?;
186 Ok(())
187}
188
189pub async fn relate_edge_postgres(
190 pool: &PgPool,
191 from: &RecordId,
192 edge_table: &str,
193 to: &RecordId,
194) -> Result<()> {
195 ensure_edges_postgres(pool).await?;
196 sqlx::query(
197 "INSERT INTO valence_edges (from_table, from_id, edge_type, to_table, to_id) \
198 VALUES ($1, $2, $3, $4, $5) ON CONFLICT DO NOTHING",
199 )
200 .bind(from.table())
201 .bind(from.id())
202 .bind(edge_table)
203 .bind(to.table())
204 .bind(to.id())
205 .execute(pool)
206 .await
207 .map_err(|e| Error::Database(e.to_string()))?;
208 Ok(())
209}
210
211pub async fn unrelate_edge_postgres(
212 pool: &PgPool,
213 from: &RecordId,
214 edge_table: &str,
215 to: &RecordId,
216) -> Result<()> {
217 sqlx::query(
218 "DELETE FROM valence_edges WHERE from_table = $1 AND from_id = $2 AND edge_type = $3 \
219 AND to_table = $4 AND to_id = $5",
220 )
221 .bind(from.table())
222 .bind(from.id())
223 .bind(edge_table)
224 .bind(to.table())
225 .bind(to.id())
226 .execute(pool)
227 .await
228 .map_err(|e| Error::Database(e.to_string()))?;
229 Ok(())
230}
231
232pub async fn get_edge_targets_postgres(
233 pool: &PgPool,
234 from: &RecordId,
235 edge_table: &str,
236) -> Result<Vec<RecordId>> {
237 let rows = sqlx::query(
238 "SELECT to_table, to_id FROM valence_edges \
239 WHERE from_table = $1 AND from_id = $2 AND edge_type = $3",
240 )
241 .bind(from.table())
242 .bind(from.id())
243 .bind(edge_table)
244 .fetch_all(pool)
245 .await
246 .map_err(|e| Error::Database(e.to_string()))?;
247 Ok(rows
248 .iter()
249 .map(|r| RecordId::new(r.get::<String, _>(0), r.get::<String, _>(1)))
250 .collect())
251}
252
253pub async fn define_unique_index_postgres(pool: &PgPool, table: &str, field: &str) -> Result<()> {
254 assert_safe_table(table)?;
255 ensure_table_postgres(pool, table).await?;
256 let idx = format!("valence_unique_{table}_{field}");
257 let q = format!("CREATE UNIQUE INDEX IF NOT EXISTS {idx} ON {table} ((body->>'{field}'))");
258 sqlx::query(&q)
259 .execute(pool)
260 .await
261 .map_err(|e| Error::Database(e.to_string()))?;
262 Ok(())
263}