connectorx/sources/mysql/
mod.rs

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

connectorx/sources/mysql/
mod.rs
