spring_batch_rs/item/rdbc/postgres_reader.rs
1use std::cell::{Cell, RefCell};
2
3use sqlx::{FromRow, Pool, Postgres, QueryBuilder, postgres::PgRow};
4
5use super::reader_common::read_item;
6use crate::BatchError;
7use crate::core::item::{ItemReader, ItemReaderResult};
8
9/// PostgreSQL RDBC Item Reader for batch processing.
10///
11/// Supports two pagination strategies:
12///
13/// - **LIMIT/OFFSET** (default): simple but degrades at large offsets — use for small datasets.
14/// - **Keyset pagination** (recommended for large datasets): uses `WHERE col > :last ORDER BY col LIMIT n`,
15/// O(log n) per page regardless of dataset size. Enable with
16/// [`RdbcItemReaderBuilder::with_keyset`](crate::item::rdbc::RdbcItemReaderBuilder::with_keyset).
17///
18/// # Type Parameters
19///
20/// * `I` - Must implement `FromRow<PgRow> + Send + Unpin + Clone`.
21///
22/// # Construction
23///
24/// Prefer [`RdbcItemReaderBuilder`](crate::item::rdbc::RdbcItemReaderBuilder) for ergonomic construction.
25pub struct PostgresRdbcItemReader<I>
26where
27 for<'r> I: FromRow<'r, PgRow> + Send + Unpin + Clone,
28{
29 pub(crate) pool: Pool<Postgres>,
30 pub(crate) query: String,
31 pub(crate) page_size: Option<i32>,
32 pub(crate) offset: Cell<i32>,
33 pub(crate) buffer: RefCell<Vec<I>>,
34 /// Column name used as the keyset cursor (e.g. `"id"`).
35 pub(crate) keyset_column: Option<String>,
36 /// Extracts the cursor value from an item for use in the next page's WHERE clause.
37 #[allow(clippy::type_complexity)]
38 pub(crate) keyset_key: Option<Box<dyn Fn(&I) -> String>>,
39 /// Last cursor value seen; drives the WHERE clause on subsequent pages.
40 pub(crate) last_cursor: RefCell<Option<String>>,
41}
42
43impl<I> PostgresRdbcItemReader<I>
44where
45 for<'r> I: FromRow<'r, PgRow> + Send + Unpin + Clone,
46{
47 /// Creates a new `PostgresRdbcItemReader` with the specified parameters.
48 ///
49 /// Prefer [`RdbcItemReaderBuilder`](crate::item::rdbc::RdbcItemReaderBuilder) for a more
50 /// ergonomic construction API.
51 ///
52 /// # Arguments
53 ///
54 /// * `pool` - PostgreSQL connection pool for database operations
55 /// * `query` - SQL query to execute (without LIMIT/OFFSET)
56 /// * `page_size` - Optional page size for pagination. None means read all at once.
57 /// * `keyset_column` - Optional column name for keyset pagination.
58 /// * `keyset_key` - Optional closure to extract the cursor value from an item.
59 ///
60 /// # Returns
61 ///
62 /// A new `PostgresRdbcItemReader` instance ready for use.
63 #[allow(clippy::type_complexity)]
64 pub fn new(
65 pool: Pool<Postgres>,
66 query: String,
67 page_size: Option<i32>,
68 keyset_column: Option<String>,
69 keyset_key: Option<Box<dyn Fn(&I) -> String>>,
70 ) -> Self {
71 Self {
72 pool,
73 query,
74 page_size,
75 offset: Cell::new(0),
76 buffer: RefCell::new(vec![]),
77 keyset_column,
78 keyset_key,
79 last_cursor: RefCell::new(None),
80 }
81 }
82
83 /// Fetches the next page from the database into the internal buffer.
84 ///
85 /// Uses keyset pagination when `keyset_column` is set, otherwise falls back
86 /// to LIMIT/OFFSET.
87 ///
88 /// # Errors
89 ///
90 /// Returns [`BatchError::ItemReader`] if the query fails.
91 fn read_page(&self) -> Result<(), BatchError> {
92 let mut query_builder = QueryBuilder::<Postgres>::new(&self.query);
93
94 if let Some(page_size) = self.page_size {
95 if let Some(ref col) = self.keyset_column {
96 let last = self.last_cursor.borrow();
97 if let Some(ref cursor_val) = *last {
98 let escaped = cursor_val.replace('\'', "''");
99 query_builder.push(format!(" WHERE {} > '{}'", col, escaped));
100 }
101 drop(last);
102 query_builder.push(format!(" ORDER BY {} LIMIT {}", col, page_size));
103 } else {
104 query_builder.push(format!(" LIMIT {} OFFSET {}", page_size, self.offset.get()));
105 }
106 }
107
108 let query = query_builder.build_query_as::<I>();
109 let items = tokio::task::block_in_place(|| {
110 tokio::runtime::Handle::current().block_on(async {
111 query
112 .fetch_all(&self.pool)
113 .await
114 .map_err(|e| BatchError::ItemReader(e.to_string()))
115 })
116 })?;
117
118 *self.buffer.borrow_mut() = items;
119 Ok(())
120 }
121}
122
123impl<I> ItemReader<I> for PostgresRdbcItemReader<I>
124where
125 for<'r> I: FromRow<'r, PgRow> + Send + Unpin + Clone,
126{
127 /// Reads the next item from the PostgreSQL database.
128 ///
129 /// Manages automatic pagination: loads a new page when the buffer is exhausted,
130 /// handles both LIMIT/OFFSET and keyset pagination transparently.
131 ///
132 /// # Returns
133 ///
134 /// - `Ok(Some(item))` if an item was successfully read
135 /// - `Ok(None)` if there are no more items to read (end of result set)
136 /// - `Err(BatchError::ItemReader)` if a database error occurred
137 ///
138 /// # Examples
139 ///
140 /// ```no_run
141 /// use spring_batch_rs::core::item::ItemReader;
142 /// use spring_batch_rs::item::rdbc::RdbcItemReaderBuilder;
143 /// use sqlx::PgPool;
144 /// # use serde::Deserialize;
145 /// # #[derive(sqlx::FromRow, Clone, Deserialize)]
146 /// # struct User { id: i32, name: String }
147 ///
148 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
149 /// let pool = PgPool::connect("postgresql://user:pass@localhost/db").await?;
150 /// let reader = RdbcItemReaderBuilder::<User>::new()
151 /// .postgres(pool)
152 /// .query("SELECT id, name FROM users ORDER BY id")
153 /// .with_page_size(100)
154 /// .build_postgres();
155 ///
156 /// // Read items one by one
157 /// let mut count = 0;
158 /// while let Some(user) = reader.read()? {
159 /// println!("User: {} - {}", user.id, user.name);
160 /// count += 1;
161 /// }
162 /// println!("Processed {} users", count);
163 /// # Ok(())
164 /// # }
165 /// ```
166 fn read(&self) -> ItemReaderResult<I> {
167 read_item(
168 &self.offset,
169 self.page_size,
170 &self.buffer,
171 &self.keyset_key,
172 &self.last_cursor,
173 || self.read_page(),
174 )
175 }
176}
177
178#[cfg(test)]
179mod tests {
180 use super::*;
181 use sqlx::PgPool;
182
183 #[derive(Clone)]
184 struct Dummy;
185
186 impl<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> for Dummy {
187 fn from_row(_row: &'r sqlx::postgres::PgRow) -> Result<Self, sqlx::Error> {
188 Ok(Dummy)
189 }
190 }
191
192 fn reader_with_keyset(keyset: bool) -> PostgresRdbcItemReader<Dummy> {
193 let pool = PgPool::connect_lazy("postgres://postgres:postgres@localhost/test")
194 .expect("lazy pool creation should not fail");
195 let (col, key): (Option<String>, Option<Box<dyn Fn(&Dummy) -> String>>) = if keyset {
196 (
197 Some("id".to_string()),
198 Some(Box::new(|_: &Dummy| "0".to_string())),
199 )
200 } else {
201 (None, None)
202 };
203 PostgresRdbcItemReader::new(pool, "SELECT 1".to_string(), Some(10), col, key)
204 }
205
206 #[tokio::test(flavor = "multi_thread")]
207 async fn should_initialize_without_keyset() {
208 let reader = reader_with_keyset(false);
209 assert!(reader.keyset_column.is_none(), "no keyset column expected");
210 assert!(reader.keyset_key.is_none(), "no keyset key fn expected");
211 assert!(
212 reader.last_cursor.borrow().is_none(),
213 "cursor must start as None"
214 );
215 assert_eq!(reader.offset.get(), 0, "initial offset should be 0");
216 assert!(
217 reader.buffer.borrow().is_empty(),
218 "buffer should start empty"
219 );
220 assert_eq!(reader.page_size, Some(10));
221 }
222
223 #[tokio::test(flavor = "multi_thread")]
224 async fn should_initialize_with_keyset_column_and_none_cursor() {
225 let reader = reader_with_keyset(true);
226 assert_eq!(
227 reader.keyset_column.as_deref(),
228 Some("id"),
229 "keyset column should be stored"
230 );
231 assert!(
232 reader.keyset_key.is_some(),
233 "keyset key fn should be stored"
234 );
235 assert!(
236 reader.last_cursor.borrow().is_none(),
237 "cursor must start as None before first read"
238 );
239 }
240}