1use persistant_relational::{Dialect, RelError, RelationalBackend, RelationalTx, Row, Value};
11use sqlx::postgres::{PgArguments, PgPoolOptions, PgRow};
12use sqlx::query::Query;
13use sqlx::{Column, Postgres, Row as SqlxRow, TypeInfo, ValueRef};
14
15fn classify(e: sqlx::Error) -> RelError {
17 match &e {
18 sqlx::Error::Database(db) => {
19 let constraint = db.constraint().map(str::to_string);
20 match db.code().as_deref() {
21 Some("23505") => RelError::UniqueViolation { constraint },
22 Some("23503") => RelError::ForeignKeyViolation { constraint },
23 Some("23514") | Some("23502") => RelError::CheckViolation { constraint },
24 Some("40001") => RelError::Serialization,
25 Some("40P01") => RelError::Deadlock,
26 Some(c) if c.starts_with("08") || c.starts_with("57P") => {
27 RelError::Unavailable(e.to_string())
28 }
29 _ => RelError::Backend(e.to_string()),
30 }
31 }
32 sqlx::Error::Io(_) | sqlx::Error::PoolTimedOut | sqlx::Error::PoolClosed => {
33 RelError::Unavailable(e.to_string())
34 }
35 _ => RelError::Backend(e.to_string()),
36 }
37}
38
39fn bind<'q>(
40 mut q: Query<'q, Postgres, PgArguments>,
41 params: &'q [Value],
42) -> Result<Query<'q, Postgres, PgArguments>, RelError> {
43 for p in params {
44 q = match p {
45 Value::Null => q.bind(Option::<String>::None),
48 Value::Bool(b) => q.bind(*b),
49 Value::Int(n) => q.bind(*n),
50 Value::Float(f) => q.bind(*f),
51 Value::Text(s) => q.bind(s.as_str()),
52 Value::Bytes(b) => q.bind(b.as_slice()),
53 other => {
54 return Err(RelError::Unsupported(format!(
55 "parameter {other:?} has no PostgreSQL binding in this driver"
56 )));
57 }
58 };
59 }
60 Ok(q)
61}
62
63fn decode(row: &PgRow) -> Result<Row, RelError> {
64 let mut values = Vec::with_capacity(row.columns().len());
65 for (i, col) in row.columns().iter().enumerate() {
66 let raw = row.try_get_raw(i).map_err(classify)?;
67 if raw.is_null() {
68 values.push(Value::Null);
69 continue;
70 }
71 let name = col.type_info().name().to_string();
72 let v = match name.as_str() {
73 "BOOL" => Value::Bool(row.try_get::<bool, _>(i).map_err(classify)?),
74 "INT2" => Value::Int(i64::from(row.try_get::<i16, _>(i).map_err(classify)?)),
75 "INT4" => Value::Int(i64::from(row.try_get::<i32, _>(i).map_err(classify)?)),
76 "INT8" => Value::Int(row.try_get::<i64, _>(i).map_err(classify)?),
77 "FLOAT4" => Value::Float(f64::from(row.try_get::<f32, _>(i).map_err(classify)?)),
78 "FLOAT8" => Value::Float(row.try_get::<f64, _>(i).map_err(classify)?),
79 "TEXT" | "VARCHAR" | "NAME" | "BPCHAR" => {
80 Value::Text(row.try_get::<String, _>(i).map_err(classify)?)
81 }
82 "BYTEA" => Value::Bytes(row.try_get::<Vec<u8>, _>(i).map_err(classify)?),
83 other => {
84 return Err(RelError::Unsupported(format!(
85 "column {:?} has type {other}, which Value cannot carry; use PgBackend::pool",
86 col.name()
87 )));
88 }
89 };
90 values.push(v);
91 }
92 Ok(Row::new(values))
93}
94
95#[derive(Debug, Clone)]
97pub struct PgBackend {
98 pool: sqlx::PgPool,
99 version: String,
100}
101
102impl PgBackend {
103 pub async fn connect(url: &str) -> Result<Self, RelError> {
105 let pool = PgPoolOptions::new()
106 .max_connections(16)
107 .connect(url)
108 .await
109 .map_err(classify)?;
110 Self::from_pool(pool).await
111 }
112
113 pub async fn from_pool(pool: sqlx::PgPool) -> Result<Self, RelError> {
115 let full: String = sqlx::query_scalar("SHOW server_version")
116 .fetch_one(&pool)
117 .await
118 .map_err(classify)?;
119 let version = full.split_whitespace().next().unwrap_or("").to_string();
121 Ok(PgBackend { pool, version })
122 }
123
124 pub fn pool(&self) -> &sqlx::PgPool {
126 &self.pool
127 }
128}
129
130impl RelationalBackend for PgBackend {
131 type Tx<'a> = PgTx<'a>;
132
133 fn backend(&self) -> &str {
134 "postgres"
135 }
136
137 fn version(&self) -> &str {
138 &self.version
139 }
140
141 fn dialect(&self) -> Dialect {
142 Dialect::Postgres
143 }
144
145 async fn execute(&self, sql: &str, params: &[Value]) -> Result<u64, RelError> {
146 let q = bind(sqlx::query(sqlx::AssertSqlSafe(sql.to_string())), params)?;
147 q.execute(&self.pool)
148 .await
149 .map(|r| r.rows_affected())
150 .map_err(classify)
151 }
152
153 async fn query(&self, sql: &str, params: &[Value]) -> Result<Vec<Row>, RelError> {
154 let q = bind(sqlx::query(sqlx::AssertSqlSafe(sql.to_string())), params)?;
155 let rows = q.fetch_all(&self.pool).await.map_err(classify)?;
156 rows.iter().map(decode).collect()
157 }
158
159 async fn begin(&self) -> Result<PgTx<'_>, RelError> {
160 Ok(PgTx {
161 tx: self.pool.begin().await.map_err(classify)?,
162 })
163 }
164}
165
166#[derive(Debug)]
168pub struct PgTx<'a> {
169 tx: sqlx::Transaction<'a, Postgres>,
170}
171
172impl RelationalTx for PgTx<'_> {
173 async fn execute(&mut self, sql: &str, params: &[Value]) -> Result<u64, RelError> {
174 let q = bind(sqlx::query(sqlx::AssertSqlSafe(sql.to_string())), params)?;
175 q.execute(&mut *self.tx)
176 .await
177 .map(|r| r.rows_affected())
178 .map_err(classify)
179 }
180
181 async fn query(&mut self, sql: &str, params: &[Value]) -> Result<Vec<Row>, RelError> {
182 let q = bind(sqlx::query(sqlx::AssertSqlSafe(sql.to_string())), params)?;
183 let rows = q.fetch_all(&mut *self.tx).await.map_err(classify)?;
184 rows.iter().map(decode).collect()
185 }
186
187 async fn commit(self) -> Result<(), RelError> {
188 self.tx.commit().await.map_err(classify)
189 }
190
191 async fn rollback(self) -> Result<(), RelError> {
192 self.tx.rollback().await.map_err(classify)
193 }
194}