spring_batch_rs/item/rdbc/
sqlite_reader.rs

1use std::cell::{Cell, RefCell};
2
3use sqlx::{FromRow, Pool, QueryBuilder, Sqlite, sqlite::SqliteRow};
4
5use super::reader_common::read_item;
6use crate::BatchError;
7use crate::core::item::{ItemReader, ItemReaderResult};
8
9/// SQLite RDBC Item Reader for batch processing.
10///
11/// Supports LIMIT/OFFSET pagination (default) and keyset pagination
12/// (enabled via [`RdbcItemReaderBuilder::with_keyset`](crate::item::rdbc::RdbcItemReaderBuilder::with_keyset)).
13///
14/// # Construction
15///
16/// Prefer [`RdbcItemReaderBuilder`](crate::item::rdbc::RdbcItemReaderBuilder) for ergonomic construction.
17pub struct SqliteRdbcItemReader<I>
18where
19    for<'r> I: FromRow<'r, SqliteRow> + Send + Unpin + Clone,
20{
21    pub(crate) pool: Pool<Sqlite>,
22    pub(crate) query: String,
23    pub(crate) page_size: Option<i32>,
24    pub(crate) offset: Cell<i32>,
25    pub(crate) buffer: RefCell<Vec<I>>,
26    pub(crate) keyset_column: Option<String>,
27    #[allow(clippy::type_complexity)]
28    pub(crate) keyset_key: Option<Box<dyn Fn(&I) -> String>>,
29    pub(crate) last_cursor: RefCell<Option<String>>,
30}
31
32impl<I> SqliteRdbcItemReader<I>
33where
34    for<'r> I: FromRow<'r, SqliteRow> + Send + Unpin + Clone,
35{
36    /// Creates a new SqliteRdbcItemReader with the specified parameters
37    ///
38    /// This constructor is only accessible within the crate to enforce the use
39    /// of `RdbcItemReaderBuilder` for creating reader instances.
40    #[allow(clippy::type_complexity)]
41    pub(crate) fn new(
42        pool: Pool<Sqlite>,
43        query: String,
44        page_size: Option<i32>,
45        keyset_column: Option<String>,
46        keyset_key: Option<Box<dyn Fn(&I) -> String>>,
47    ) -> Self {
48        Self {
49            pool,
50            query,
51            page_size,
52            offset: Cell::new(0),
53            buffer: RefCell::new(vec![]),
54            keyset_column,
55            keyset_key,
56            last_cursor: RefCell::new(None),
57        }
58    }
59
60    /// Fetches the next page from the database into the internal buffer.
61    ///
62    /// # Errors
63    ///
64    /// Returns [`BatchError::ItemReader`] if the query fails.
65    fn read_page(&self) -> Result<(), BatchError> {
66        let mut query_builder = QueryBuilder::<Sqlite>::new(&self.query);
67
68        if let Some(page_size) = self.page_size {
69            if let Some(ref col) = self.keyset_column {
70                let last = self.last_cursor.borrow();
71                if let Some(ref cursor_val) = *last {
72                    let escaped = cursor_val.replace('\'', "''");
73                    query_builder.push(format!(" WHERE {} > '{}'", col, escaped));
74                }
75                drop(last);
76                query_builder.push(format!(" ORDER BY {} LIMIT {}", col, page_size));
77            } else {
78                query_builder.push(format!(" LIMIT {} OFFSET {}", page_size, self.offset.get()));
79            }
80        }
81
82        let query = query_builder.build_query_as::<I>();
83        let items = tokio::task::block_in_place(|| {
84            tokio::runtime::Handle::current().block_on(async {
85                query
86                    .fetch_all(&self.pool)
87                    .await
88                    .map_err(|e| BatchError::ItemReader(e.to_string()))
89            })
90        })?;
91
92        *self.buffer.borrow_mut() = items;
93        Ok(())
94    }
95}
96
97impl<I> ItemReader<I> for SqliteRdbcItemReader<I>
98where
99    for<'r> I: FromRow<'r, SqliteRow> + Send + Unpin + Clone,
100{
101    fn read(&self) -> ItemReaderResult<I> {
102        read_item(
103            &self.offset,
104            self.page_size,
105            &self.buffer,
106            &self.keyset_key,
107            &self.last_cursor,
108            || self.read_page(),
109        )
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use crate::core::item::ItemReader;
117    use sqlx::{FromRow, SqlitePool};
118
119    #[derive(Clone, FromRow)]
120    struct Row {
121        id: i32,
122        name: String,
123    }
124
125    async fn pool_with_rows(rows: &[(i32, &str)]) -> SqlitePool {
126        let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
127        sqlx::query("CREATE TABLE items (id INTEGER, name TEXT)")
128            .execute(&pool)
129            .await
130            .unwrap();
131        for (id, name) in rows {
132            sqlx::query("INSERT INTO items (id, name) VALUES (?, ?)")
133                .bind(id)
134                .bind(name)
135                .execute(&pool)
136                .await
137                .unwrap();
138        }
139        pool
140    }
141
142    fn make_reader(
143        pool: SqlitePool,
144        query: &str,
145        page_size: Option<i32>,
146    ) -> SqliteRdbcItemReader<Row> {
147        SqliteRdbcItemReader::<Row>::new(pool, query.to_string(), page_size, None, None)
148    }
149
150    #[tokio::test(flavor = "multi_thread")]
151    async fn should_start_with_offset_zero_and_empty_buffer() {
152        let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
153        let reader = make_reader(pool, "SELECT id, name FROM items", None);
154        assert_eq!(reader.offset.get(), 0, "initial offset should be 0");
155        assert!(
156            reader.buffer.borrow().is_empty(),
157            "initial buffer should be empty"
158        );
159        assert_eq!(reader.page_size, None);
160    }
161
162    #[tokio::test(flavor = "multi_thread")]
163    async fn should_return_none_when_table_is_empty() {
164        let pool = pool_with_rows(&[]).await;
165        let reader = make_reader(pool, "SELECT id, name FROM items", None);
166        let result = reader.read().unwrap();
167        assert!(result.is_none(), "empty table should yield None");
168    }
169
170    #[tokio::test(flavor = "multi_thread")]
171    async fn should_read_all_items_without_pagination() {
172        let pool = pool_with_rows(&[(1, "alice"), (2, "bob")]).await;
173        let reader = make_reader(pool, "SELECT id, name FROM items ORDER BY id", None);
174
175        let first = reader.read().unwrap().expect("first item should exist");
176        assert_eq!(first.name, "alice");
177
178        let second = reader.read().unwrap().expect("second item should exist");
179        assert_eq!(second.name, "bob");
180
181        assert!(
182            reader.read().unwrap().is_none(),
183            "should return None after all items"
184        );
185    }
186
187    #[tokio::test(flavor = "multi_thread")]
188    async fn should_advance_offset_on_each_read() {
189        let pool = pool_with_rows(&[(1, "x"), (2, "y")]).await;
190        let reader = make_reader(pool, "SELECT id, name FROM items ORDER BY id", None);
191
192        assert_eq!(reader.offset.get(), 0);
193        reader.read().unwrap();
194        assert_eq!(
195            reader.offset.get(),
196            1,
197            "offset should increment after each read"
198        );
199        reader.read().unwrap();
200        assert_eq!(reader.offset.get(), 2);
201    }
202
203    #[tokio::test(flavor = "multi_thread")]
204    async fn should_read_all_items_with_pagination() {
205        let pool = pool_with_rows(&[(1, "a"), (2, "b"), (3, "c"), (4, "d")]).await;
206        let reader = make_reader(pool, "SELECT id, name FROM items ORDER BY id", Some(2));
207
208        let mut count = 0;
209        while reader.read().unwrap().is_some() {
210            count += 1;
211        }
212        assert_eq!(count, 4, "should read all 4 items across 2 pages");
213    }
214
215    #[tokio::test(flavor = "multi_thread")]
216    async fn should_read_single_item() {
217        let pool = pool_with_rows(&[(42, "only")]).await;
218        let reader = make_reader(pool, "SELECT id, name FROM items", None);
219
220        let item = reader
221            .read()
222            .unwrap()
223            .expect("should return the single item");
224        assert_eq!(item.id, 42);
225        assert_eq!(item.name, "only");
226        assert!(
227            reader.read().unwrap().is_none(),
228            "should return None after the only item"
229        );
230    }
231
232    #[tokio::test(flavor = "multi_thread")]
233    async fn should_read_all_items_with_keyset_pagination() {
234        let pool = pool_with_rows(&[(1, "a"), (2, "b"), (3, "c"), (4, "d"), (5, "e")]).await;
235        let reader = SqliteRdbcItemReader::<Row>::new(
236            pool,
237            "SELECT id, name FROM items".to_string(),
238            Some(2),
239            Some("id".to_string()),
240            Some(Box::new(|r: &Row| r.id.to_string())),
241        );
242
243        let mut names = vec![];
244        while let Some(item) = reader.read().unwrap() {
245            names.push(item.name.clone());
246        }
247        assert_eq!(
248            names,
249            vec!["a", "b", "c", "d", "e"],
250            "keyset should return all items in order"
251        );
252    }
253
254    #[tokio::test(flavor = "multi_thread")]
255    async fn should_update_last_cursor_after_each_read_with_keyset() {
256        let pool = pool_with_rows(&[(10, "x"), (20, "y")]).await;
257        let reader = SqliteRdbcItemReader::<Row>::new(
258            pool,
259            "SELECT id, name FROM items".to_string(),
260            Some(2),
261            Some("id".to_string()),
262            Some(Box::new(|r: &Row| r.id.to_string())),
263        );
264
265        assert!(
266            reader.last_cursor.borrow().is_none(),
267            "cursor should be None before first read"
268        );
269        reader.read().unwrap();
270        assert_eq!(
271            reader.last_cursor.borrow().as_deref(),
272            Some("10"),
273            "cursor should be updated after first read"
274        );
275        reader.read().unwrap();
276        assert_eq!(
277            reader.last_cursor.borrow().as_deref(),
278            Some("20"),
279            "cursor should reflect last read item"
280        );
281    }
282
283    #[tokio::test(flavor = "multi_thread")]
284    async fn should_return_none_for_empty_table_with_keyset() {
285        let pool = pool_with_rows(&[]).await;
286        let reader = SqliteRdbcItemReader::<Row>::new(
287            pool,
288            "SELECT id, name FROM items".to_string(),
289            Some(2),
290            Some("id".to_string()),
291            Some(Box::new(|r: &Row| r.id.to_string())),
292        );
293        assert!(
294            reader.read().unwrap().is_none(),
295            "empty table should yield None with keyset"
296        );
297    }
298}
spring_batch_rs/item/rdbc/sqlite_reader.rs

spring_batch_rs/item/rdbc/
sqlite_reader.rs
