Skip to main content

pgorm_check/
client.rs

1//! Database client trait for pgorm-check
2//!
3//! This module defines a minimal trait for database operations,
4//! allowing pgorm-check to work independently of pgorm.
5
6use crate::error::{CheckError, CheckResult};
7use tokio_postgres::Row;
8use tokio_postgres::types::ToSql;
9
10/// A trait for types that can execute PostgreSQL queries.
11///
12/// This is implemented for `tokio_postgres::Client`, `tokio_postgres::Transaction`,
13/// and connection pool clients.
14#[async_trait::async_trait]
15pub trait CheckClient: Sync {
16    /// Execute a query and return all rows.
17    async fn query(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) -> CheckResult<Vec<Row>>;
18
19    /// Execute a query and return exactly one row.
20    async fn query_one(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) -> CheckResult<Row>;
21
22    /// Execute a statement and return the number of affected rows.
23    async fn execute(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) -> CheckResult<u64>;
24}
25
26#[async_trait::async_trait]
27impl CheckClient for tokio_postgres::Client {
28    async fn query(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) -> CheckResult<Vec<Row>> {
29        self.query(sql, params).await.map_err(CheckError::from)
30    }
31
32    async fn query_one(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) -> CheckResult<Row> {
33        self.query_one(sql, params).await.map_err(CheckError::from)
34    }
35
36    async fn execute(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) -> CheckResult<u64> {
37        self.execute(sql, params).await.map_err(CheckError::from)
38    }
39}
40
41#[async_trait::async_trait]
42impl<'a> CheckClient for tokio_postgres::Transaction<'a> {
43    async fn query(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) -> CheckResult<Vec<Row>> {
44        self.query(sql, params).await.map_err(CheckError::from)
45    }
46
47    async fn query_one(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) -> CheckResult<Row> {
48        self.query_one(sql, params).await.map_err(CheckError::from)
49    }
50
51    async fn execute(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) -> CheckResult<u64> {
52        self.execute(sql, params).await.map_err(CheckError::from)
53    }
54}
55
56/// Extension trait for accessing row columns with better error handling.
57pub trait RowExt {
58    /// Get a column value by name, returning a CheckError on failure.
59    fn try_get_column<'a, T>(&'a self, column: &str) -> CheckResult<T>
60    where
61        T: tokio_postgres::types::FromSql<'a>;
62}
63
64impl RowExt for Row {
65    fn try_get_column<'a, T>(&'a self, column: &str) -> CheckResult<T>
66    where
67        T: tokio_postgres::types::FromSql<'a>,
68    {
69        self.try_get(column)
70            .map_err(|e| CheckError::decode(column, e.to_string()))
71    }
72}