spring_batch_rs/item/rdbc/unified_reader_builder.rs
1use sqlx::{
2 FromRow, MySql, Pool, Postgres, Sqlite, mysql::MySqlRow, postgres::PgRow, sqlite::SqliteRow,
3};
4use std::marker::PhantomData;
5
6use super::database_type::DatabaseType;
7use super::mysql_reader::MySqlRdbcItemReader;
8use super::postgres_reader::PostgresRdbcItemReader;
9use super::select_builder::SelectBuilder;
10use super::sqlite_reader::SqliteRdbcItemReader;
11
12/// Source of the SQL query for an RDBC item reader.
13///
14/// This is an internal type used by [`RdbcItemReaderBuilder`] to track whether
15/// the query was provided as a raw string via [`.query()`] or constructed via
16/// a [`SelectBuilder`] with [`.select()`].
17enum QuerySource<'a> {
18 /// A raw SQL string provided directly by the caller.
19 Raw(&'a str),
20 /// A SQL string built by [`SelectBuilder`].
21 Built(String),
22}
23
24/// Unified builder for creating RDBC item readers for any supported database type.
25///
26/// This builder provides a consistent API for constructing database readers
27/// regardless of the underlying database (PostgreSQL, MySQL, or SQLite).
28/// Users specify the database type and connection pool, and the builder
29/// handles the creation of the appropriate reader implementation.
30///
31/// # Type Parameters
32///
33/// * `I` - The item type that implements the appropriate `FromRow` trait for the database
34///
35/// # Examples
36///
37/// ## PostgreSQL
38/// ```no_run
39/// use spring_batch_rs::item::rdbc::{RdbcItemReaderBuilder, DatabaseType};
40/// use sqlx::PgPool;
41/// # use serde::Deserialize;
42/// # #[derive(sqlx::FromRow, Clone, Deserialize)]
43/// # struct User { id: i32, name: String }
44///
45/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
46/// let pool = PgPool::connect("postgresql://user:pass@localhost/db").await?;
47///
48/// let reader = RdbcItemReaderBuilder::<User>::new()
49/// .postgres(pool)
50/// .query("SELECT id, name FROM users")
51/// .with_page_size(100)
52/// .build_postgres();
53/// # Ok(())
54/// # }
55/// ```
56///
57/// ## MySQL
58/// ```no_run
59/// use spring_batch_rs::item::rdbc::{RdbcItemReaderBuilder, DatabaseType};
60/// use sqlx::MySqlPool;
61/// # use serde::Deserialize;
62/// # #[derive(sqlx::FromRow, Clone, Deserialize)]
63/// # struct Product { id: i32, name: String }
64///
65/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
66/// let pool = MySqlPool::connect("mysql://user:pass@localhost/db").await?;
67///
68/// let reader = RdbcItemReaderBuilder::<Product>::new()
69/// .mysql(pool)
70/// .query("SELECT id, name FROM products")
71/// .with_page_size(100)
72/// .build_mysql();
73/// # Ok(())
74/// # }
75/// ```
76///
77/// ## SQLite
78/// ```no_run
79/// use spring_batch_rs::item::rdbc::{RdbcItemReaderBuilder, DatabaseType};
80/// use sqlx::SqlitePool;
81/// # use serde::Deserialize;
82/// # #[derive(sqlx::FromRow, Clone, Deserialize)]
83/// # struct Task { id: i32, title: String }
84///
85/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
86/// let pool = SqlitePool::connect("sqlite::memory:").await?;
87///
88/// let reader = RdbcItemReaderBuilder::<Task>::new()
89/// .sqlite(pool)
90/// .query("SELECT id, title FROM tasks")
91/// .with_page_size(100)
92/// .build_sqlite();
93/// # Ok(())
94/// # }
95/// ```
96pub struct RdbcItemReaderBuilder<'a, I> {
97 postgres_pool: Option<Pool<Postgres>>,
98 mysql_pool: Option<Pool<MySql>>,
99 sqlite_pool: Option<Pool<Sqlite>>,
100 query_source: Option<QuerySource<'a>>,
101 page_size: Option<i32>,
102 db_type: Option<DatabaseType>,
103 keyset_column: Option<String>,
104 #[allow(clippy::type_complexity)]
105 keyset_key_fn: Option<Box<dyn Fn(&I) -> String>>,
106 _phantom: PhantomData<I>,
107}
108
109impl<'a, I> RdbcItemReaderBuilder<'a, I> {
110 /// Creates a new unified reader builder with default configuration.
111 pub fn new() -> Self {
112 Self {
113 postgres_pool: None,
114 mysql_pool: None,
115 sqlite_pool: None,
116 query_source: None,
117 page_size: None,
118 db_type: None,
119 keyset_column: None,
120 keyset_key_fn: None,
121 _phantom: PhantomData,
122 }
123 }
124
125 /// Sets the PostgreSQL connection pool and database type.
126 ///
127 /// # Arguments
128 /// * `pool` - The PostgreSQL connection pool
129 ///
130 /// # Returns
131 /// The updated builder instance configured for PostgreSQL
132 pub fn postgres(mut self, pool: Pool<Postgres>) -> Self {
133 self.postgres_pool = Some(pool);
134 self.db_type = Some(DatabaseType::Postgres);
135 self
136 }
137
138 /// Sets the MySQL connection pool and database type.
139 ///
140 /// # Arguments
141 /// * `pool` - The MySQL connection pool
142 ///
143 /// # Returns
144 /// The updated builder instance configured for MySQL
145 pub fn mysql(mut self, pool: Pool<MySql>) -> Self {
146 self.mysql_pool = Some(pool);
147 self.db_type = Some(DatabaseType::MySql);
148 self
149 }
150
151 /// Sets the SQLite connection pool and database type.
152 ///
153 /// # Arguments
154 /// * `pool` - The SQLite connection pool
155 ///
156 /// # Returns
157 /// The updated builder instance configured for SQLite
158 pub fn sqlite(mut self, pool: Pool<Sqlite>) -> Self {
159 self.sqlite_pool = Some(pool);
160 self.db_type = Some(DatabaseType::Sqlite);
161 self
162 }
163
164 /// Sets the SQL query for the reader.
165 ///
166 /// The query should not include LIMIT/OFFSET clauses as these are handled
167 /// automatically when page_size is configured.
168 ///
169 /// # Arguments
170 /// * `query` - The SQL query to execute
171 ///
172 /// # Returns
173 /// The updated builder instance
174 pub fn query(mut self, query: &'a str) -> Self {
175 self.query_source = Some(QuerySource::Raw(query));
176 self
177 }
178
179 /// Configures the reader query using a [`SelectBuilder`].
180 ///
181 /// This is an ergonomic alternative to [`Self::query`] that lets you build the
182 /// SQL statement through a fluent API instead of writing raw SQL. If the
183 /// [`SelectBuilder`] was configured with [`SelectBuilder::order_by_keyset`], the keyset
184 /// column and key function are automatically propagated to the reader.
185 ///
186 /// Calling `.select()` after `.query()` (or vice-versa) is allowed; the
187 /// **last** call wins.
188 ///
189 /// # Arguments
190 ///
191 /// * `builder` - A [`SelectBuilder`] instance ready to be compiled.
192 ///
193 /// # Examples
194 ///
195 /// ```no_run
196 /// use spring_batch_rs::item::rdbc::{RdbcItemReaderBuilder, SelectBuilder};
197 /// use sqlx::SqlitePool;
198 /// # use serde::Deserialize;
199 /// # #[derive(sqlx::FromRow, Clone, Deserialize)]
200 /// # struct Task { id: i32, title: String }
201 ///
202 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
203 /// let pool = SqlitePool::connect("sqlite::memory:").await?;
204 ///
205 /// let reader = RdbcItemReaderBuilder::<Task>::new()
206 /// .sqlite(pool)
207 /// .select(
208 /// SelectBuilder::from("tasks")
209 /// .columns(&["id", "title"])
210 /// .where_eq("done", false)
211 /// .order_by_asc("id"),
212 /// )
213 /// .with_page_size(50)
214 /// .build_sqlite();
215 /// # Ok(())
216 /// # }
217 /// ```
218 pub fn select(mut self, builder: SelectBuilder<I>) -> Self {
219 let sql = if builder.keyset_column.is_some() {
220 builder.build_sql_no_order()
221 } else {
222 builder.build_sql()
223 };
224 if let Some(col) = builder.keyset_column {
225 self.keyset_column = Some(col);
226 }
227 if let Some(key_fn) = builder.keyset_key_fn {
228 self.keyset_key_fn = Some(key_fn);
229 }
230 self.query_source = Some(QuerySource::Built(sql));
231 self
232 }
233
234 /// Sets the page size for paginated reading.
235 ///
236 /// When set, the reader will fetch data in chunks of this size to manage
237 /// memory usage efficiently.
238 ///
239 /// # Arguments
240 /// * `page_size` - Number of items to read per page
241 ///
242 /// # Returns
243 /// The updated builder instance
244 pub fn with_page_size(mut self, page_size: i32) -> Self {
245 self.page_size = Some(page_size);
246 self
247 }
248
249 /// Enables keyset (cursor) pagination instead of LIMIT/OFFSET.
250 ///
251 /// Keyset pagination is O(log n) per page regardless of dataset size, making it
252 /// significantly faster than LIMIT/OFFSET for large tables.
253 ///
254 /// # Requirements
255 ///
256 /// - The query must **not** include `WHERE`, `ORDER BY`, or `LIMIT` clauses — the
257 /// framework appends them automatically.
258 /// - The keyset column must be indexed and have unique, sortable values (e.g.
259 /// primary key, UUID, zero-padded string ID).
260 /// - [`Self::with_page_size`] must also be set.
261 ///
262 /// # Arguments
263 ///
264 /// * `column` - Column name used as the cursor (appended to `WHERE` and `ORDER BY`).
265 /// * `key_fn` - Closure that extracts the cursor value from an item as a `String`.
266 ///
267 /// # Examples
268 ///
269 /// ```no_run
270 /// use spring_batch_rs::item::rdbc::RdbcItemReaderBuilder;
271 /// use sqlx::PgPool;
272 /// # use serde::Deserialize;
273 /// # #[derive(sqlx::FromRow, Clone, Deserialize)]
274 /// # struct Order { order_id: String, amount: f64 }
275 ///
276 /// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
277 /// let pool = PgPool::connect("postgresql://user:pass@localhost/db").await?;
278 ///
279 /// let reader = RdbcItemReaderBuilder::<Order>::new()
280 /// .postgres(pool)
281 /// .query("SELECT order_id, amount FROM orders")
282 /// .with_page_size(1_000)
283 /// .with_keyset("order_id", |o: &Order| o.order_id.clone())
284 /// .build_postgres();
285 /// # Ok(())
286 /// # }
287 /// ```
288 pub fn with_keyset(mut self, column: &str, key_fn: impl Fn(&I) -> String + 'static) -> Self {
289 self.keyset_column = Some(column.to_string());
290 self.keyset_key_fn = Some(Box::new(key_fn));
291 self
292 }
293}
294
295impl<'a, I> RdbcItemReaderBuilder<'a, I>
296where
297 for<'r> I: FromRow<'r, PgRow> + Send + Unpin + Clone,
298{
299 /// Builds a PostgreSQL reader.
300 ///
301 /// # Returns
302 /// A configured `PostgresRdbcItemReader` instance
303 ///
304 /// # Panics
305 /// Panics if PostgreSQL pool or query are missing
306 pub fn build_postgres(self) -> PostgresRdbcItemReader<I> {
307 let query = match self
308 .query_source
309 .expect("Query is required — call .query() or .select()")
310 {
311 QuerySource::Raw(s) => s.to_string(),
312 QuerySource::Built(s) => s,
313 };
314 PostgresRdbcItemReader::new(
315 self.postgres_pool.expect("PostgreSQL pool is required"),
316 query,
317 self.page_size,
318 self.keyset_column,
319 self.keyset_key_fn,
320 )
321 }
322}
323
324impl<'a, I> RdbcItemReaderBuilder<'a, I>
325where
326 for<'r> I: FromRow<'r, MySqlRow> + Send + Unpin + Clone,
327{
328 /// Builds a MySQL reader.
329 ///
330 /// # Returns
331 /// A configured `MySqlRdbcItemReader` instance
332 ///
333 /// # Panics
334 /// Panics if MySQL pool or query are missing
335 pub fn build_mysql(self) -> MySqlRdbcItemReader<I> {
336 let query = match self
337 .query_source
338 .expect("Query is required — call .query() or .select()")
339 {
340 QuerySource::Raw(s) => s.to_string(),
341 QuerySource::Built(s) => s,
342 };
343 MySqlRdbcItemReader::new(
344 self.mysql_pool.expect("MySQL pool is required"),
345 query,
346 self.page_size,
347 self.keyset_column,
348 self.keyset_key_fn,
349 )
350 }
351}
352
353impl<'a, I> RdbcItemReaderBuilder<'a, I>
354where
355 for<'r> I: FromRow<'r, SqliteRow> + Send + Unpin + Clone,
356{
357 /// Builds a SQLite reader.
358 ///
359 /// # Returns
360 /// A configured `SqliteRdbcItemReader` instance
361 ///
362 /// # Panics
363 /// Panics if SQLite pool or query are missing
364 pub fn build_sqlite(self) -> SqliteRdbcItemReader<I> {
365 let query = match self
366 .query_source
367 .expect("Query is required — call .query() or .select()")
368 {
369 QuerySource::Raw(s) => s.to_string(),
370 QuerySource::Built(s) => s,
371 };
372 SqliteRdbcItemReader::new(
373 self.sqlite_pool.expect("SQLite pool is required"),
374 query,
375 self.page_size,
376 self.keyset_column,
377 self.keyset_key_fn,
378 )
379 }
380}
381
382impl<'a, I> Default for RdbcItemReaderBuilder<'a, I> {
383 fn default() -> Self {
384 Self::new()
385 }
386}
387
388#[cfg(test)]
389mod tests {
390 use super::super::select_builder::SelectBuilder;
391 use super::*;
392 use sqlx::{FromRow, SqlitePool};
393
394 #[derive(Clone, FromRow)]
395 struct Dummy {
396 id: i32,
397 }
398
399 #[test]
400 fn should_create_via_default() {
401 // Default == new(), both should produce identical builders
402 let _b = RdbcItemReaderBuilder::<Dummy>::default();
403 }
404
405 #[test]
406 #[should_panic(expected = "SQLite pool is required")]
407 fn should_panic_when_building_sqlite_without_pool() {
408 let _ = RdbcItemReaderBuilder::<Dummy>::new()
409 .query("SELECT id FROM t")
410 .build_sqlite();
411 }
412
413 #[tokio::test(flavor = "multi_thread")]
414 async fn should_build_sqlite_reader_with_pool_and_query() {
415 let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
416 let reader = RdbcItemReaderBuilder::<Dummy>::new()
417 .sqlite(pool)
418 .query("SELECT 1 AS id")
419 .build_sqlite();
420 assert_eq!(reader.query, "SELECT 1 AS id");
421 assert_eq!(reader.page_size, None);
422 assert_eq!(reader.offset.get(), 0);
423 }
424
425 #[tokio::test(flavor = "multi_thread")]
426 async fn should_propagate_page_size_to_sqlite_reader() {
427 let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
428 let reader = RdbcItemReaderBuilder::<Dummy>::new()
429 .sqlite(pool)
430 .query("SELECT 1 AS id")
431 .with_page_size(25)
432 .build_sqlite();
433 assert_eq!(reader.page_size, Some(25));
434 }
435
436 #[test]
437 #[should_panic(expected = "PostgreSQL pool is required")]
438 fn should_panic_when_building_postgres_without_pool() {
439 let _ = RdbcItemReaderBuilder::<Dummy>::new()
440 .query("SELECT id FROM t")
441 .build_postgres();
442 }
443
444 #[test]
445 #[should_panic(expected = "MySQL pool is required")]
446 fn should_panic_when_building_mysql_without_pool() {
447 let _ = RdbcItemReaderBuilder::<Dummy>::new()
448 .query("SELECT id FROM t")
449 .build_mysql();
450 }
451
452 #[tokio::test(flavor = "multi_thread")]
453 async fn should_propagate_keyset_to_sqlite_reader() {
454 let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
455 let reader = RdbcItemReaderBuilder::<Dummy>::new()
456 .sqlite(pool)
457 .query("SELECT 1 AS id")
458 .with_page_size(5)
459 .with_keyset("id", |d: &Dummy| d.id.to_string())
460 .build_sqlite();
461 assert_eq!(
462 reader.keyset_column.as_deref(),
463 Some("id"),
464 "keyset column should be propagated to reader"
465 );
466 assert!(
467 reader.keyset_key.is_some(),
468 "keyset key fn should be propagated to reader"
469 );
470 assert!(
471 reader.last_cursor.borrow().is_none(),
472 "cursor starts as None"
473 );
474 }
475
476 #[tokio::test(flavor = "multi_thread")]
477 async fn should_build_sqlite_reader_from_select_builder() {
478 let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
479 let reader = RdbcItemReaderBuilder::<Dummy>::new()
480 .sqlite(pool)
481 .select(
482 SelectBuilder::from("items")
483 .columns(&["id"])
484 .where_eq("active", true)
485 .order_by_asc("id"),
486 )
487 .build_sqlite();
488 assert_eq!(
489 reader.query, "SELECT id FROM items WHERE active = true ORDER BY id ASC",
490 "select builder SQL should be stored in reader"
491 );
492 }
493
494 #[tokio::test(flavor = "multi_thread")]
495 async fn should_propagate_keyset_from_select_builder_to_sqlite_reader() {
496 let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
497 let reader = RdbcItemReaderBuilder::<Dummy>::new()
498 .sqlite(pool)
499 .select(
500 SelectBuilder::from("items").order_by_keyset("id", |d: &Dummy| d.id.to_string()),
501 )
502 .with_page_size(10)
503 .build_sqlite();
504 assert_eq!(
505 reader.keyset_column.as_deref(),
506 Some("id"),
507 "keyset column should propagate from SelectBuilder"
508 );
509 assert!(
510 reader.keyset_key.is_some(),
511 "keyset key fn should propagate from SelectBuilder"
512 );
513 assert_eq!(
514 reader.query, "SELECT * FROM items",
515 "keyset select builder must store SQL without ORDER BY to avoid double ORDER BY in read_page"
516 );
517 }
518
519 #[tokio::test(flavor = "multi_thread")]
520 async fn should_prefer_select_over_query_when_called_last() {
521 let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
522 let reader = RdbcItemReaderBuilder::<Dummy>::new()
523 .sqlite(pool)
524 .query("SELECT id FROM old_table")
525 .select(SelectBuilder::from("new_table").columns(&["id"]))
526 .build_sqlite();
527 assert_eq!(
528 reader.query, "SELECT id FROM new_table",
529 "select() called last should win"
530 );
531 }
532
533 #[tokio::test(flavor = "multi_thread")]
534 async fn should_prefer_query_over_select_when_called_last() {
535 let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
536 let reader = RdbcItemReaderBuilder::<Dummy>::new()
537 .sqlite(pool)
538 .select(SelectBuilder::from("old_table").columns(&["id"]))
539 .query("SELECT id FROM new_table")
540 .build_sqlite();
541 assert_eq!(
542 reader.query, "SELECT id FROM new_table",
543 "query() called last should win"
544 );
545 }
546}