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}