Skip to main content

ferriorm_runtime/
client.rs

1//! Database connection pool wrapper.
2//!
3//! [`DatabaseClient`] is an enum that wraps either a PostgreSQL or SQLite
4//! connection pool (via sqlx). It provides auto-detection from the connection
5//! URL and exposes typed `fetch_all`, `fetch_optional`, `fetch_one`, and
6//! `execute` helpers used by the generated query builders.
7
8use crate::error::FerriormError;
9
10/// The database client, wrapping an sqlx connection pool.
11///
12/// Supports PostgreSQL and SQLite via feature flags.
13#[derive(Debug, Clone)]
14pub enum DatabaseClient {
15    #[cfg(feature = "postgres")]
16    Postgres(sqlx::PgPool),
17    #[cfg(feature = "sqlite")]
18    Sqlite(sqlx::SqlitePool),
19}
20
21impl DatabaseClient {
22    /// Connect to a PostgreSQL database.
23    #[cfg(feature = "postgres")]
24    pub async fn connect_postgres(url: &str) -> Result<Self, FerriormError> {
25        let pool = sqlx::PgPool::connect(url).await?;
26        Ok(Self::Postgres(pool))
27    }
28
29    /// Connect to a SQLite database.
30    #[cfg(feature = "sqlite")]
31    pub async fn connect_sqlite(url: &str) -> Result<Self, FerriormError> {
32        let pool = sqlx::SqlitePool::connect(url).await?;
33        Ok(Self::Sqlite(pool))
34    }
35
36    /// Connect by auto-detecting the database type from the URL.
37    pub async fn connect(url: &str) -> Result<Self, FerriormError> {
38        #[cfg(feature = "sqlite")]
39        if url.starts_with("sqlite:") || url.starts_with("file:") || url.ends_with(".db") {
40            return Self::connect_sqlite(url).await;
41        }
42
43        #[cfg(feature = "postgres")]
44        {
45            return Self::connect_postgres(url).await;
46        }
47
48        #[allow(unreachable_code)]
49        Err(FerriormError::Connection(
50            "No database backend enabled. Enable 'postgres' or 'sqlite' feature.".into(),
51        ))
52    }
53
54    /// Close the connection pool.
55    pub async fn disconnect(self) {
56        match self {
57            #[cfg(feature = "postgres")]
58            Self::Postgres(pool) => pool.close().await,
59            #[cfg(feature = "sqlite")]
60            Self::Sqlite(pool) => pool.close().await,
61        }
62    }
63
64    /// Execute a query builder against the appropriate pool, returning all rows.
65    #[cfg(feature = "postgres")]
66    pub async fn fetch_all_pg<'q, T>(
67        &self,
68        mut qb: sqlx::QueryBuilder<'q, sqlx::Postgres>,
69    ) -> Result<Vec<T>, FerriormError>
70    where
71        T: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
72    {
73        match self {
74            Self::Postgres(pool) => Ok(qb.build_query_as::<T>().fetch_all(pool).await?),
75            #[cfg(feature = "sqlite")]
76            _ => Err(FerriormError::Query(
77                "Expected PostgreSQL connection".into(),
78            )),
79        }
80    }
81
82    #[cfg(feature = "postgres")]
83    pub async fn fetch_optional_pg<'q, T>(
84        &self,
85        mut qb: sqlx::QueryBuilder<'q, sqlx::Postgres>,
86    ) -> Result<Option<T>, FerriormError>
87    where
88        T: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
89    {
90        match self {
91            Self::Postgres(pool) => Ok(qb.build_query_as::<T>().fetch_optional(pool).await?),
92            #[cfg(feature = "sqlite")]
93            _ => Err(FerriormError::Query(
94                "Expected PostgreSQL connection".into(),
95            )),
96        }
97    }
98
99    #[cfg(feature = "postgres")]
100    pub async fn fetch_one_pg<'q, T>(
101        &self,
102        mut qb: sqlx::QueryBuilder<'q, sqlx::Postgres>,
103    ) -> Result<T, FerriormError>
104    where
105        T: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin,
106    {
107        match self {
108            Self::Postgres(pool) => Ok(qb.build_query_as::<T>().fetch_one(pool).await?),
109            #[cfg(feature = "sqlite")]
110            _ => Err(FerriormError::Query(
111                "Expected PostgreSQL connection".into(),
112            )),
113        }
114    }
115
116    #[cfg(feature = "postgres")]
117    pub async fn execute_pg<'q>(
118        &self,
119        mut qb: sqlx::QueryBuilder<'q, sqlx::Postgres>,
120    ) -> Result<u64, FerriormError> {
121        match self {
122            Self::Postgres(pool) => Ok(qb.build().execute(pool).await?.rows_affected()),
123            #[cfg(feature = "sqlite")]
124            _ => Err(FerriormError::Query(
125                "Expected PostgreSQL connection".into(),
126            )),
127        }
128    }
129
130    // SQLite variants
131    #[cfg(feature = "sqlite")]
132    pub async fn fetch_all_sqlite<'q, T>(
133        &self,
134        mut qb: sqlx::QueryBuilder<'q, sqlx::Sqlite>,
135    ) -> Result<Vec<T>, FerriormError>
136    where
137        T: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow> + Send + Unpin,
138    {
139        match self {
140            Self::Sqlite(pool) => Ok(qb.build_query_as::<T>().fetch_all(pool).await?),
141            #[cfg(feature = "postgres")]
142            _ => Err(FerriormError::Query("Expected SQLite connection".into())),
143        }
144    }
145
146    #[cfg(feature = "sqlite")]
147    pub async fn fetch_optional_sqlite<'q, T>(
148        &self,
149        mut qb: sqlx::QueryBuilder<'q, sqlx::Sqlite>,
150    ) -> Result<Option<T>, FerriormError>
151    where
152        T: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow> + Send + Unpin,
153    {
154        match self {
155            Self::Sqlite(pool) => Ok(qb.build_query_as::<T>().fetch_optional(pool).await?),
156            #[cfg(feature = "postgres")]
157            _ => Err(FerriormError::Query("Expected SQLite connection".into())),
158        }
159    }
160
161    #[cfg(feature = "sqlite")]
162    pub async fn fetch_one_sqlite<'q, T>(
163        &self,
164        mut qb: sqlx::QueryBuilder<'q, sqlx::Sqlite>,
165    ) -> Result<T, FerriormError>
166    where
167        T: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow> + Send + Unpin,
168    {
169        match self {
170            Self::Sqlite(pool) => Ok(qb.build_query_as::<T>().fetch_one(pool).await?),
171            #[cfg(feature = "postgres")]
172            _ => Err(FerriormError::Query("Expected SQLite connection".into())),
173        }
174    }
175
176    #[cfg(feature = "sqlite")]
177    pub async fn execute_sqlite<'q>(
178        &self,
179        mut qb: sqlx::QueryBuilder<'q, sqlx::Sqlite>,
180    ) -> Result<u64, FerriormError> {
181        match self {
182            Self::Sqlite(pool) => Ok(qb.build().execute(pool).await?.rows_affected()),
183            #[cfg(feature = "postgres")]
184            _ => Err(FerriormError::Query("Expected SQLite connection".into())),
185        }
186    }
187}