use sqlx::postgres::{PgPool, PgPoolOptions, PgRow};
use tokio::sync::Mutex;
use crate::error::FlozError;
use crate::value::Value;
#[derive(Debug, Clone)]
pub struct Db {
pool: PgPool,
}
impl Db {
pub async fn connect(url: &str) -> Result<Self, FlozError> {
let pool = PgPoolOptions::new()
.max_connections(10)
.connect(url)
.await?;
Ok(Self { pool })
}
pub async fn connect_env(env_var: &str, default_url: &str) -> Result<Self, FlozError> {
let url = std::env::var(env_var).unwrap_or_else(|_| default_url.to_string());
Self::connect(&url).await
}
pub fn from_pool(pool: PgPool) -> Self {
Self { pool }
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
pub async fn begin(&self) -> Result<Tx, FlozError> {
let mut conn = self.pool.acquire().await?;
sqlx::query("BEGIN").execute(&mut *conn).await?;
Ok(Tx {
conn: Mutex::new(conn),
committed: false,
})
}
}
pub struct Tx {
conn: Mutex<sqlx::pool::PoolConnection<sqlx::Postgres>>,
committed: bool,
}
impl Tx {
pub async fn commit(mut self) -> Result<(), FlozError> {
let mut conn = self.conn.lock().await;
sqlx::query("COMMIT").execute(&mut **conn).await?;
self.committed = true;
Ok(())
}
pub async fn rollback(self) -> Result<(), FlozError> {
let mut conn = self.conn.lock().await;
sqlx::query("ROLLBACK").execute(&mut **conn).await?;
Ok(())
}
}
impl std::fmt::Debug for Tx {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Tx")
.field("committed", &self.committed)
.finish()
}
}
pub trait Executor {
fn execute_raw(
&self,
sql: &str,
params: Vec<Value>,
) -> impl std::future::Future<Output = Result<u64, FlozError>> + Send;
fn fetch_all<T>(
&self,
sql: &str,
params: Vec<Value>,
) -> impl std::future::Future<Output = Result<Vec<T>, FlozError>> + Send
where
T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>;
fn fetch_one<T>(
&self,
sql: &str,
params: Vec<Value>,
) -> impl std::future::Future<Output = Result<T, FlozError>> + Send
where
T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>;
fn fetch_optional<T>(
&self,
sql: &str,
params: Vec<Value>,
) -> impl std::future::Future<Output = Result<Option<T>, FlozError>> + Send
where
T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>;
}
macro_rules! bind_params {
($query:expr, $params:expr) => {{
let mut q = $query;
for param in $params {
q = match param {
Value::Short(v) => q.bind(v),
Value::Int(v) => q.bind(v),
Value::BigInt(v) => q.bind(v),
Value::Real(v) => q.bind(v),
Value::Double(v) => q.bind(v),
Value::Bool(v) => q.bind(v),
Value::String(v) => q.bind(v),
Value::Bytes(v) => q.bind(v),
Value::Uuid(v) => q.bind(v),
Value::DateTime(v) => q.bind(v),
Value::NaiveDateTime(v) => q.bind(v),
Value::NaiveDate(v) => q.bind(v),
Value::NaiveTime(v) => q.bind(v),
Value::Json(v) => q.bind(v),
Value::Jsonb(v) => q.bind(v),
Value::OptionShort(v) => q.bind(v),
Value::OptionInt(v) => q.bind(v),
Value::OptionBigInt(v) => q.bind(v),
Value::OptionReal(v) => q.bind(v),
Value::OptionDouble(v) => q.bind(v),
Value::OptionBool(v) => q.bind(v),
Value::OptionString(v) => q.bind(v),
Value::OptionBytes(v) => q.bind(v),
Value::OptionUuid(v) => q.bind(v),
Value::OptionDateTime(v) => q.bind(v),
Value::OptionNaiveDateTime(v) => q.bind(v),
Value::OptionNaiveDate(v) => q.bind(v),
Value::OptionNaiveTime(v) => q.bind(v),
Value::OptionJson(v) => q.bind(v),
Value::OptionJsonb(v) => q.bind(v),
};
}
q
}};
}
impl Executor for Db {
async fn execute_raw(&self, sql: &str, params: Vec<Value>) -> Result<u64, FlozError> {
let query = bind_params!(sqlx::query(sql), params);
let result = query.execute(&self.pool).await?;
Ok(result.rows_affected())
}
async fn fetch_all<T>(&self, sql: &str, params: Vec<Value>) -> Result<Vec<T>, FlozError>
where
T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>,
{
let query = bind_params!(sqlx::query_as::<_, T>(sql), params);
let rows = query.fetch_all(&self.pool).await?;
Ok(rows)
}
async fn fetch_one<T>(&self, sql: &str, params: Vec<Value>) -> Result<T, FlozError>
where
T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>,
{
let query = bind_params!(sqlx::query_as::<_, T>(sql), params);
let row = query
.fetch_optional(&self.pool)
.await?
.ok_or(FlozError::NotFound)?;
Ok(row)
}
async fn fetch_optional<T>(&self, sql: &str, params: Vec<Value>) -> Result<Option<T>, FlozError>
where
T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>,
{
let query = bind_params!(sqlx::query_as::<_, T>(sql), params);
let row = query.fetch_optional(&self.pool).await?;
Ok(row)
}
}
impl Executor for Tx {
async fn execute_raw(&self, sql: &str, params: Vec<Value>) -> Result<u64, FlozError> {
let mut conn = self.conn.lock().await;
let query = bind_params!(sqlx::query(sql), params);
let result = query.execute(&mut **conn).await?;
Ok(result.rows_affected())
}
async fn fetch_all<T>(&self, sql: &str, params: Vec<Value>) -> Result<Vec<T>, FlozError>
where
T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>,
{
let mut conn = self.conn.lock().await;
let query = bind_params!(sqlx::query_as::<_, T>(sql), params);
let rows = query.fetch_all(&mut **conn).await?;
Ok(rows)
}
async fn fetch_one<T>(&self, sql: &str, params: Vec<Value>) -> Result<T, FlozError>
where
T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>,
{
let mut conn = self.conn.lock().await;
let query = bind_params!(sqlx::query_as::<_, T>(sql), params);
let row = query
.fetch_optional(&mut **conn)
.await?
.ok_or(FlozError::NotFound)?;
Ok(row)
}
async fn fetch_optional<T>(&self, sql: &str, params: Vec<Value>) -> Result<Option<T>, FlozError>
where
T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>,
{
let mut conn = self.conn.lock().await;
let query = bind_params!(sqlx::query_as::<_, T>(sql), params);
let row = query.fetch_optional(&mut **conn).await?;
Ok(row)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn db_from_pool_type_check() {
fn _assert_send<T: Send>() {}
fn _assert_sync<T: Sync>() {}
_assert_send::<Db>();
_assert_sync::<Db>();
}
#[test]
fn tx_is_debug() {
fn _assert_debug<T: std::fmt::Debug>() {}
_assert_debug::<Tx>();
}
fn test_url() -> Option<String> {
std::env::var("DATABASE_URL").ok()
}
#[tokio::test]
async fn connect_to_postgres() {
let Some(url) = test_url() else { return };
let db = Db::connect(&url).await.unwrap();
assert!(db.pool().size() > 0 || true); }
#[tokio::test]
async fn execute_raw_select() {
let Some(url) = test_url() else { return };
let db = Db::connect(&url).await.unwrap();
let affected = db.execute_raw("SELECT 1", vec![]).await.unwrap();
assert_eq!(affected, 0);
}
#[tokio::test]
async fn execute_with_params() {
let Some(url) = test_url() else { return };
let db = Db::connect(&url).await.unwrap();
db.execute_raw("CREATE TEMP TABLE xql_test_params (id INT, name TEXT)", vec![])
.await
.unwrap();
let affected = db
.execute_raw(
"INSERT INTO xql_test_params (id, name) VALUES ($1, $2)",
vec![Value::Int(1), Value::String("Alice".into())],
)
.await
.unwrap();
assert_eq!(affected, 1);
}
#[tokio::test]
async fn fetch_all_rows() {
let Some(url) = test_url() else { return };
let db = Db::connect(&url).await.unwrap();
db.execute_raw("CREATE TEMP TABLE xql_test_fetch (id INT, name TEXT)", vec![])
.await
.unwrap();
db.execute_raw(
"INSERT INTO xql_test_fetch VALUES (1, 'Alice'), (2, 'Bob')",
vec![],
)
.await
.unwrap();
#[derive(Debug, sqlx::FromRow)]
struct Row {
id: i32,
name: String,
}
let rows: Vec<Row> = db
.fetch_all("SELECT id, name FROM xql_test_fetch ORDER BY id", vec![])
.await
.unwrap();
assert_eq!(rows.len(), 2);
assert_eq!(rows[0].name, "Alice");
assert_eq!(rows[1].name, "Bob");
}
#[tokio::test]
async fn fetch_one_found() {
let Some(url) = test_url() else { return };
let db = Db::connect(&url).await.unwrap();
db.execute_raw("CREATE TEMP TABLE xql_test_one (id INT, name TEXT)", vec![])
.await
.unwrap();
db.execute_raw("INSERT INTO xql_test_one VALUES (1, 'Alice')", vec![])
.await
.unwrap();
#[derive(Debug, sqlx::FromRow)]
struct Row {
id: i32,
name: String,
}
let row: Row = db
.fetch_one(
"SELECT id, name FROM xql_test_one WHERE id = $1",
vec![Value::Int(1)],
)
.await
.unwrap();
assert_eq!(row.name, "Alice");
}
#[tokio::test]
async fn fetch_one_not_found() {
let Some(url) = test_url() else { return };
let db = Db::connect(&url).await.unwrap();
db.execute_raw("CREATE TEMP TABLE xql_test_nf (id INT)", vec![])
.await
.unwrap();
#[derive(Debug, sqlx::FromRow)]
struct Row {
id: i32,
}
let result: Result<Row, _> = db
.fetch_one(
"SELECT id FROM xql_test_nf WHERE id = $1",
vec![Value::Int(999)],
)
.await;
assert!(matches!(result, Err(FlozError::NotFound)));
}
#[tokio::test]
async fn fetch_optional_some() {
let Some(url) = test_url() else { return };
let db = Db::connect(&url).await.unwrap();
db.execute_raw("CREATE TEMP TABLE xql_test_opt (id INT)", vec![])
.await
.unwrap();
db.execute_raw("INSERT INTO xql_test_opt VALUES (42)", vec![])
.await
.unwrap();
#[derive(Debug, sqlx::FromRow)]
struct Row {
id: i32,
}
let row: Option<Row> = db
.fetch_optional("SELECT id FROM xql_test_opt WHERE id = $1", vec![Value::Int(42)])
.await
.unwrap();
assert!(row.is_some());
assert_eq!(row.unwrap().id, 42);
}
#[tokio::test]
async fn fetch_optional_none() {
let Some(url) = test_url() else { return };
let db = Db::connect(&url).await.unwrap();
db.execute_raw("CREATE TEMP TABLE xql_test_opt_n (id INT)", vec![])
.await
.unwrap();
#[derive(Debug, sqlx::FromRow)]
struct Row {
id: i32,
}
let row: Option<Row> = db
.fetch_optional(
"SELECT id FROM xql_test_opt_n WHERE id = $1",
vec![Value::Int(999)],
)
.await
.unwrap();
assert!(row.is_none());
}
#[tokio::test]
async fn transaction_commit() {
let Some(url) = test_url() else { return };
let db = Db::connect(&url).await.unwrap();
db.execute_raw("CREATE TEMP TABLE xql_test_tx (id INT)", vec![])
.await
.unwrap();
let tx = db.begin().await.unwrap();
tx.execute_raw(
"INSERT INTO xql_test_tx VALUES ($1)",
vec![Value::Int(1)],
)
.await
.unwrap();
tx.commit().await.unwrap();
#[derive(Debug, sqlx::FromRow)]
struct Row { id: i32 }
let rows: Vec<Row> = db
.fetch_all("SELECT id FROM xql_test_tx", vec![])
.await
.unwrap();
assert_eq!(rows.len(), 1);
}
#[tokio::test]
async fn transaction_rollback_on_drop() {
let Some(url) = test_url() else { return };
let db = Db::connect(&url).await.unwrap();
db.execute_raw("CREATE TABLE IF NOT EXISTS xql_test_txd (id INT)", vec![])
.await
.unwrap();
db.execute_raw("DELETE FROM xql_test_txd", vec![])
.await
.unwrap();
{
let tx = db.begin().await.unwrap();
tx.execute_raw(
"INSERT INTO xql_test_txd VALUES ($1)",
vec![Value::Int(99)],
)
.await
.unwrap();
}
#[derive(Debug, sqlx::FromRow)]
struct Row { id: i32 }
let rows: Vec<Row> = db
.fetch_all("SELECT id FROM xql_test_txd WHERE id = 99", vec![])
.await
.unwrap();
assert_eq!(rows.len(), 0, "Row should not exist after rollback-on-drop");
db.execute_raw("DROP TABLE IF EXISTS xql_test_txd", vec![])
.await
.unwrap();
}
#[tokio::test]
async fn transaction_explicit_rollback() {
let Some(url) = test_url() else { return };
let db = Db::connect(&url).await.unwrap();
db.execute_raw("CREATE TEMP TABLE xql_test_rb (id INT)", vec![])
.await
.unwrap();
let tx = db.begin().await.unwrap();
tx.execute_raw(
"INSERT INTO xql_test_rb VALUES ($1)",
vec![Value::Int(1)],
)
.await
.unwrap();
tx.rollback().await.unwrap();
#[derive(Debug, sqlx::FromRow)]
struct Row { id: i32 }
let rows: Vec<Row> = db
.fetch_all("SELECT id FROM xql_test_rb", vec![])
.await
.unwrap();
assert_eq!(rows.len(), 0);
}
#[tokio::test]
async fn transaction_fetch() {
let Some(url) = test_url() else { return };
let db = Db::connect(&url).await.unwrap();
db.execute_raw("CREATE TEMP TABLE xql_test_txf (id INT, name TEXT)", vec![])
.await
.unwrap();
let tx = db.begin().await.unwrap();
tx.execute_raw(
"INSERT INTO xql_test_txf VALUES (1, 'Inside Tx')",
vec![],
)
.await
.unwrap();
#[derive(Debug, sqlx::FromRow)]
struct Row { id: i32, name: String }
let rows: Vec<Row> = tx
.fetch_all("SELECT id, name FROM xql_test_txf", vec![])
.await
.unwrap();
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].name, "Inside Tx");
tx.rollback().await.unwrap();
}
#[tokio::test]
async fn bind_typed_null() {
let Some(url) = test_url() else { return };
let db = Db::connect(&url).await.unwrap();
db.execute_raw(
"CREATE TEMP TABLE xql_test_null (id INT, name TEXT)",
vec![],
)
.await
.unwrap();
let affected = db
.execute_raw(
"INSERT INTO xql_test_null (id, name) VALUES ($1, $2)",
vec![Value::Int(1), Value::OptionString(None)],
)
.await
.unwrap();
assert_eq!(affected, 1);
#[derive(Debug, sqlx::FromRow)]
struct Row {
id: i32,
name: Option<String>,
}
let row: Row = db
.fetch_one("SELECT id, name FROM xql_test_null WHERE id = 1", vec![])
.await
.unwrap();
assert_eq!(row.id, 1);
assert!(row.name.is_none());
}
}