rust_ef_postgres/
provider.rs1use 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 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 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 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}