pg_worm/
config.rs

1//! This module contains the code for configuring the connection pool.
2
3use std::str::FromStr;
4
5use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod, Runtime};
6use tokio_postgres::{
7    tls::{MakeTlsConnect, TlsConnect},
8    Config, NoTls, Socket,
9};
10
11/// An empty struct that only provides the `build` method.
12pub struct Connection;
13
14///
15pub struct ConnectionBuilder<Tls>
16where
17    Tls: MakeTlsConnect<Socket> + Clone + Send + Sync + 'static,
18    Tls::Stream: Sync + Send,
19    Tls::TlsConnect: Sync + Send,
20    <Tls::TlsConnect as TlsConnect<Socket>>::Future: Send,
21{
22    url: String,
23    tls: Tls,
24    recycling_method: RecyclingMethod,
25    max_pool_size: Option<usize>,
26    runtime: Option<Runtime>,
27}
28
29impl<Tls> ConnectionBuilder<Tls>
30where
31    Tls: MakeTlsConnect<Socket> + Clone + Send + Sync + 'static,
32    Tls::Stream: Sync + Send,
33    Tls::TlsConnect: Sync + Send,
34    <Tls::TlsConnect as TlsConnect<Socket>>::Future: Send,
35{
36    fn to(url: impl Into<String>) -> ConnectionBuilder<NoTls> {
37        ConnectionBuilder {
38            url: url.into(),
39            tls: NoTls,
40            recycling_method: RecyclingMethod::Fast,
41            max_pool_size: None,
42            runtime: None,
43        }
44    }
45
46    /// Set the Tls method.
47    ///
48    /// Use either [`postgres-openssl`](https://crates.io/crates/postgres-openssl)
49    /// or [`postgres-nativ-tls`](https://crates.io/crates/postgres-native-tls)
50    /// and their respective documentation.
51    /// This function accepts the same types as `tokio-postgres`.
52    pub fn tls<NewTls>(self, tls: NewTls) -> ConnectionBuilder<NewTls>
53    where
54        NewTls: MakeTlsConnect<Socket> + Clone + Send + Sync + 'static,
55        NewTls::Stream: Sync + Send,
56        NewTls::TlsConnect: Sync + Send,
57        <NewTls::TlsConnect as TlsConnect<Socket>>::Future: Send,
58    {
59        ConnectionBuilder {
60            tls,
61            url: self.url,
62            recycling_method: self.recycling_method,
63            runtime: self.runtime,
64            max_pool_size: self.max_pool_size,
65        }
66    }
67
68    /// Set the maximum size of the connection pool,
69    /// i.e. the maximum amount of concurrent connections to the database server.
70    ///
71    /// The default is `num_cpus * 4`, ignoring hyperthreading, etc.
72    pub fn max_pool_size(mut self, n: usize) -> Self {
73        self.max_pool_size = Some(n);
74
75        self
76    }
77
78    /// Finish the setup and build the pool.
79    ///
80    /// Fails if
81    ///  - the url couldn't be parsed, or
82    ///  - some other configuration error has been made.
83    pub fn connect(self) -> Result<(), crate::Error> {
84        let config = Config::from_str(&self.url)?;
85        let manager_config = ManagerConfig {
86            recycling_method: self.recycling_method,
87        };
88
89        let manager = Manager::from_config(config, self.tls, manager_config);
90        let mut builder = Pool::builder(manager).runtime(Runtime::Tokio1);
91
92        if let Some(n) = self.max_pool_size {
93            builder = builder.max_size(n);
94        }
95
96        let pool = builder.build()?;
97
98        crate::set_pool(pool)?;
99
100        Ok(())
101    }
102}
103
104impl Connection {
105    /// Start building a new connection (pool).
106    ///
107    /// This returns a [`ConnectionBuilder`] which can be configured
108    /// using the builder pattern.
109    ///
110    /// If you are fine with the default configuration
111    /// (`max_pool_size = num_cpus * 4` and no Tls) or have
112    /// configured to your needs you  can finish the setup
113    /// by calling `.connect()`.
114    ///  
115    /// A connection must be created before executing any
116    /// queries or the like.
117    /// Doing otherwise will result in a runime error.
118    ///
119    /// # Example
120    /// ```ignore
121    /// use pg_worm::prelude::*;
122    ///
123    /// Connection::build("postgres://postgres").connect()?;    
124    /// ```
125    ///
126    pub fn build(connection_string: impl Into<String>) -> ConnectionBuilder<NoTls> {
127        ConnectionBuilder::<NoTls>::to(connection_string)
128    }
129}