connectorx/sources/mysql/
mod.rs

1//! Source implementation for MySQL database.
2
3mod errors;
4mod typesystem;
5
6pub use self::errors::MySQLSourceError;
7use crate::constants::DB_BUFFER_SIZE;
8use crate::{
9    data_order::DataOrder,
10    errors::ConnectorXError,
11    sources::{PartitionParser, Produce, Source, SourcePartition},
12    sql::{count_query, limit0_query, CXQuery},
13};
14use anyhow::anyhow;
15use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
16use fehler::{throw, throws};
17use log::{debug, warn};
18use r2d2::{Pool, PooledConnection};
19use r2d2_mysql::{
20    mysql::{prelude::Queryable, Binary, Opts, OptsBuilder, QueryResult, Row, Text},
21    MySqlConnectionManager,
22};
23use rust_decimal::Decimal;
24use serde_json::Value;
25use sqlparser::dialect::MySqlDialect;
26use std::marker::PhantomData;
27pub use typesystem::MySQLTypeSystem;
28
29type MysqlConn = PooledConnection<MySqlConnectionManager>;
30
31pub enum BinaryProtocol {}
32pub enum TextProtocol {}
33
34#[throws(MySQLSourceError)]
35fn get_total_rows(conn: &mut MysqlConn, query: &CXQuery<String>) -> usize {
36    conn.query_first(&count_query(query, &MySqlDialect {})?)?
37        .ok_or_else(|| anyhow!("mysql failed to get the count of query: {}", query))?
38}
39
40pub struct MySQLSource<P> {
41    pool: Pool<MySqlConnectionManager>,
42    origin_query: Option<String>,
43    queries: Vec<CXQuery<String>>,
44    names: Vec<String>,
45    schema: Vec<MySQLTypeSystem>,
46    pre_execution_queries: Option<Vec<String>>,
47    _protocol: PhantomData<P>,
48}
49
50impl<P> MySQLSource<P> {
51    #[throws(MySQLSourceError)]
52    pub fn new(conn: &str, nconn: usize) -> Self {
53        let manager = MySqlConnectionManager::new(OptsBuilder::from_opts(Opts::from_url(conn)?));
54        let pool = r2d2::Pool::builder()
55            .max_size(nconn as u32)
56            .build(manager)?;
57
58        Self {
59            pool,
60            origin_query: None,
61            queries: vec![],
62            names: vec![],
63            schema: vec![],
64            pre_execution_queries: None,
65            _protocol: PhantomData,
66        }
67    }
68}
69
70impl<P> Source for MySQLSource<P>
71where
72    MySQLSourcePartition<P>:
73        SourcePartition<TypeSystem = MySQLTypeSystem, Error = MySQLSourceError>,
74    P: Send,
75{
76    const DATA_ORDERS: &'static [DataOrder] = &[DataOrder::RowMajor];
77    type Partition = MySQLSourcePartition<P>;
78    type TypeSystem = MySQLTypeSystem;
79    type Error = MySQLSourceError;
80
81    #[throws(MySQLSourceError)]
82    fn set_data_order(&mut self, data_order: DataOrder) {
83        if !matches!(data_order, DataOrder::RowMajor) {
84            throw!(ConnectorXError::UnsupportedDataOrder(data_order));
85        }
86    }
87
88    fn set_queries<Q: ToString>(&mut self, queries: &[CXQuery<Q>]) {
89        self.queries = queries.iter().map(|q| q.map(Q::to_string)).collect();
90    }
91
92    fn set_origin_query(&mut self, query: Option<String>) {
93        self.origin_query = query;
94    }
95
96    fn set_pre_execution_queries(&mut self, pre_execution_queries: Option<&[String]>) {
97        self.pre_execution_queries = pre_execution_queries.map(|s| s.to_vec());
98    }
99
100    #[throws(MySQLSourceError)]
101    fn fetch_metadata(&mut self) {
102        assert!(!self.queries.is_empty());
103
104        let mut conn = self.pool.get()?;
105        let first_query = &self.queries[0];
106
107        match conn.prep(first_query) {
108            Ok(stmt) => {
109                let (names, types) = stmt
110                    .columns()
111                    .iter()
112                    .map(|col| {
113                        let col_name = col.name_str().to_string();
114                        let d = MySQLTypeSystem::from((
115                            &col.column_type(),
116                            &col.flags(),
117                            col.character_set(),
118                        ));
119                        (col_name, d)
120                    })
121                    .unzip();
122                self.names = names;
123                self.schema = types;
124            }
125            Err(e) => {
126                warn!(
127                    "mysql text prepared statement error: {:?}, switch to limit1 method",
128                    e
129                );
130                for (i, query) in self.queries.iter().enumerate() {
131                    // assuming all the partition queries yield same schema
132                    match conn
133                        .query_first::<Row, _>(limit0_query(query, &MySqlDialect {})?.as_str())
134                    {
135                        Ok(Some(row)) => {
136                            let (names, types) = row
137                                .columns_ref()
138                                .iter()
139                                .map(|col| {
140                                    (
141                                        col.name_str().to_string(),
142                                        MySQLTypeSystem::from((
143                                            &col.column_type(),
144                                            &col.flags(),
145                                            col.character_set(),
146                                        )),
147                                    )
148                                })
149                                .unzip();
150                            self.names = names;
151                            self.schema = types;
152                            return;
153                        }
154                        Ok(None) => {}
155                        Err(e) if i == self.queries.len() - 1 => {
156                            // tried the last query but still get an error
157                            debug!("cannot get metadata for '{}', try next query: {}", query, e);
158                            throw!(e)
159                        }
160                        Err(_) => {}
161                    }
162                }
163
164                // tried all queries but all get empty result set
165                let iter = conn.query_iter(self.queries[0].as_str())?;
166                let (names, types) = iter
167                    .columns()
168                    .as_ref()
169                    .iter()
170                    .map(|col| {
171                        (
172                            col.name_str().to_string(),
173                            MySQLTypeSystem::VarChar(false), // set all columns as string (align with pandas)
174                        )
175                    })
176                    .unzip();
177                self.names = names;
178                self.schema = types;
179            }
180        }
181    }
182
183    #[throws(MySQLSourceError)]
184    fn result_rows(&mut self) -> Option<usize> {
185        match &self.origin_query {
186            Some(q) => {
187                let cxq = CXQuery::Naked(q.clone());
188                let mut conn = self.pool.get()?;
189                let nrows = get_total_rows(&mut conn, &cxq)?;
190                Some(nrows)
191            }
192            None => None,
193        }
194    }
195
196    fn names(&self) -> Vec<String> {
197        self.names.clone()
198    }
199
200    fn schema(&self) -> Vec<Self::TypeSystem> {
201        self.schema.clone()
202    }
203
204    #[throws(MySQLSourceError)]
205    fn partition(self) -> Vec<Self::Partition> {
206        let mut ret = vec![];
207        for query in self.queries {
208            let mut conn = self.pool.get()?;
209
210            if let Some(pre_queries) = &self.pre_execution_queries {
211                for pre_query in pre_queries {
212                    conn.query_drop(pre_query)?;
213                }
214            }
215
216            ret.push(MySQLSourcePartition::new(conn, &query, &self.schema));
217        }
218        ret
219    }
220}
221
222pub struct MySQLSourcePartition<P> {
223    conn: MysqlConn,
224    query: CXQuery<String>,
225    schema: Vec<MySQLTypeSystem>,
226    nrows: usize,
227    ncols: usize,
228    _protocol: PhantomData<P>,
229}
230
231impl<P> MySQLSourcePartition<P> {
232    pub fn new(conn: MysqlConn, query: &CXQuery<String>, schema: &[MySQLTypeSystem]) -> Self {
233        Self {
234            conn,
235            query: query.clone(),
236            schema: schema.to_vec(),
237            nrows: 0,
238            ncols: schema.len(),
239            _protocol: PhantomData,
240        }
241    }
242}
243
244impl SourcePartition for MySQLSourcePartition<BinaryProtocol> {
245    type TypeSystem = MySQLTypeSystem;
246    type Parser<'a> = MySQLBinarySourceParser<'a>;
247    type Error = MySQLSourceError;
248
249    #[throws(MySQLSourceError)]
250    fn result_rows(&mut self) {
251        self.nrows = get_total_rows(&mut self.conn, &self.query)?;
252    }
253
254    #[throws(MySQLSourceError)]
255    fn parser(&mut self) -> Self::Parser<'_> {
256        let stmt = self.conn.prep(self.query.as_str())?;
257        let iter = self.conn.exec_iter(stmt, ())?;
258        MySQLBinarySourceParser::new(iter, &self.schema)
259    }
260
261    fn nrows(&self) -> usize {
262        self.nrows
263    }
264
265    fn ncols(&self) -> usize {
266        self.ncols
267    }
268}
269
270impl SourcePartition for MySQLSourcePartition<TextProtocol> {
271    type TypeSystem = MySQLTypeSystem;
272    type Parser<'a> = MySQLTextSourceParser<'a>;
273    type Error = MySQLSourceError;
274
275    #[throws(MySQLSourceError)]
276    fn result_rows(&mut self) {
277        self.nrows = get_total_rows(&mut self.conn, &self.query)?;
278    }
279
280    #[throws(MySQLSourceError)]
281    fn parser(&mut self) -> Self::Parser<'_> {
282        let query = self.query.clone();
283        let iter = self.conn.query_iter(query)?;
284        MySQLTextSourceParser::new(iter, &self.schema)
285    }
286
287    fn nrows(&self) -> usize {
288        self.nrows
289    }
290
291    fn ncols(&self) -> usize {
292        self.ncols
293    }
294}
295
296pub struct MySQLBinarySourceParser<'a> {
297    iter: QueryResult<'a, 'a, 'a, Binary>,
298    rowbuf: Vec<Row>,
299    ncols: usize,
300    current_col: usize,
301    current_row: usize,
302    is_finished: bool,
303}
304
305impl<'a> MySQLBinarySourceParser<'a> {
306    pub fn new(iter: QueryResult<'a, 'a, 'a, Binary>, schema: &[MySQLTypeSystem]) -> Self {
307        Self {
308            iter,
309            rowbuf: Vec::with_capacity(DB_BUFFER_SIZE),
310            ncols: schema.len(),
311            current_row: 0,
312            current_col: 0,
313            is_finished: false,
314        }
315    }
316
317    #[throws(MySQLSourceError)]
318    fn next_loc(&mut self) -> (usize, usize) {
319        let ret = (self.current_row, self.current_col);
320        self.current_row += (self.current_col + 1) / self.ncols;
321        self.current_col = (self.current_col + 1) % self.ncols;
322        ret
323    }
324}
325
326impl<'a> PartitionParser<'a> for MySQLBinarySourceParser<'a> {
327    type TypeSystem = MySQLTypeSystem;
328    type Error = MySQLSourceError;
329
330    #[throws(MySQLSourceError)]
331    fn fetch_next(&mut self) -> (usize, bool) {
332        assert!(self.current_col == 0);
333        let remaining_rows = self.rowbuf.len() - self.current_row;
334        if remaining_rows > 0 {
335            return (remaining_rows, self.is_finished);
336        } else if self.is_finished {
337            return (0, self.is_finished);
338        }
339
340        if !self.rowbuf.is_empty() {
341            self.rowbuf.drain(..);
342        }
343
344        for _ in 0..DB_BUFFER_SIZE {
345            if let Some(item) = self.iter.next() {
346                self.rowbuf.push(item?);
347            } else {
348                self.is_finished = true;
349                break;
350            }
351        }
352        self.current_row = 0;
353        self.current_col = 0;
354
355        (self.rowbuf.len(), self.is_finished)
356    }
357}
358
359macro_rules! impl_produce_binary {
360    ($($t: ty,)+) => {
361        $(
362            impl<'r, 'a> Produce<'r, $t> for MySQLBinarySourceParser<'a> {
363                type Error = MySQLSourceError;
364
365                #[throws(MySQLSourceError)]
366                fn produce(&'r mut self) -> $t {
367                    let (ridx, cidx) = self.next_loc()?;
368                    let res = self.rowbuf[ridx].take(cidx).ok_or_else(|| anyhow!("mysql cannot parse at position: ({}, {})", ridx, cidx))?;
369                    res
370                }
371            }
372
373            impl<'r, 'a> Produce<'r, Option<$t>> for MySQLBinarySourceParser<'a> {
374                type Error = MySQLSourceError;
375
376                #[throws(MySQLSourceError)]
377                fn produce(&'r mut self) -> Option<$t> {
378                    let (ridx, cidx) = self.next_loc()?;
379                    let res = self.rowbuf[ridx].take(cidx).ok_or_else(|| anyhow!("mysql cannot parse at position: ({}, {})", ridx, cidx))?;
380                    res
381                }
382            }
383        )+
384    };
385}
386
387impl_produce_binary!(
388    i8,
389    i16,
390    i32,
391    i64,
392    u8,
393    u16,
394    u32,
395    u64,
396    f32,
397    f64,
398    NaiveDate,
399    NaiveTime,
400    NaiveDateTime,
401    Decimal,
402    String,
403    Vec<u8>,
404    Value,
405);
406
407pub struct MySQLTextSourceParser<'a> {
408    iter: QueryResult<'a, 'a, 'a, Text>,
409    rowbuf: Vec<Row>,
410    ncols: usize,
411    current_col: usize,
412    current_row: usize,
413    is_finished: bool,
414}
415
416impl<'a> MySQLTextSourceParser<'a> {
417    pub fn new(iter: QueryResult<'a, 'a, 'a, Text>, schema: &[MySQLTypeSystem]) -> Self {
418        Self {
419            iter,
420            rowbuf: Vec::with_capacity(DB_BUFFER_SIZE),
421            ncols: schema.len(),
422            current_row: 0,
423            current_col: 0,
424            is_finished: false,
425        }
426    }
427
428    #[throws(MySQLSourceError)]
429    fn next_loc(&mut self) -> (usize, usize) {
430        let ret = (self.current_row, self.current_col);
431        self.current_row += (self.current_col + 1) / self.ncols;
432        self.current_col = (self.current_col + 1) % self.ncols;
433        ret
434    }
435}
436
437impl<'a> PartitionParser<'a> for MySQLTextSourceParser<'a> {
438    type TypeSystem = MySQLTypeSystem;
439    type Error = MySQLSourceError;
440
441    #[throws(MySQLSourceError)]
442    fn fetch_next(&mut self) -> (usize, bool) {
443        assert!(self.current_col == 0);
444        let remaining_rows = self.rowbuf.len() - self.current_row;
445        if remaining_rows > 0 {
446            return (remaining_rows, self.is_finished);
447        } else if self.is_finished {
448            return (0, self.is_finished);
449        }
450
451        if !self.rowbuf.is_empty() {
452            self.rowbuf.drain(..);
453        }
454        for _ in 0..DB_BUFFER_SIZE {
455            if let Some(item) = self.iter.next() {
456                self.rowbuf.push(item?);
457            } else {
458                self.is_finished = true;
459                break;
460            }
461        }
462        self.current_row = 0;
463        self.current_col = 0;
464        (self.rowbuf.len(), self.is_finished)
465    }
466}
467
468macro_rules! impl_produce_text {
469    ($($t: ty,)+) => {
470        $(
471            impl<'r, 'a> Produce<'r, $t> for MySQLTextSourceParser<'a> {
472                type Error = MySQLSourceError;
473
474                #[throws(MySQLSourceError)]
475                fn produce(&'r mut self) -> $t {
476                    let (ridx, cidx) = self.next_loc()?;
477                    let res = self.rowbuf[ridx].take(cidx).ok_or_else(|| anyhow!("mysql cannot parse at position: ({}, {})", ridx, cidx))?;
478                    res
479                }
480            }
481
482            impl<'r, 'a> Produce<'r, Option<$t>> for MySQLTextSourceParser<'a> {
483                type Error = MySQLSourceError;
484
485                #[throws(MySQLSourceError)]
486                fn produce(&'r mut self) -> Option<$t> {
487                    let (ridx, cidx) = self.next_loc()?;
488                    let res = self.rowbuf[ridx].take(cidx).ok_or_else(|| anyhow!("mysql cannot parse at position: ({}, {})", ridx, cidx))?;
489                    res
490                }
491            }
492        )+
493    };
494}
495
496impl_produce_text!(
497    i8,
498    i16,
499    i32,
500    i64,
501    u8,
502    u16,
503    u32,
504    u64,
505    f32,
506    f64,
507    NaiveDate,
508    NaiveTime,
509    NaiveDateTime,
510    Decimal,
511    String,
512    Vec<u8>,
513    Value,
514);
connectorx/sources/mysql/mod.rs

connectorx/sources/mysql/
mod.rs
