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}