Skip to main content

rust_ef_postgres/
connection.rs

1use async_trait::async_trait;
2use rust_ef::error::{EFError, EFResult};
3use rust_ef::provider::{DbValue, IAsyncConnection, IsolationLevel};
4use tokio_postgres::types::ToSql;
5
6pub struct PostgresConnection {
7    pub(crate) client: deadpool_postgres::Client,
8    #[cfg(feature = "tracing")]
9    slow_query_threshold: Option<std::time::Duration>,
10}
11
12impl PostgresConnection {
13    pub(crate) fn new(client: deadpool_postgres::Client) -> Self {
14        Self {
15            client,
16            #[cfg(feature = "tracing")]
17            slow_query_threshold: None,
18        }
19    }
20
21    fn threshold(&self) -> Option<std::time::Duration> {
22        #[cfg(feature = "tracing")]
23        {
24            self.slow_query_threshold
25        }
26        #[cfg(not(feature = "tracing"))]
27        {
28            None
29        }
30    }
31}
32
33#[async_trait]
34impl IAsyncConnection for PostgresConnection {
35    async fn execute(&mut self, sql: &str, params: &[DbValue]) -> EFResult<u64> {
36        let _guard = rust_ef::observability::QueryGuard::new(sql, self.threshold());
37        let pgp = crate::type_conversion::db_values_to_pg_params(params);
38        let refs: Vec<&(dyn ToSql + Sync)> = pgp
39            .iter()
40            .map(|p| p.as_ref() as &(dyn ToSql + Sync))
41            .collect();
42        self.client
43            .execute(sql, &refs)
44            .await
45            .map_err(|e| EFError::query(format!("Execution error: {}", e)))
46    }
47
48    async fn query(&mut self, sql: &str, params: &[DbValue]) -> EFResult<Vec<Vec<DbValue>>> {
49        let _guard = rust_ef::observability::QueryGuard::new(sql, self.threshold());
50        let pgp = crate::type_conversion::db_values_to_pg_params(params);
51        let refs: Vec<&(dyn ToSql + Sync)> = pgp
52            .iter()
53            .map(|p| p.as_ref() as &(dyn ToSql + Sync))
54            .collect();
55        let rows = self
56            .client
57            .query(sql, &refs)
58            .await
59            .map_err(|e| EFError::query(format!("Query error: {}", e)))?;
60        let columns: Vec<&tokio_postgres::Column> = if !rows.is_empty() {
61            rows[0].columns().iter().collect()
62        } else {
63            Vec::new()
64        };
65        let result = rows
66            .iter()
67            .map(|row| {
68                columns
69                    .iter()
70                    .enumerate()
71                    .map(|(i, col)| crate::row_conversion::cell_to_db_value(row, i, col.type_()))
72                    .collect()
73            })
74            .collect();
75        Ok(result)
76    }
77
78    async fn begin_transaction(&mut self) -> EFResult<()> {
79        self.client
80            .simple_query("BEGIN")
81            .await
82            .map_err(|e| EFError::transaction(format!("BEGIN failed: {}", e)))?;
83        Ok(())
84    }
85
86    async fn commit_transaction(&mut self) -> EFResult<()> {
87        self.client
88            .simple_query("COMMIT")
89            .await
90            .map_err(|e| EFError::transaction(format!("COMMIT failed: {}", e)))?;
91        Ok(())
92    }
93
94    async fn rollback_transaction(&mut self) -> EFResult<()> {
95        self.client
96            .simple_query("ROLLBACK")
97            .await
98            .map_err(|e| EFError::transaction(format!("ROLLBACK failed: {}", e)))?;
99        Ok(())
100    }
101
102    async fn create_savepoint(&mut self, name: &str) -> EFResult<()> {
103        self.client
104            .simple_query(&format!("SAVEPOINT {}", name))
105            .await
106            .map_err(|e| EFError::transaction(format!("SAVEPOINT failed: {}", e)))?;
107        Ok(())
108    }
109
110    async fn release_savepoint(&mut self, name: &str) -> EFResult<()> {
111        self.client
112            .simple_query(&format!("RELEASE SAVEPOINT {}", name))
113            .await
114            .map_err(|e| EFError::transaction(format!("RELEASE failed: {}", e)))?;
115        Ok(())
116    }
117
118    async fn rollback_to_savepoint(&mut self, name: &str) -> EFResult<()> {
119        self.client
120            .simple_query(&format!("ROLLBACK TO SAVEPOINT {}", name))
121            .await
122            .map_err(|e| EFError::transaction(format!("ROLLBACK TO failed: {}", e)))?;
123        Ok(())
124    }
125
126    async fn set_transaction_isolation(&mut self, level: IsolationLevel) -> EFResult<()> {
127        let sql = format!(
128            "SET TRANSACTION ISOLATION LEVEL {}",
129            match level {
130                IsolationLevel::ReadUncommitted => "READ UNCOMMITTED",
131                IsolationLevel::ReadCommitted => "READ COMMITTED",
132                IsolationLevel::RepeatableRead => "REPEATABLE READ",
133                IsolationLevel::Serializable => "SERIALIZABLE",
134            }
135        );
136        self.client
137            .simple_query(&sql)
138            .await
139            .map_err(|e| EFError::transaction(format!("SET ISOLATION failed: {}", e)))?;
140        Ok(())
141    }
142
143    #[cfg(feature = "tracing")]
144    fn set_slow_query_threshold(&mut self, threshold: std::time::Duration) {
145        self.slow_query_threshold = Some(threshold);
146    }
147}