deadpool_postgres/
lib.rs

1#![doc = include_str!("../README.md")]
2#![cfg_attr(docsrs, feature(doc_cfg))]
3#![deny(
4    nonstandard_style,
5    rust_2018_idioms,
6    rustdoc::broken_intra_doc_links,
7    rustdoc::private_intra_doc_links
8)]
9#![forbid(non_ascii_idents, unsafe_code)]
10#![warn(
11    deprecated_in_future,
12    missing_copy_implementations,
13    missing_debug_implementations,
14    missing_docs,
15    unreachable_pub,
16    unused_import_braces,
17    unused_labels,
18    unused_lifetimes,
19    unused_qualifications,
20    unused_results
21)]
22
23mod config;
24mod generic_client;
25mod statement_cache;
26
27use std::{
28    fmt,
29    future::Future,
30    ops::{Deref, DerefMut},
31    pin::Pin,
32    sync::Arc,
33};
34
35use deadpool::managed;
36#[cfg(not(target_arch = "wasm32"))]
37use tokio::spawn;
38use tokio::task::JoinHandle;
39use tokio_postgres::{
40    Client as PgClient, Config as PgConfig, Error, IsolationLevel, Statement,
41    Transaction as PgTransaction, TransactionBuilder as PgTransactionBuilder, types::Type,
42};
43
44#[cfg(not(target_arch = "wasm32"))]
45use tokio_postgres::{
46    Socket,
47    tls::{MakeTlsConnect, TlsConnect},
48};
49
50pub use tokio_postgres;
51
52pub use crate::statement_cache::{StatementCache, StatementCaches};
53
54pub use self::config::{
55    ChannelBinding, Config, ConfigError, LoadBalanceHosts, ManagerConfig, RecyclingMethod, SslMode,
56    TargetSessionAttrs,
57};
58
59pub use self::generic_client::GenericClient;
60
61pub use deadpool::managed::reexports::*;
62deadpool::managed_reexports!(
63    "tokio_postgres",
64    Manager,
65    managed::Object<Manager>,
66    Error,
67    ConfigError
68);
69
70type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
71
72/// Type alias for [`Object`]
73pub type Client = Object;
74
75type RecycleResult = managed::RecycleResult<Error>;
76type RecycleError = managed::RecycleError<Error>;
77
78/// [`Manager`] for creating and recycling PostgreSQL connections.
79///
80/// [`Manager`]: managed::Manager
81pub struct Manager {
82    config: ManagerConfig,
83    pg_config: PgConfig,
84    connect: Box<dyn Connect>,
85    /// [`StatementCaches`] of [`Client`]s handed out by the [`Pool`].
86    pub statement_caches: StatementCaches,
87}
88
89impl Manager {
90    #[cfg(not(target_arch = "wasm32"))]
91    /// Creates a new [`Manager`] using the given [`tokio_postgres::Config`] and
92    /// `tls` connector.
93    pub fn new<T>(pg_config: tokio_postgres::Config, tls: T) -> Self
94    where
95        T: MakeTlsConnect<Socket> + Clone + Sync + Send + 'static,
96        T::Stream: Sync + Send,
97        T::TlsConnect: Sync + Send,
98        <T::TlsConnect as TlsConnect<Socket>>::Future: Send,
99    {
100        Self::from_config(pg_config, tls, ManagerConfig::default())
101    }
102
103    #[cfg(not(target_arch = "wasm32"))]
104    /// Create a new [`Manager`] using the given [`tokio_postgres::Config`], and
105    /// `tls` connector and [`ManagerConfig`].
106    pub fn from_config<T>(pg_config: tokio_postgres::Config, tls: T, config: ManagerConfig) -> Self
107    where
108        T: MakeTlsConnect<Socket> + Clone + Sync + Send + 'static,
109        T::Stream: Sync + Send,
110        T::TlsConnect: Sync + Send,
111        <T::TlsConnect as TlsConnect<Socket>>::Future: Send,
112    {
113        Self::from_connect(pg_config, ConfigConnectImpl { tls }, config)
114    }
115
116    /// Create a new [`Manager`] using the given [`tokio_postgres::Config`], and
117    /// `connect` impl and [`ManagerConfig`].
118    pub fn from_connect(
119        pg_config: tokio_postgres::Config,
120        connect: impl Connect + 'static,
121        config: ManagerConfig,
122    ) -> Self {
123        Self {
124            config,
125            pg_config,
126            connect: Box::new(connect),
127            statement_caches: StatementCaches::default(),
128        }
129    }
130}
131
132impl fmt::Debug for Manager {
133    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
134        f.debug_struct("Manager")
135            .field("config", &self.config)
136            .field("pg_config", &self.pg_config)
137            //.field("connect", &self.connect)
138            .field("statement_caches", &self.statement_caches)
139            .finish()
140    }
141}
142
143impl managed::Manager for Manager {
144    type Type = ClientWrapper;
145    type Error = Error;
146
147    async fn create(&self) -> Result<ClientWrapper, Error> {
148        let (client, conn_task) = self.connect.connect(&self.pg_config).await?;
149        let client_wrapper = ClientWrapper::new(client, conn_task);
150        self.statement_caches
151            .attach(&client_wrapper.statement_cache);
152        Ok(client_wrapper)
153    }
154
155    async fn recycle(&self, client: &mut ClientWrapper, _: &Metrics) -> RecycleResult {
156        if client.is_closed() {
157            tracing::warn!(target: "deadpool.postgres", "Connection could not be recycled: Connection closed");
158            return Err(RecycleError::message("Connection closed"));
159        }
160        match self.config.recycling_method.query() {
161            Some(sql) => match client.simple_query(sql).await {
162                Ok(_) => Ok(()),
163                Err(e) => {
164                    tracing::warn!(target: "deadpool.postgres", "Connection could not be recycled: {}", e);
165                    Err(e.into())
166                }
167            },
168            None => Ok(()),
169        }
170    }
171
172    fn detach(&self, object: &mut ClientWrapper) {
173        self.statement_caches.detach(&object.statement_cache);
174    }
175}
176
177/// Describes a mechanism for establishing a connection to a PostgreSQL
178/// server via `tokio_postgres`.
179pub trait Connect: Sync + Send {
180    /// Establishes a new `tokio_postgres` connection, returning
181    /// the associated `Client` and a `JoinHandle` to a tokio task
182    /// for processing the connection.
183    fn connect(
184        &self,
185        pg_config: &PgConfig,
186    ) -> BoxFuture<'_, Result<(PgClient, JoinHandle<()>), Error>>;
187}
188
189#[cfg(not(target_arch = "wasm32"))]
190/// Provides an implementation of [`Connect`] that establishes the connection
191/// using the `tokio_postgres` configuration itself.
192#[derive(Debug)]
193pub struct ConfigConnectImpl<T>
194where
195    T: MakeTlsConnect<Socket> + Clone + Sync + Send + 'static,
196    T::Stream: Sync + Send,
197    T::TlsConnect: Sync + Send,
198    <T::TlsConnect as TlsConnect<Socket>>::Future: Send,
199{
200    /// The TLS connector to use for the connection.
201    pub tls: T,
202}
203
204#[cfg(not(target_arch = "wasm32"))]
205impl<T> Connect for ConfigConnectImpl<T>
206where
207    T: MakeTlsConnect<Socket> + Clone + Sync + Send + 'static,
208    T::Stream: Sync + Send,
209    T::TlsConnect: Sync + Send,
210    <T::TlsConnect as TlsConnect<Socket>>::Future: Send,
211{
212    fn connect(
213        &self,
214        pg_config: &PgConfig,
215    ) -> BoxFuture<'_, Result<(PgClient, JoinHandle<()>), Error>> {
216        let tls = self.tls.clone();
217        let pg_config = pg_config.clone();
218        Box::pin(async move {
219            let fut = pg_config.connect(tls);
220            let (client, connection) = fut.await?;
221            let conn_task = spawn(async move {
222                if let Err(e) = connection.await {
223                    tracing::warn!(target: "deadpool.postgres", "Connection error: {}", e);
224                }
225            });
226            Ok((client, conn_task))
227        })
228    }
229}
230
231/// Wrapper around [`tokio_postgres::Client`] with a [`StatementCache`].
232#[derive(Debug)]
233pub struct ClientWrapper {
234    /// Original [`PgClient`].
235    client: PgClient,
236
237    /// A handle to the connection task that should be aborted when the client
238    /// wrapper is dropped.
239    conn_task: JoinHandle<()>,
240
241    /// [`StatementCache`] of this client.
242    pub statement_cache: Arc<StatementCache>,
243}
244
245impl ClientWrapper {
246    /// Create a new [`ClientWrapper`] instance using the given
247    /// [`tokio_postgres::Client`] and handle to the connection task.
248    #[must_use]
249    pub fn new(client: PgClient, conn_task: JoinHandle<()>) -> Self {
250        Self {
251            client,
252            conn_task,
253            statement_cache: Arc::new(StatementCache::new()),
254        }
255    }
256
257    /// Like [`tokio_postgres::Client::prepare()`], but uses an existing
258    /// [`Statement`] from the [`StatementCache`] if possible.
259    pub async fn prepare_cached(&self, query: &str) -> Result<Statement, Error> {
260        self.statement_cache.prepare(&self.client, query).await
261    }
262
263    /// Like [`tokio_postgres::Client::prepare_typed()`], but uses an
264    /// existing [`Statement`] from the [`StatementCache`] if possible.
265    pub async fn prepare_typed_cached(
266        &self,
267        query: &str,
268        types: &[Type],
269    ) -> Result<Statement, Error> {
270        self.statement_cache
271            .prepare_typed(&self.client, query, types)
272            .await
273    }
274
275    /// Like [`tokio_postgres::Client::transaction()`], but returns a wrapped
276    /// [`Transaction`] with a [`StatementCache`].
277    #[allow(unused_lifetimes)] // false positive
278    pub async fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
279        Ok(Transaction {
280            txn: PgClient::transaction(&mut self.client).await?,
281            statement_cache: self.statement_cache.clone(),
282        })
283    }
284
285    /// Like [`tokio_postgres::Client::build_transaction()`], but creates a
286    /// wrapped [`Transaction`] with a [`StatementCache`].
287    pub fn build_transaction(&mut self) -> TransactionBuilder<'_> {
288        TransactionBuilder {
289            builder: self.client.build_transaction(),
290            statement_cache: self.statement_cache.clone(),
291        }
292    }
293}
294
295impl Deref for ClientWrapper {
296    type Target = PgClient;
297
298    fn deref(&self) -> &PgClient {
299        &self.client
300    }
301}
302
303impl DerefMut for ClientWrapper {
304    fn deref_mut(&mut self) -> &mut PgClient {
305        &mut self.client
306    }
307}
308
309impl Drop for ClientWrapper {
310    fn drop(&mut self) {
311        self.conn_task.abort()
312    }
313}
314
315/// Wrapper around [`tokio_postgres::Transaction`] with a [`StatementCache`]
316/// from the [`Client`] object it was created by.
317pub struct Transaction<'a> {
318    /// Original [`PgTransaction`].
319    txn: PgTransaction<'a>,
320
321    /// [`StatementCache`] of this [`Transaction`].
322    pub statement_cache: Arc<StatementCache>,
323}
324
325impl Transaction<'_> {
326    /// Like [`tokio_postgres::Transaction::prepare()`], but uses an existing
327    /// [`Statement`] from the [`StatementCache`] if possible.
328    pub async fn prepare_cached(&self, query: &str) -> Result<Statement, Error> {
329        self.statement_cache.prepare(self.client(), query).await
330    }
331
332    /// Like [`tokio_postgres::Transaction::prepare_typed()`], but uses an
333    /// existing [`Statement`] from the [`StatementCache`] if possible.
334    pub async fn prepare_typed_cached(
335        &self,
336        query: &str,
337        types: &[Type],
338    ) -> Result<Statement, Error> {
339        self.statement_cache
340            .prepare_typed(self.client(), query, types)
341            .await
342    }
343
344    /// Like [`tokio_postgres::Transaction::commit()`].
345    pub async fn commit(self) -> Result<(), Error> {
346        self.txn.commit().await
347    }
348
349    /// Like [`tokio_postgres::Transaction::rollback()`].
350    pub async fn rollback(self) -> Result<(), Error> {
351        self.txn.rollback().await
352    }
353
354    /// Like [`tokio_postgres::Transaction::transaction()`], but returns a
355    /// wrapped [`Transaction`] with a [`StatementCache`].
356    #[allow(unused_lifetimes)] // false positive
357    pub async fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
358        Ok(Transaction {
359            txn: PgTransaction::transaction(&mut self.txn).await?,
360            statement_cache: self.statement_cache.clone(),
361        })
362    }
363
364    /// Like [`tokio_postgres::Transaction::savepoint()`], but returns a wrapped
365    /// [`Transaction`] with a [`StatementCache`].
366    #[allow(unused_lifetimes)] // false positive
367    pub async fn savepoint<I>(&mut self, name: I) -> Result<Transaction<'_>, Error>
368    where
369        I: Into<String>,
370    {
371        Ok(Transaction {
372            txn: PgTransaction::savepoint(&mut self.txn, name).await?,
373            statement_cache: self.statement_cache.clone(),
374        })
375    }
376}
377
378impl fmt::Debug for Transaction<'_> {
379    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
380        f.debug_struct("Transaction")
381            //.field("txn", &self.txn)
382            .field("statement_cache", &self.statement_cache)
383            .finish()
384    }
385}
386
387impl<'a> Deref for Transaction<'a> {
388    type Target = PgTransaction<'a>;
389
390    fn deref(&self) -> &PgTransaction<'a> {
391        &self.txn
392    }
393}
394
395impl<'a> DerefMut for Transaction<'a> {
396    fn deref_mut(&mut self) -> &mut PgTransaction<'a> {
397        &mut self.txn
398    }
399}
400
401/// Wrapper around [`tokio_postgres::TransactionBuilder`] with a
402/// [`StatementCache`] from the [`Client`] object it was created by.
403#[must_use = "builder does nothing itself, use `.start()` to use it"]
404pub struct TransactionBuilder<'a> {
405    /// Original [`PgTransactionBuilder`].
406    builder: PgTransactionBuilder<'a>,
407
408    /// [`StatementCache`] of this [`TransactionBuilder`].
409    statement_cache: Arc<StatementCache>,
410}
411
412impl<'a> TransactionBuilder<'a> {
413    /// Sets the isolation level of the transaction.
414    ///
415    /// Like [`tokio_postgres::TransactionBuilder::isolation_level()`].
416    pub fn isolation_level(self, isolation_level: IsolationLevel) -> Self {
417        Self {
418            builder: self.builder.isolation_level(isolation_level),
419            statement_cache: self.statement_cache,
420        }
421    }
422
423    /// Sets the access mode of the transaction.
424    ///
425    /// Like [`tokio_postgres::TransactionBuilder::read_only()`].
426    pub fn read_only(self, read_only: bool) -> Self {
427        Self {
428            builder: self.builder.read_only(read_only),
429            statement_cache: self.statement_cache,
430        }
431    }
432
433    /// Sets the deferrability of the transaction.
434    ///
435    /// If the transaction is also serializable and read only, creation
436    /// of the transaction may block, but when it completes the transaction
437    /// is able to run with less overhead and a guarantee that it will not
438    /// be aborted due to serialization failure.
439    ///
440    /// Like [`tokio_postgres::TransactionBuilder::deferrable()`].
441    pub fn deferrable(self, deferrable: bool) -> Self {
442        Self {
443            builder: self.builder.deferrable(deferrable),
444            statement_cache: self.statement_cache,
445        }
446    }
447
448    /// Begins the [`Transaction`].
449    ///
450    /// The transaction will roll back by default - use the commit method
451    /// to commit it.
452    ///
453    /// Like [`tokio_postgres::TransactionBuilder::start()`].
454    pub async fn start(self) -> Result<Transaction<'a>, Error> {
455        Ok(Transaction {
456            txn: self.builder.start().await?,
457            statement_cache: self.statement_cache,
458        })
459    }
460}
461
462impl fmt::Debug for TransactionBuilder<'_> {
463    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
464        f.debug_struct("TransactionBuilder")
465            //.field("builder", &self.builder)
466            .field("statement_cache", &self.statement_cache)
467            .finish()
468    }
469}
470
471impl<'a> Deref for TransactionBuilder<'a> {
472    type Target = PgTransactionBuilder<'a>;
473
474    fn deref(&self) -> &Self::Target {
475        &self.builder
476    }
477}
478
479impl DerefMut for TransactionBuilder<'_> {
480    fn deref_mut(&mut self) -> &mut Self::Target {
481        &mut self.builder
482    }
483}
deadpool_postgres/lib.rs

deadpool_postgres/
lib.rs
