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 introspection;
15pub mod transaction;
16
17use async_trait::async_trait;
18use sqlx::Executor;
19use sqlx::postgres::PgPool;
20use std::sync::Arc;
21
22use super::provider::DatabaseProvider;
23use crate::error::DatabaseResult;
24use crate::models::{
25    DatabaseInfo, DatabaseTransaction, JsonRow, QueryResult, QuerySelector, ToDbValue,
26};
27use conversion::{bind_params, row_to_json, rows_to_result};
28use transaction::PostgresTransaction;
29
30#[derive(Debug)]
31pub struct PostgresProvider {
32    pool: Arc<PgPool>,
33}
34
35impl PostgresProvider {
36    pub async fn new(database_url: &str) -> DatabaseResult<Self> {
37        Self::new_with_pool(
38            database_url,
39            &connection::PoolConfig::default(),
40            &connection::SchemaEpoch::new(),
41        )
42        .await
43    }
44
45    pub async fn new_with_pool(
46        database_url: &str,
47        pool_config: &connection::PoolConfig,
48        epoch: &connection::SchemaEpoch,
49    ) -> DatabaseResult<Self> {
50        let connect_options = connection::connect_options(database_url)?
51            .statement_cache_capacity(pool_config.statement_cache_capacity);
52
53        let pool = connection::connect_with_retry(
54            connection::build_pool_options(pool_config, epoch),
55            connect_options,
56        )
57        .await?;
58
59        Ok(Self {
60            pool: Arc::new(pool),
61        })
62    }
63
64    #[must_use]
65    pub const fn from_pool(pool: Arc<PgPool>) -> Self {
66        Self { pool }
67    }
68
69    #[must_use]
70    pub fn pool(&self) -> &PgPool {
71        &self.pool
72    }
73}
74
75#[async_trait]
76impl DatabaseProvider for PostgresProvider {
77    fn get_postgres_pool(&self) -> Arc<PgPool> {
78        Arc::clone(&self.pool)
79    }
80
81    async fn execute(
82        &self,
83        query: &dyn QuerySelector,
84        params: &[&dyn ToDbValue],
85    ) -> DatabaseResult<u64> {
86        let sql = query.select_query();
87        let query_obj = sqlx::query(sqlx::AssertSqlSafe(sql));
88        let query_obj = bind_params(query_obj, params);
89
90        let result = query_obj.execute(&*self.pool).await?;
91
92        Ok(result.rows_affected())
93    }
94
95    async fn execute_raw(&self, sql: &str) -> DatabaseResult<()> {
96        let mut conn = self.pool.acquire().await?;
97
98        conn.execute(sqlx::AssertSqlSafe(sql.to_owned())).await?;
99
100        Ok(())
101    }
102
103    async fn fetch_all(
104        &self,
105        query: &dyn QuerySelector,
106        params: &[&dyn ToDbValue],
107    ) -> DatabaseResult<Vec<JsonRow>> {
108        let sql = query.select_query();
109        let query_obj = sqlx::query(sqlx::AssertSqlSafe(sql));
110        let query_obj = bind_params(query_obj, params);
111
112        let rows = query_obj.fetch_all(&*self.pool).await?;
113
114        Ok(rows.iter().map(row_to_json).collect())
115    }
116
117    async fn fetch_one(
118        &self,
119        query: &dyn QuerySelector,
120        params: &[&dyn ToDbValue],
121    ) -> DatabaseResult<JsonRow> {
122        let sql = query.select_query();
123        let query_obj = sqlx::query(sqlx::AssertSqlSafe(sql));
124        let query_obj = bind_params(query_obj, params);
125
126        let row = query_obj.fetch_one(&*self.pool).await?;
127
128        Ok(row_to_json(&row))
129    }
130
131    async fn fetch_optional(
132        &self,
133        query: &dyn QuerySelector,
134        params: &[&dyn ToDbValue],
135    ) -> DatabaseResult<Option<JsonRow>> {
136        let sql = query.select_query();
137        let query_obj = sqlx::query(sqlx::AssertSqlSafe(sql));
138        let query_obj = bind_params(query_obj, params);
139
140        let row = query_obj.fetch_optional(&*self.pool).await?;
141
142        Ok(row.map(|r| row_to_json(&r)))
143    }
144
145    async fn begin_transaction(&self) -> DatabaseResult<Box<dyn DatabaseTransaction>> {
146        let tx = self.pool.begin().await?;
147
148        Ok(Box::new(PostgresTransaction::new(tx)))
149    }
150
151    async fn get_database_info(&self) -> DatabaseResult<DatabaseInfo> {
152        introspection::get_database_info(&self.pool).await
153    }
154
155    async fn test_connection(&self) -> DatabaseResult<()> {
156        sqlx::query("SELECT 1").fetch_one(&*self.pool).await?;
157        Ok(())
158    }
159
160    async fn execute_batch(&self, sql: &str) -> DatabaseResult<()> {
161        let statements = crate::services::SqlExecutor::parse_sql_statements(sql)?;
162        for statement in statements {
163            sqlx::query(sqlx::AssertSqlSafe(statement))
164                .execute(&*self.pool)
165                .await?;
166        }
167        Ok(())
168    }
169
170    async fn query_raw(&self, query: &dyn QuerySelector) -> DatabaseResult<QueryResult> {
171        let sql = query.select_query();
172        let start = std::time::Instant::now();
173
174        let rows = sqlx::query(sqlx::AssertSqlSafe(sql))
175            .fetch_all(&*self.pool)
176            .await?;
177
178        Ok(rows_to_result(rows, start))
179    }
180
181    async fn query_raw_with(
182        &self,
183        query: &dyn QuerySelector,
184        params: &[&dyn ToDbValue],
185    ) -> DatabaseResult<QueryResult> {
186        let sql = query.select_query();
187        let start = std::time::Instant::now();
188
189        let query_obj = bind_params(sqlx::query(sqlx::AssertSqlSafe(sql)), params);
190        let rows = query_obj.fetch_all(&*self.pool).await?;
191
192        Ok(rows_to_result(rows, start))
193    }
194}
systemprompt_database/services/postgres/mod.rs

systemprompt_database/services/postgres/
mod.rs
