Skip to main content

rust_ef_postgres/
provider.rs

1use crate::sql_generator::PostgresSqlGenerator;
2use crate::tls::PgTlsMode;
3use async_trait::async_trait;
4use deadpool_postgres::{Config, Pool, Runtime};
5use rust_ef::error::{EFError, EFResult};
6#[cfg(feature = "tracing")]
7use rust_ef::provider::IAsyncConnection;
8use rust_ef::provider::{IDatabaseProvider, ISqlGenerator};
9use tokio_postgres::NoTls;
10
11pub struct PostgresProvider {
12    pool: Pool,
13    #[cfg(feature = "tracing")]
14    slow_query_threshold_ms: std::sync::atomic::AtomicU64,
15}
16
17impl PostgresProvider {
18    /// Creates a provider with TLS required (secure-by-default, v1.6+).
19    ///
20    /// Uses a `native_tls::TlsConnector` built from the platform's default
21    /// root certificate store. For plaintext connections (local dev only),
22    /// use [`PostgresProvider::new_insecure`]. For custom CA certificates,
23    /// use [`PostgresProvider::new_with_tls`].
24    pub fn new(connection_string: &str, pool_size: usize) -> EFResult<Self> {
25        let connector = native_tls::TlsConnector::builder()
26            .build()
27            .map_err(|e| EFError::connection(format!("TLS connector init failed: {}", e)))?;
28        Self::new_with_tls(connection_string, pool_size, PgTlsMode::Require(connector))
29    }
30
31    /// Creates a provider with TLS disabled (plaintext, local dev only).
32    pub fn new_insecure(connection_string: &str, pool_size: usize) -> EFResult<Self> {
33        Self::new_with_tls(connection_string, pool_size, PgTlsMode::Disable)
34    }
35
36    /// Creates a provider with configurable TLS.
37    ///
38    /// `PgTlsMode::Disable` is equivalent to [`PostgresProvider::new_insecure`].
39    /// `PgTlsMode::Require(connector)` enforces TLS for all pooled connections.
40    ///
41    /// The TLS connector type is erased inside `deadpool_postgres::Manager`
42    /// (via `Box<dyn Connect>`), so [`Pool`] and [`PostgresConnection`] are
43    /// non-generic — TLS configuration is a per-pool construction-time
44    /// decision, not a type-level parameter.
45    pub fn new_with_tls(
46        connection_string: &str,
47        pool_size: usize,
48        tls: PgTlsMode,
49    ) -> EFResult<Self> {
50        let config: tokio_postgres::Config = connection_string
51            .parse()
52            .map_err(|e| EFError::connection(format!("Invalid connection string: {}", e)))?;
53        let mut cfg = Config::new();
54        if let Some(tokio_postgres::config::Host::Tcp(h)) = config.get_hosts().first() {
55            cfg.host = Some(h.clone());
56        }
57        cfg.port = Some(config.get_ports().first().copied().unwrap_or(5432));
58        cfg.dbname = Some(config.get_dbname().unwrap_or("postgres").to_string());
59        cfg.user = Some(config.get_user().unwrap_or("postgres").to_string());
60        cfg.password = config
61            .get_password()
62            .map(|p| String::from_utf8_lossy(p).to_string());
63        if pool_size > 0 {
64            cfg.pool.get_or_insert_default().max_size = pool_size;
65        }
66
67        let pool = match tls {
68            PgTlsMode::Disable => cfg
69                .create_pool(Some(Runtime::Tokio1), NoTls)
70                .map_err(|e| EFError::connection(format!("Failed to create pool: {}", e)))?,
71            PgTlsMode::Require(connector) => {
72                let tls = postgres_native_tls::MakeTlsConnector::new(connector);
73                cfg.create_pool(Some(Runtime::Tokio1), tls)
74                    .map_err(|e| EFError::connection(format!("Failed to create pool: {}", e)))?
75            }
76        };
77        Ok(Self {
78            pool,
79            #[cfg(feature = "tracing")]
80            slow_query_threshold_ms: std::sync::atomic::AtomicU64::new(0),
81        })
82    }
83}
84
85#[async_trait]
86impl IDatabaseProvider for PostgresProvider {
87    fn sql_generator(&self) -> &'static dyn ISqlGenerator {
88        &PostgresSqlGenerator
89    }
90
91    async fn get_connection(&self) -> EFResult<Box<dyn rust_ef::provider::IAsyncConnection>> {
92        let _guard = rust_ef::observability::PoolAcquireGuard::new("PostgreSQL");
93        let client = self
94            .pool
95            .get()
96            .await
97            .map_err(|e| EFError::connection(format!("Pool error: {}", e)))?;
98        #[cfg_attr(not(feature = "tracing"), allow(unused_mut))]
99        let mut conn = Box::new(crate::connection::PostgresConnection::new(client));
100        #[cfg(feature = "tracing")]
101        {
102            let ms = self
103                .slow_query_threshold_ms
104                .load(std::sync::atomic::Ordering::Relaxed);
105            if ms > 0 {
106                conn.set_slow_query_threshold(std::time::Duration::from_millis(ms));
107            }
108        }
109        Ok(conn)
110    }
111
112    async fn execute_migration_command(&self, sql: &str) -> EFResult<()> {
113        let client = self
114            .pool
115            .get()
116            .await
117            .map_err(|e| EFError::connection(format!("Pool error: {}", e)))?;
118        client
119            .batch_execute(sql)
120            .await
121            .map_err(|e| EFError::migration(format!("Migration execution failed: {}", e)))?;
122        Ok(())
123    }
124
125    fn name(&self) -> &str {
126        "PostgreSQL"
127    }
128
129    fn migration_dialect(&self) -> rust_ef::migration::MigrationDialect {
130        rust_ef::migration::MigrationDialect::Postgres
131    }
132
133    #[cfg(feature = "tracing")]
134    fn set_slow_query_threshold(&self, threshold: std::time::Duration) {
135        self.slow_query_threshold_ms.store(
136            threshold.as_millis() as u64,
137            std::sync::atomic::Ordering::Relaxed,
138        );
139    }
140}
141
142impl std::fmt::Debug for PostgresProvider {
143    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
144        f.debug_struct("PostgresProvider")
145            .field("name", &self.name())
146            .finish()
147    }
148}