systemprompt_database/services/postgres/
mod.rs

1//! `PostgreSQL` implementation of [`crate::services::DatabaseProvider`].
2//!
3//! This module is part of the documented sqlx allowlist: every `sqlx::query(_)`
4//! call here either binds a [`crate::models::QuerySelector`] string supplied
5//! at runtime (extension-defined SQL, dynamic admin queries) or executes
6//! `SELECT 1` for connection probing. Static SQL goes through the verified
7//! macros elsewhere.
8//!
9//! Copyright (c) systemprompt.io — Business Source License 1.1.
10//! See <https://systemprompt.io> for licensing details.
11
12pub mod connection;
13pub mod conversion;
14mod ext;
15mod introspection;
16pub mod transaction;
17
18use async_trait::async_trait;
19use sqlx::Executor;
20use sqlx::postgres::{PgConnectOptions, PgPool, PgSslMode};
21use std::str::FromStr;
22use std::sync::Arc;
23
24use super::provider::DatabaseProvider;
25use crate::error::{DatabaseResult, RepositoryError};
26use crate::models::{
27    DatabaseInfo, DatabaseTransaction, DbValue, JsonRow, QueryResult, QuerySelector, ToDbValue,
28};
29use conversion::{bind_params, row_to_json, rows_to_result};
30use transaction::PostgresTransaction;
31
32#[derive(Debug)]
33pub struct PostgresProvider {
34    pool: Arc<PgPool>,
35}
36
37impl PostgresProvider {
38    pub async fn new(database_url: &str) -> DatabaseResult<Self> {
39        Self::new_with_pool(database_url, &connection::PoolConfig::default()).await
40    }
41
42    pub async fn new_with_pool(
43        database_url: &str,
44        pool_config: &connection::PoolConfig,
45    ) -> DatabaseResult<Self> {
46        let mut connect_options = PgConnectOptions::from_str(database_url)?;
47
48        let ssl_mode = if database_url.contains("sslmode=require") {
49            PgSslMode::Require
50        } else if database_url.contains("sslmode=disable") {
51            PgSslMode::Disable
52        } else {
53            PgSslMode::Prefer
54        };
55
56        connect_options = connect_options
57            .application_name("systemprompt")
58            .statement_cache_capacity(0)
59            .ssl_mode(ssl_mode)
60            .options([("client_min_messages", "warning")]);
61
62        if let Some(ca_cert_path) = Self::get_cert_path() {
63            connect_options = connect_options.ssl_root_cert(&ca_cert_path);
64        }
65
66        let pool = connection::connect_with_retry(
67            connection::build_pool_options(pool_config),
68            connect_options,
69        )
70        .await?;
71
72        Ok(Self {
73            pool: Arc::new(pool),
74        })
75    }
76
77    #[must_use]
78    pub const fn from_pool(pool: Arc<PgPool>) -> Self {
79        Self { pool }
80    }
81
82    fn get_cert_path() -> Option<std::path::PathBuf> {
83        std::env::var("PGCA_CERT_PATH")
84            .ok()
85            .map(std::path::PathBuf::from)
86    }
87
88    #[must_use]
89    pub fn pool(&self) -> &PgPool {
90        &self.pool
91    }
92}
93
94#[async_trait]
95impl DatabaseProvider for PostgresProvider {
96    fn get_postgres_pool(&self) -> Option<Arc<PgPool>> {
97        Some(Arc::clone(&self.pool))
98    }
99
100    async fn execute(
101        &self,
102        query: &dyn QuerySelector,
103        params: &[&dyn ToDbValue],
104    ) -> DatabaseResult<u64> {
105        let sql = query.select_query();
106        let query_obj = sqlx::query(sqlx::AssertSqlSafe(sql));
107        let query_obj = bind_params(query_obj, params);
108
109        let result = query_obj.execute(&*self.pool).await?;
110
111        Ok(result.rows_affected())
112    }
113
114    async fn execute_raw(&self, sql: &str) -> DatabaseResult<()> {
115        let mut conn = self.pool.acquire().await?;
116
117        conn.execute(sqlx::AssertSqlSafe(sql.to_owned())).await?;
118
119        Ok(())
120    }
121
122    async fn fetch_all(
123        &self,
124        query: &dyn QuerySelector,
125        params: &[&dyn ToDbValue],
126    ) -> DatabaseResult<Vec<JsonRow>> {
127        let sql = query.select_query();
128        let query_obj = sqlx::query(sqlx::AssertSqlSafe(sql));
129        let query_obj = bind_params(query_obj, params);
130
131        let rows = query_obj.fetch_all(&*self.pool).await?;
132
133        Ok(rows.iter().map(row_to_json).collect())
134    }
135
136    async fn fetch_one(
137        &self,
138        query: &dyn QuerySelector,
139        params: &[&dyn ToDbValue],
140    ) -> DatabaseResult<JsonRow> {
141        let sql = query.select_query();
142        let query_obj = sqlx::query(sqlx::AssertSqlSafe(sql));
143        let query_obj = bind_params(query_obj, params);
144
145        let row = query_obj.fetch_one(&*self.pool).await?;
146
147        Ok(row_to_json(&row))
148    }
149
150    async fn fetch_optional(
151        &self,
152        query: &dyn QuerySelector,
153        params: &[&dyn ToDbValue],
154    ) -> DatabaseResult<Option<JsonRow>> {
155        let sql = query.select_query();
156        let query_obj = sqlx::query(sqlx::AssertSqlSafe(sql));
157        let query_obj = bind_params(query_obj, params);
158
159        let row = query_obj.fetch_optional(&*self.pool).await?;
160
161        Ok(row.map(|r| row_to_json(&r)))
162    }
163
164    async fn fetch_scalar_value(
165        &self,
166        query: &dyn QuerySelector,
167        params: &[&dyn ToDbValue],
168    ) -> DatabaseResult<DbValue> {
169        let row = self.fetch_one(query, params).await?;
170
171        let first_value = row
172            .values()
173            .next()
174            .ok_or_else(|| RepositoryError::invalid_state("No columns in result"))?;
175
176        let db_value = match first_value {
177            serde_json::Value::String(s) => DbValue::String(s.clone()),
178            serde_json::Value::Number(n) => n
179                .as_i64()
180                .map(DbValue::Int)
181                .or_else(|| n.as_f64().map(DbValue::Float))
182                .unwrap_or(DbValue::NullFloat),
183            serde_json::Value::Bool(b) => DbValue::Bool(*b),
184            serde_json::Value::Null => DbValue::NullString,
185            serde_json::Value::Array(_) | serde_json::Value::Object(_) => {
186                return Err(RepositoryError::invalid_state("Unsupported value type"));
187            },
188        };
189
190        Ok(db_value)
191    }
192
193    async fn begin_transaction(&self) -> DatabaseResult<Box<dyn DatabaseTransaction>> {
194        let tx = self.pool.begin().await?;
195
196        Ok(Box::new(PostgresTransaction::new(tx)))
197    }
198
199    async fn get_database_info(&self) -> DatabaseResult<DatabaseInfo> {
200        introspection::get_database_info(&self.pool).await
201    }
202
203    async fn test_connection(&self) -> DatabaseResult<()> {
204        sqlx::query("SELECT 1").fetch_one(&*self.pool).await?;
205        Ok(())
206    }
207
208    async fn execute_batch(&self, sql: &str) -> DatabaseResult<()> {
209        let statements = crate::services::SqlExecutor::parse_sql_statements(sql)?;
210        for statement in statements {
211            sqlx::query(sqlx::AssertSqlSafe(statement))
212                .execute(&*self.pool)
213                .await?;
214        }
215        Ok(())
216    }
217
218    async fn query_raw(&self, query: &dyn QuerySelector) -> DatabaseResult<QueryResult> {
219        let sql = query.select_query();
220        let start = std::time::Instant::now();
221
222        let rows = sqlx::query(sqlx::AssertSqlSafe(sql))
223            .fetch_all(&*self.pool)
224            .await?;
225
226        Ok(rows_to_result(rows, start))
227    }
228
229    async fn query_raw_with(
230        &self,
231        query: &dyn QuerySelector,
232        params: &[&dyn ToDbValue],
233    ) -> DatabaseResult<QueryResult> {
234        let sql = query.select_query();
235        let start = std::time::Instant::now();
236
237        let query_obj = bind_params(sqlx::query(sqlx::AssertSqlSafe(sql)), params);
238        let rows = query_obj.fetch_all(&*self.pool).await?;
239
240        Ok(rows_to_result(rows, start))
241    }
242}
systemprompt_database/services/postgres/mod.rs

systemprompt_database/services/postgres/
mod.rs
