Skip to main content

sqlx_pool_registry/
db_pools.rs

1use std::ops::Deref;
2
3use crate::{provider::PoolProvider, sqlx::PgPool};
4
5/// Database pool abstraction supporting read replicas.
6///
7/// Wraps primary and optional replica pools, providing methods for
8/// explicit read/write routing while maintaining backwards compatibility
9/// through `Deref<Target = PgPool>`.
10///
11/// # Examples
12///
13/// ## Single Pool Configuration
14///
15/// ```rust,no_run
16/// use sqlx_pool_registry::sqlx::{self, PgPool};
17/// use sqlx_pool_registry::DbPools;
18///
19/// # async fn example() -> Result<(), sqlx::Error> {
20/// let pool = PgPool::connect("postgresql://localhost/db").await?;
21/// let pools = DbPools::new(pool);
22///
23/// // Both read() and write() return the same pool
24/// assert!(!pools.has_replica());
25/// # Ok(())
26/// # }
27/// ```
28///
29/// ## Primary/Replica Configuration
30///
31/// ```rust,no_run
32/// use sqlx_pool_registry::sqlx::{self, postgres::PgPoolOptions};
33/// use sqlx_pool_registry::DbPools;
34///
35/// # async fn example() -> Result<(), sqlx::Error> {
36/// let primary = PgPoolOptions::new()
37///     .max_connections(5)
38///     .connect("postgresql://primary/db")
39///     .await?;
40///
41/// let replica = PgPoolOptions::new()
42///     .max_connections(10)
43///     .connect("postgresql://replica/db")
44///     .await?;
45///
46/// let pools = DbPools::with_replica(primary, replica);
47/// assert!(pools.has_replica());
48/// # Ok(())
49/// # }
50/// ```
51#[derive(Clone, Debug)]
52pub struct DbPools {
53    primary: PgPool,
54    replica: Option<PgPool>,
55}
56
57impl DbPools {
58    /// Create a new DbPools with only a primary pool.
59    ///
60    /// This is useful for development or when you don't have a read replica configured.
61    /// All read and write operations will route to the primary pool.
62    ///
63    /// # Example
64    ///
65    /// ```rust,no_run
66    /// use sqlx_pool_registry::sqlx::{self, PgPool};
67    /// use sqlx_pool_registry::DbPools;
68    ///
69    /// # async fn example() -> Result<(), sqlx::Error> {
70    /// let pool = PgPool::connect("postgresql://localhost/db").await?;
71    /// let pools = DbPools::new(pool);
72    /// # Ok(())
73    /// # }
74    /// ```
75    pub fn new(primary: PgPool) -> Self {
76        Self {
77            primary,
78            replica: None,
79        }
80    }
81
82    /// Create a new DbPools with primary and replica pools.
83    ///
84    /// Read operations will route to the replica pool for load distribution,
85    /// while write operations always use the primary pool.
86    ///
87    /// # Example
88    ///
89    /// ```rust,no_run
90    /// use sqlx_pool_registry::sqlx::{self, postgres::PgPoolOptions};
91    /// use sqlx_pool_registry::DbPools;
92    ///
93    /// # async fn example() -> Result<(), sqlx::Error> {
94    /// let primary = PgPoolOptions::new()
95    ///     .max_connections(5)
96    ///     .connect("postgresql://primary/db")
97    ///     .await?;
98    ///
99    /// let replica = PgPoolOptions::new()
100    ///     .max_connections(10)
101    ///     .connect("postgresql://replica/db")
102    ///     .await?;
103    ///
104    /// let pools = DbPools::with_replica(primary, replica);
105    /// # Ok(())
106    /// # }
107    /// ```
108    pub fn with_replica(primary: PgPool, replica: PgPool) -> Self {
109        Self {
110            primary,
111            replica: Some(replica),
112        }
113    }
114
115    /// Get the primary pool.
116    ///
117    /// Unlike [`write`](PoolProvider::write), this method describes the physical
118    /// pool topology rather than the operation being routed. Both methods return
119    /// the same pool for `DbPools`.
120    pub fn primary(&self) -> &PgPool {
121        &self.primary
122    }
123
124    /// Get the configured replica pool, if one exists.
125    ///
126    /// This method does not fall back to the primary pool. Use
127    /// [`read`](PoolProvider::read) when a routable read pool is required.
128    pub fn replica(&self) -> Option<&PgPool> {
129        self.replica.as_ref()
130    }
131
132    /// Check if a replica pool is configured.
133    ///
134    /// Returns `true` if a replica pool was provided via [`with_replica`](Self::with_replica).
135    ///
136    /// # Example
137    ///
138    /// ```rust,no_run
139    /// use sqlx_pool_registry::sqlx::{self, PgPool};
140    /// use sqlx_pool_registry::DbPools;
141    ///
142    /// # async fn example() -> Result<(), sqlx::Error> {
143    /// let pool = PgPool::connect("postgresql://localhost/db").await?;
144    /// let pools = DbPools::new(pool);
145    /// assert!(!pools.has_replica());
146    /// # Ok(())
147    /// # }
148    /// ```
149    pub fn has_replica(&self) -> bool {
150        self.replica().is_some()
151    }
152
153    /// Close all database connections.
154    ///
155    /// Closes both primary and replica pools (if configured).
156    ///
157    /// # Example
158    ///
159    /// ```rust,no_run
160    /// use sqlx_pool_registry::sqlx::{self, PgPool};
161    /// use sqlx_pool_registry::DbPools;
162    ///
163    /// # async fn example() -> Result<(), sqlx::Error> {
164    /// let pool = PgPool::connect("postgresql://localhost/db").await?;
165    /// let pools = DbPools::new(pool);
166    /// pools.close().await;
167    /// # Ok(())
168    /// # }
169    /// ```
170    pub async fn close(&self) {
171        self.primary().close().await;
172        if let Some(replica) = self.replica() {
173            replica.close().await;
174        }
175    }
176}
177
178impl PoolProvider for DbPools {
179    fn read(&self) -> &PgPool {
180        self.replica().unwrap_or(self.primary())
181    }
182
183    fn write(&self) -> &PgPool {
184        self.primary()
185    }
186}
187
188/// Dereferences to the primary pool.
189///
190/// This allows natural usage like `&*pools` when you need a `&PgPool`.
191/// For explicit routing, use `.read()` or `.write()` methods.
192impl Deref for DbPools {
193    type Target = PgPool;
194
195    fn deref(&self) -> &Self::Target {
196        self.primary()
197    }
198}