spring_batch_rs/item/rdbc/
mysql_reader.rs

1use std::cell::{Cell, RefCell};
2
3use sqlx::{FromRow, MySql, Pool, QueryBuilder, mysql::MySqlRow};
4
5use super::reader_common::read_item;
6use crate::BatchError;
7use crate::core::item::{ItemReader, ItemReaderResult};
8
9/// MySQL 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 MySqlRdbcItemReader<I>
18where
19    for<'r> I: FromRow<'r, MySqlRow> + Send + Unpin + Clone,
20{
21    pub(crate) pool: Pool<MySql>,
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> MySqlRdbcItemReader<I>
33where
34    for<'r> I: FromRow<'r, MySqlRow> + Send + Unpin + Clone,
35{
36    /// Creates a new `MySqlRdbcItemReader` with the specified parameters.
37    ///
38    /// Prefer [`RdbcItemReaderBuilder`](crate::item::rdbc::RdbcItemReaderBuilder) for a more
39    /// ergonomic construction API.
40    #[allow(clippy::type_complexity)]
41    pub fn new(
42        pool: Pool<MySql>,
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::<MySql>::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 MySqlRdbcItemReader<I>
98where
99    for<'r> I: FromRow<'r, MySqlRow> + 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 sqlx::MySqlPool;
117
118    #[derive(Clone)]
119    struct Dummy;
120
121    impl<'r> sqlx::FromRow<'r, sqlx::mysql::MySqlRow> for Dummy {
122        fn from_row(_row: &'r sqlx::mysql::MySqlRow) -> Result<Self, sqlx::Error> {
123            Ok(Dummy)
124        }
125    }
126
127    fn reader_with_keyset(keyset: bool) -> MySqlRdbcItemReader<Dummy> {
128        let pool = MySqlPool::connect_lazy("mysql://root:root@localhost/test")
129            .expect("lazy pool creation should not fail");
130        let (col, key): (Option<String>, Option<Box<dyn Fn(&Dummy) -> String>>) = if keyset {
131            (
132                Some("id".to_string()),
133                Some(Box::new(|_: &Dummy| "0".to_string())),
134            )
135        } else {
136            (None, None)
137        };
138        MySqlRdbcItemReader::new(pool, "SELECT 1".to_string(), Some(10), col, key)
139    }
140
141    #[tokio::test(flavor = "multi_thread")]
142    async fn should_initialize_without_keyset() {
143        let reader = reader_with_keyset(false);
144        assert!(reader.keyset_column.is_none(), "no keyset column expected");
145        assert!(reader.keyset_key.is_none(), "no keyset key fn expected");
146        assert!(
147            reader.last_cursor.borrow().is_none(),
148            "cursor must start as None"
149        );
150        assert_eq!(reader.offset.get(), 0, "initial offset should be 0");
151        assert!(
152            reader.buffer.borrow().is_empty(),
153            "buffer should start empty"
154        );
155        assert_eq!(reader.page_size, Some(10));
156    }
157
158    #[tokio::test(flavor = "multi_thread")]
159    async fn should_initialize_with_keyset_column_and_none_cursor() {
160        let reader = reader_with_keyset(true);
161        assert_eq!(
162            reader.keyset_column.as_deref(),
163            Some("id"),
164            "keyset column should be stored"
165        );
166        assert!(
167            reader.keyset_key.is_some(),
168            "keyset key fn should be stored"
169        );
170        assert!(
171            reader.last_cursor.borrow().is_none(),
172            "cursor must start as None before first read"
173        );
174    }
175}
spring_batch_rs/item/rdbc/mysql_reader.rs

spring_batch_rs/item/rdbc/
mysql_reader.rs
