Skip to main content

sqlx_pool_registry/
test_db_pools.rs

1use crate::{
2    provider::PoolProvider,
3    sqlx::{self, PgPool},
4};
5
6/// Test pool provider with a replica that is read-only by default.
7///
8/// This creates two separate connection pools from the same database:
9/// - Primary pool for writes (normal permissions)
10/// - Replica pool for reads (sets `default_transaction_read_only = on`)
11///
12/// This helps tests catch bugs where ordinary write operations are incorrectly
13/// routed through `.read()`. PostgreSQL will reject writes to non-temporary
14/// tables by default with errors such as:
15/// "cannot execute INSERT/UPDATE/DELETE in a read-only transaction"
16///
17/// This helper is not a security boundary. A client can override the default
18/// for an individual transaction or session, and PostgreSQL read-only
19/// transactions do not prohibit every possible write. See PostgreSQL's
20/// [`SET TRANSACTION`](https://www.postgresql.org/docs/current/sql-set-transaction.html)
21/// documentation for the exact restrictions.
22///
23/// # Usage with `#[sqlx::test]`
24///
25/// ```rust,no_run
26/// use sqlx_pool_registry::sqlx::{self, PgPool};
27/// use sqlx_pool_registry::{PoolProvider, TestDbPools};
28///
29/// #[sqlx::test]
30/// async fn test_read_write_routing(pool: PgPool) {
31///     let pools = TestDbPools::new(pool).await.unwrap();
32///
33///     // Write operations work on .write()
34///     sqlx::query("CREATE TEMP TABLE users (id INT)")
35///         .execute(pools.write())
36///         .await
37///         .expect("Write pool should allow writes");
38///
39///     // Write operations FAIL on .read()
40///     let result = sqlx::query("INSERT INTO users VALUES (1)")
41///         .execute(pools.read())
42///         .await;
43///     assert!(result.is_err(), "Read pool should reject writes");
44///
45///     // Read operations work on .read()
46///     let count: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM users")
47///         .fetch_one(pools.read())
48///         .await
49///         .expect("Read pool should allow reads");
50/// }
51/// ```
52///
53/// # Why This Matters
54///
55/// Without this test helper, you might accidentally route write operations through
56/// `.read()` and not catch the bug until production when you have an actual replica
57/// with replication lag. This helper makes ordinary routing mistakes easier to
58/// detect in tests.
59///
60/// # Example
61///
62/// ```rust,no_run
63/// use sqlx_pool_registry::sqlx::{self, PgPool};
64/// use sqlx_pool_registry::{PoolProvider, TestDbPools};
65///
66/// struct Repository<P: PoolProvider> {
67///     pools: P,
68/// }
69///
70/// impl<P: PoolProvider> Repository<P> {
71///     async fn get_user(&self, id: i64) -> Result<String, sqlx::Error> {
72///         sqlx::query_scalar("SELECT name FROM users WHERE id = $1")
73///             .bind(id)
74///             .fetch_one(self.pools.read())
75///             .await
76///     }
77///
78///     async fn create_user(&self, name: &str) -> Result<i64, sqlx::Error> {
79///         sqlx::query_scalar("INSERT INTO users (name) VALUES ($1) RETURNING id")
80///             .bind(name)
81///             .fetch_one(self.pools.write())
82///             .await
83///     }
84/// }
85///
86/// #[sqlx::test]
87/// async fn test_repository_routing(pool: PgPool) {
88///     let pools = TestDbPools::new(pool).await.unwrap();
89///     let repo = Repository { pools };
90///
91///     // Test will fail if create_user incorrectly uses .read()
92///     sqlx::query("CREATE TEMP TABLE users (id SERIAL PRIMARY KEY, name TEXT)")
93///         .execute(repo.pools.write())
94///         .await
95///         .unwrap();
96///
97///     let user_id = repo.create_user("Alice").await.unwrap();
98///     let name = repo.get_user(user_id).await.unwrap();
99///     assert_eq!(name, "Alice");
100/// }
101/// ```
102#[derive(Clone, Debug)]
103pub struct TestDbPools {
104    primary: PgPool,
105    replica: PgPool,
106}
107
108impl TestDbPools {
109    /// Create test pools from a single database pool.
110    ///
111    /// This creates:
112    /// - A primary pool (clone of input) for writes
113    /// - A replica pool (new connection) configured as read-only by default
114    ///
115    /// The replica pool sets `default_transaction_read_only = on`, so ordinary
116    /// writes to non-temporary tables fail unless the client overrides that
117    /// default. This is a testing aid, not an authorization mechanism.
118    ///
119    /// # Example
120    ///
121    /// ```rust,no_run
122    /// use sqlx_pool_registry::sqlx::{self, PgPool};
123    /// use sqlx_pool_registry::TestDbPools;
124    ///
125    /// # async fn example(pool: PgPool) -> Result<(), sqlx::Error> {
126    /// let pools = TestDbPools::new(pool).await?;
127    ///
128    /// // Now ordinary write-through-read mistakes fail during tests
129    /// # Ok(())
130    /// # }
131    /// ```
132    pub async fn new(pool: PgPool) -> Result<Self, sqlx::Error> {
133        use crate::sqlx::postgres::PgPoolOptions;
134
135        let primary = pool.clone();
136
137        // Create a separate pool that is read-only by default
138        let replica = PgPoolOptions::new()
139            .max_connections(pool.options().get_max_connections())
140            .after_connect(|conn, _meta| {
141                Box::pin(async move {
142                    // Set all transactions to read-only by default
143                    sqlx::query("SET default_transaction_read_only = on")
144                        .execute(&mut *conn)
145                        .await?;
146                    Ok(())
147                })
148            })
149            .connect_with(pool.connect_options().as_ref().clone())
150            .await?;
151
152        Ok(Self { primary, replica })
153    }
154}
155
156impl PoolProvider for TestDbPools {
157    fn read(&self) -> &PgPool {
158        &self.replica
159    }
160
161    fn write(&self) -> &PgPool {
162        &self.primary
163    }
164}