grafbase_sdk/host_io/postgres/
pool.rs

1use std::time::Duration;
2
3use crate::{SdkError, wit};
4
5use super::{Connection, Transaction};
6
7/// Configuration options for a Postgres connection pool.
8///
9/// This struct provides various configuration options for controlling the behavior
10/// of a Postgres connection pool, such as connection limits and timeouts.
11pub struct PoolOptions(wit::PgPoolOptions);
12
13impl Default for PoolOptions {
14    fn default() -> Self {
15        Self(wit::PgPoolOptions {
16            max_connections: None,
17            min_connections: None,
18            idle_timeout_ms: None,
19            acquisition_timeout_ms: None,
20            max_lifetime_ms: None,
21        })
22    }
23}
24
25impl PoolOptions {
26    /// Creates a new `PoolOptions` instance with default settings.
27    pub fn new() -> Self {
28        Self::default()
29    }
30
31    /// Sets the maximum number of connections in the pool.
32    ///
33    /// This limits how many database connections the pool can open simultaneously.
34    pub fn max_connections(mut self, max_connections: u32) -> Self {
35        self.0.max_connections = Some(max_connections);
36        self
37    }
38
39    /// Sets the minimum number of connections the pool will maintain.
40    ///
41    /// The pool will attempt to maintain this many idle connections at all times.
42    pub fn min_connections(mut self, min_connections: u32) -> Self {
43        self.0.min_connections = Some(min_connections);
44        self
45    }
46
47    /// Sets the idle timeout for connections in the pool.
48    ///
49    /// Connections that remain idle for longer than this duration may be closed.
50    pub fn idle_timeout(mut self, idle_timeout: Duration) -> Self {
51        self.0.idle_timeout_ms = Some(idle_timeout.as_millis() as u64);
52        self
53    }
54
55    /// Sets the maximum time to wait when acquiring a connection from the pool.
56    ///
57    /// If a connection cannot be acquired within this time, the acquisition attempt fails.
58    pub fn acquire_timeout(mut self, acquire_timeout: Duration) -> Self {
59        self.0.acquisition_timeout_ms = Some(acquire_timeout.as_millis() as u64);
60        self
61    }
62
63    /// Sets the maximum lifetime of connections in the pool.
64    ///
65    /// Connections will be closed after this duration regardless of whether they're idle.
66    pub fn max_lifetime(mut self, max_lifetime: Duration) -> Self {
67        self.0.max_lifetime_ms = Some(max_lifetime.as_millis() as u64);
68        self
69    }
70}
71
72/// A Postgres connection pool.
73///
74/// This pool manages a set of database connections that can be reused across
75/// different operations, improving performance by avoiding the overhead of
76/// establishing new connections for each database operation.
77pub struct Pool(wit::PgPool);
78
79impl Pool {
80    /// Creates a new connection pool with default options.
81    ///
82    /// # Parameters
83    /// * `identifier` - A unique identifier for the pool
84    /// * `url` - The Postgres connection URL
85    ///
86    /// # Returns
87    /// A new connection pool or an error if the connection fails
88    pub fn connect(identifier: &str, url: &str) -> Result<Self, SdkError> {
89        Self::connect_with_options(identifier, url, Default::default())
90    }
91
92    /// Creates a new connection pool with custom options.
93    ///
94    /// # Parameters
95    /// * `identifier` - A unique identifier for the pool
96    /// * `url` - The Postgres connection URL
97    /// * `options` - Configuration options for the connection pool
98    ///
99    /// # Returns
100    /// A new connection pool or an error if the connection fails
101    pub fn connect_with_options(identifier: &str, url: &str, options: PoolOptions) -> Result<Self, SdkError> {
102        let pool = wit::PgPool::connect(identifier, url, options.0).map_err(SdkError::from)?;
103        Ok(Self(pool))
104    }
105
106    /// Acquires a connection from the pool.
107    ///
108    /// This function will wait until a connection is available or until
109    /// the configured acquisition timeout is reached.
110    ///
111    /// # Returns
112    /// A database connection or an error if no connection could be acquired
113    pub fn acquire(&self) -> Result<Connection, SdkError> {
114        self.0.acquire().map(Into::into).map_err(SdkError::from)
115    }
116
117    /// Begins a new database transaction.
118    ///
119    /// This is a convenience method that acquires a connection and starts
120    /// a transaction on it.
121    ///
122    /// # Returns
123    /// A transaction object or an error if the transaction couldn't be started
124    pub fn begin_transaction(&self) -> Result<Transaction, SdkError> {
125        self.0.begin_transaction().map(Into::into).map_err(SdkError::from)
126    }
127}