1use crate::connection::{RemoteDbType, big_decimal_to_i128, just_return, projections_contains};
2use crate::{
3 Connection, ConnectionOptions, DFResult, MysqlType, Pool, RemoteField, RemoteSchema,
4 RemoteSchemaRef, RemoteType,
5};
6use async_stream::stream;
7use bigdecimal::{BigDecimal, num_bigint};
8use chrono::Timelike;
9use datafusion::arrow::array::{
10 ArrayRef, BinaryBuilder, Date32Builder, Decimal128Builder, Decimal256Builder, Float32Builder,
11 Float64Builder, Int8Builder, Int16Builder, Int32Builder, Int64Builder, LargeBinaryBuilder,
12 LargeStringBuilder, RecordBatch, StringBuilder, Time32SecondBuilder, Time64NanosecondBuilder,
13 TimestampMicrosecondBuilder, UInt8Builder, UInt16Builder, UInt32Builder, UInt64Builder,
14 make_builder,
15};
16use datafusion::arrow::datatypes::{DataType, Date32Type, SchemaRef, TimeUnit, i256};
17use datafusion::common::{DataFusionError, project_schema};
18use datafusion::execution::SendableRecordBatchStream;
19use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
20use derive_getters::Getters;
21use derive_with::With;
22use futures::StreamExt;
23use futures::lock::Mutex;
24use mysql_async::consts::{ColumnFlags, ColumnType};
25use mysql_async::prelude::Queryable;
26use mysql_async::{Column, FromValueError, Row, Value};
27use std::sync::Arc;
28
29#[derive(Debug, Clone, With, Getters)]
30pub struct MysqlConnectionOptions {
31 pub(crate) host: String,
32 pub(crate) port: u16,
33 pub(crate) username: String,
34 pub(crate) password: String,
35 pub(crate) database: Option<String>,
36 pub(crate) pool_max_size: usize,
37 pub(crate) stream_chunk_size: usize,
38}
39
40impl MysqlConnectionOptions {
41 pub fn new(
42 host: impl Into<String>,
43 port: u16,
44 username: impl Into<String>,
45 password: impl Into<String>,
46 ) -> Self {
47 Self {
48 host: host.into(),
49 port,
50 username: username.into(),
51 password: password.into(),
52 database: None,
53 pool_max_size: 10,
54 stream_chunk_size: 2048,
55 }
56 }
57}
58
59#[derive(Debug)]
60pub struct MysqlPool {
61 pool: mysql_async::Pool,
62}
63
64pub(crate) fn connect_mysql(options: &MysqlConnectionOptions) -> DFResult<MysqlPool> {
65 let pool_opts = mysql_async::PoolOpts::new().with_constraints(
66 mysql_async::PoolConstraints::new(0, options.pool_max_size)
67 .expect("Failed to create pool constraints"),
68 );
69 let opts_builder = mysql_async::OptsBuilder::default()
70 .ip_or_hostname(options.host.clone())
71 .tcp_port(options.port)
72 .user(Some(options.username.clone()))
73 .pass(Some(options.password.clone()))
74 .db_name(options.database.clone())
75 .init(vec!["set time_zone='+00:00'".to_string()])
76 .pool_opts(pool_opts);
77 let pool = mysql_async::Pool::new(opts_builder);
78 Ok(MysqlPool { pool })
79}
80
81#[async_trait::async_trait]
82impl Pool for MysqlPool {
83 async fn get(&self) -> DFResult<Arc<dyn Connection>> {
84 let conn = self.pool.get_conn().await.map_err(|e| {
85 DataFusionError::Execution(format!("Failed to get mysql connection from pool: {:?}", e))
86 })?;
87 Ok(Arc::new(MysqlConnection {
88 conn: Arc::new(Mutex::new(conn)),
89 }))
90 }
91}
92
93#[derive(Debug)]
94pub struct MysqlConnection {
95 conn: Arc<Mutex<mysql_async::Conn>>,
96}
97
98#[async_trait::async_trait]
99impl Connection for MysqlConnection {
100 async fn infer_schema(&self, sql: &str) -> DFResult<RemoteSchemaRef> {
101 let sql = RemoteDbType::Mysql.query_limit_1(sql)?;
102 let mut conn = self.conn.lock().await;
103 let conn = &mut *conn;
104 let row: Option<Row> = conn.query_first(&sql).await.map_err(|e| {
105 DataFusionError::Execution(format!("Failed to execute query {sql} on mysql: {e:?}",))
106 })?;
107 let Some(row) = row else {
108 return Err(DataFusionError::Execution(
109 "No rows returned to infer schema".to_string(),
110 ));
111 };
112 let remote_schema = Arc::new(build_remote_schema(&row)?);
113 Ok(remote_schema)
114 }
115
116 async fn query(
117 &self,
118 conn_options: &ConnectionOptions,
119 sql: &str,
120 table_schema: SchemaRef,
121 projection: Option<&Vec<usize>>,
122 unparsed_filters: &[String],
123 limit: Option<usize>,
124 ) -> DFResult<SendableRecordBatchStream> {
125 let projected_schema = project_schema(&table_schema, projection)?;
126 let sql = RemoteDbType::Mysql.try_rewrite_query(sql, unparsed_filters, limit)?;
127 let projection = projection.cloned();
128 let chunk_size = conn_options.stream_chunk_size();
129 let conn = Arc::clone(&self.conn);
130 let stream = Box::pin(stream! {
131 let mut conn = conn.lock().await;
132 let mut query_iter = conn
133 .query_iter(sql.clone())
134 .await
135 .map_err(|e| {
136 DataFusionError::Execution(format!("Failed to execute query {sql} on mysql: {e:?}"))
137 })?;
138
139 let Some(stream) = query_iter.stream::<Row>().await.map_err(|e| {
140 DataFusionError::Execution(format!("Failed to get stream from mysql: {e:?}"))
141 })? else {
142 yield Err(DataFusionError::Execution("Get none stream from mysql".to_string()));
143 return;
144 };
145
146 let mut chunked_stream = stream.chunks(chunk_size).boxed();
147
148 while let Some(chunk) = chunked_stream.next().await {
149 let rows = chunk
150 .into_iter()
151 .collect::<Result<Vec<_>, _>>()
152 .map_err(|e| {
153 DataFusionError::Execution(format!(
154 "Failed to collect rows from mysql due to {e}",
155 ))
156 })?;
157
158 yield Ok::<_, DataFusionError>(rows)
159 }
160 });
161
162 let stream = stream.map(move |rows| {
163 let rows = rows?;
164 rows_to_batch(rows.as_slice(), &table_schema, projection.as_ref())
165 });
166
167 Ok(Box::pin(RecordBatchStreamAdapter::new(
168 projected_schema,
169 stream,
170 )))
171 }
172}
173
174fn mysql_type_to_remote_type(mysql_col: &Column) -> DFResult<MysqlType> {
175 let character_set = mysql_col.character_set();
176 let is_utf8_bin_character_set = character_set == 45;
177 let is_binary = mysql_col.flags().contains(ColumnFlags::BINARY_FLAG);
178 let is_blob = mysql_col.flags().contains(ColumnFlags::BLOB_FLAG);
179 let is_unsigned = mysql_col.flags().contains(ColumnFlags::UNSIGNED_FLAG);
180 let col_length = mysql_col.column_length();
181 match mysql_col.column_type() {
182 ColumnType::MYSQL_TYPE_TINY => {
183 if is_unsigned {
184 Ok(MysqlType::TinyIntUnsigned)
185 } else {
186 Ok(MysqlType::TinyInt)
187 }
188 }
189 ColumnType::MYSQL_TYPE_SHORT => {
190 if is_unsigned {
191 Ok(MysqlType::SmallIntUnsigned)
192 } else {
193 Ok(MysqlType::SmallInt)
194 }
195 }
196 ColumnType::MYSQL_TYPE_INT24 => {
197 if is_unsigned {
198 Ok(MysqlType::MediumIntUnsigned)
199 } else {
200 Ok(MysqlType::MediumInt)
201 }
202 }
203 ColumnType::MYSQL_TYPE_LONG => {
204 if is_unsigned {
205 Ok(MysqlType::IntegerUnsigned)
206 } else {
207 Ok(MysqlType::Integer)
208 }
209 }
210 ColumnType::MYSQL_TYPE_LONGLONG => {
211 if is_unsigned {
212 Ok(MysqlType::BigIntUnsigned)
213 } else {
214 Ok(MysqlType::BigInt)
215 }
216 }
217 ColumnType::MYSQL_TYPE_FLOAT => Ok(MysqlType::Float),
218 ColumnType::MYSQL_TYPE_DOUBLE => Ok(MysqlType::Double),
219 ColumnType::MYSQL_TYPE_NEWDECIMAL => {
220 let precision = (mysql_col.column_length() - 2) as u8;
221 let scale = mysql_col.decimals();
222 Ok(MysqlType::Decimal(precision, scale))
223 }
224 ColumnType::MYSQL_TYPE_DATE => Ok(MysqlType::Date),
225 ColumnType::MYSQL_TYPE_DATETIME => Ok(MysqlType::Datetime),
226 ColumnType::MYSQL_TYPE_TIME => Ok(MysqlType::Time),
227 ColumnType::MYSQL_TYPE_TIMESTAMP => Ok(MysqlType::Timestamp),
228 ColumnType::MYSQL_TYPE_YEAR => Ok(MysqlType::Year),
229 ColumnType::MYSQL_TYPE_STRING if !is_binary => Ok(MysqlType::Char),
230 ColumnType::MYSQL_TYPE_STRING if is_binary => {
231 if is_utf8_bin_character_set {
232 Ok(MysqlType::Char)
233 } else {
234 Ok(MysqlType::Binary)
235 }
236 }
237 ColumnType::MYSQL_TYPE_VAR_STRING if !is_binary => Ok(MysqlType::Varchar),
238 ColumnType::MYSQL_TYPE_VAR_STRING if is_binary => {
239 if is_utf8_bin_character_set {
240 Ok(MysqlType::Varchar)
241 } else {
242 Ok(MysqlType::Varbinary)
243 }
244 }
245 ColumnType::MYSQL_TYPE_VARCHAR => Ok(MysqlType::Varchar),
246 ColumnType::MYSQL_TYPE_BLOB if is_blob && !is_binary => Ok(MysqlType::Text(col_length)),
247 ColumnType::MYSQL_TYPE_BLOB if is_blob && is_binary => {
248 if is_utf8_bin_character_set {
249 Ok(MysqlType::Text(col_length))
250 } else {
251 Ok(MysqlType::Blob(col_length))
252 }
253 }
254 ColumnType::MYSQL_TYPE_JSON => Ok(MysqlType::Json),
255 ColumnType::MYSQL_TYPE_GEOMETRY => Ok(MysqlType::Geometry),
256 _ => Err(DataFusionError::NotImplemented(format!(
257 "Unsupported mysql type: {mysql_col:?}",
258 ))),
259 }
260}
261
262fn build_remote_schema(row: &Row) -> DFResult<RemoteSchema> {
263 let mut remote_fields = vec![];
264 for col in row.columns_ref() {
265 remote_fields.push(RemoteField::new(
266 col.name_str().to_string(),
267 RemoteType::Mysql(mysql_type_to_remote_type(col)?),
268 true,
269 ));
270 }
271 Ok(RemoteSchema::new(remote_fields))
272}
273
274macro_rules! handle_primitive_type {
275 ($builder:expr, $field:expr, $col:expr, $builder_ty:ty, $value_ty:ty, $row:expr, $index:expr, $convert:expr) => {{
276 let builder = $builder
277 .as_any_mut()
278 .downcast_mut::<$builder_ty>()
279 .unwrap_or_else(|| {
280 panic!(
281 "Failed to downcast builder to {} for {:?} and {:?}",
282 stringify!($builder_ty),
283 $field,
284 $col
285 )
286 });
287 let v = $row.get_opt::<$value_ty, usize>($index);
288
289 match v {
290 None => builder.append_null(),
291 Some(Ok(v)) => builder.append_value($convert(v)?),
292 Some(Err(FromValueError(Value::NULL))) => builder.append_null(),
293 Some(Err(e)) => {
294 return Err(DataFusionError::Execution(format!(
295 "Failed to get optional {:?} value for {:?} and {:?}: {e:?}",
296 stringify!($value_ty),
297 $field,
298 $col,
299 )));
300 }
301 }
302 }};
303}
304
305fn rows_to_batch(
306 rows: &[Row],
307 table_schema: &SchemaRef,
308 projection: Option<&Vec<usize>>,
309) -> DFResult<RecordBatch> {
310 let projected_schema = project_schema(table_schema, projection)?;
311 let mut array_builders = vec![];
312 for field in table_schema.fields() {
313 let builder = make_builder(field.data_type(), rows.len());
314 array_builders.push(builder);
315 }
316
317 for row in rows {
318 for (idx, field) in table_schema.fields.iter().enumerate() {
319 if !projections_contains(projection, idx) {
320 continue;
321 }
322 let builder = &mut array_builders[idx];
323 let col = row.columns_ref().get(idx);
324 match field.data_type() {
325 DataType::Int8 => {
326 handle_primitive_type!(
327 builder,
328 field,
329 col,
330 Int8Builder,
331 i8,
332 row,
333 idx,
334 just_return
335 );
336 }
337 DataType::Int16 => {
338 handle_primitive_type!(
339 builder,
340 field,
341 col,
342 Int16Builder,
343 i16,
344 row,
345 idx,
346 just_return
347 );
348 }
349 DataType::Int32 => {
350 handle_primitive_type!(
351 builder,
352 field,
353 col,
354 Int32Builder,
355 i32,
356 row,
357 idx,
358 just_return
359 );
360 }
361 DataType::Int64 => {
362 handle_primitive_type!(
363 builder,
364 field,
365 col,
366 Int64Builder,
367 i64,
368 row,
369 idx,
370 just_return
371 );
372 }
373 DataType::UInt8 => {
374 handle_primitive_type!(
375 builder,
376 field,
377 col,
378 UInt8Builder,
379 u8,
380 row,
381 idx,
382 just_return
383 );
384 }
385 DataType::UInt16 => {
386 handle_primitive_type!(
387 builder,
388 field,
389 col,
390 UInt16Builder,
391 u16,
392 row,
393 idx,
394 just_return
395 );
396 }
397 DataType::UInt32 => {
398 handle_primitive_type!(
399 builder,
400 field,
401 col,
402 UInt32Builder,
403 u32,
404 row,
405 idx,
406 just_return
407 );
408 }
409 DataType::UInt64 => {
410 handle_primitive_type!(
411 builder,
412 field,
413 col,
414 UInt64Builder,
415 u64,
416 row,
417 idx,
418 just_return
419 );
420 }
421 DataType::Float32 => {
422 handle_primitive_type!(
423 builder,
424 field,
425 col,
426 Float32Builder,
427 f32,
428 row,
429 idx,
430 just_return
431 );
432 }
433 DataType::Float64 => {
434 handle_primitive_type!(
435 builder,
436 field,
437 col,
438 Float64Builder,
439 f64,
440 row,
441 idx,
442 just_return
443 );
444 }
445 DataType::Decimal128(_precision, scale) => {
446 handle_primitive_type!(
447 builder,
448 field,
449 col,
450 Decimal128Builder,
451 BigDecimal,
452 row,
453 idx,
454 |v: BigDecimal| {
455 big_decimal_to_i128(&v, Some(*scale as i32)).ok_or_else(|| {
456 DataFusionError::Execution(format!(
457 "Failed to convert BigDecimal {v:?} to i128"
458 ))
459 })
460 }
461 );
462 }
463 DataType::Decimal256(_precision, _scale) => {
464 handle_primitive_type!(
465 builder,
466 field,
467 col,
468 Decimal256Builder,
469 BigDecimal,
470 row,
471 idx,
472 |v: BigDecimal| { Ok::<_, DataFusionError>(to_decimal_256(&v)) }
473 );
474 }
475 DataType::Date32 => {
476 handle_primitive_type!(
477 builder,
478 field,
479 col,
480 Date32Builder,
481 chrono::NaiveDate,
482 row,
483 idx,
484 |v: chrono::NaiveDate| {
485 Ok::<_, DataFusionError>(Date32Type::from_naive_date(v))
486 }
487 );
488 }
489 DataType::Timestamp(TimeUnit::Microsecond, tz_opt) => {
490 match tz_opt {
491 None => {}
492 Some(tz) => {
493 if !tz.eq_ignore_ascii_case("utc") {
494 return Err(DataFusionError::NotImplemented(format!(
495 "Unsupported data type {:?} for col: {:?}",
496 field.data_type(),
497 col
498 )));
499 }
500 }
501 }
502 handle_primitive_type!(
503 builder,
504 field,
505 col,
506 TimestampMicrosecondBuilder,
507 time::PrimitiveDateTime,
508 row,
509 idx,
510 |v: time::PrimitiveDateTime| {
511 let timestamp_micros =
512 (v.assume_utc().unix_timestamp_nanos() / 1_000) as i64;
513 Ok::<_, DataFusionError>(timestamp_micros)
514 }
515 );
516 }
517 DataType::Time32(TimeUnit::Second) => {
518 handle_primitive_type!(
519 builder,
520 field,
521 col,
522 Time32SecondBuilder,
523 chrono::NaiveTime,
524 row,
525 idx,
526 |v: chrono::NaiveTime| {
527 Ok::<_, DataFusionError>(v.num_seconds_from_midnight() as i32)
528 }
529 );
530 }
531 DataType::Time64(TimeUnit::Nanosecond) => {
532 handle_primitive_type!(
533 builder,
534 field,
535 col,
536 Time64NanosecondBuilder,
537 chrono::NaiveTime,
538 row,
539 idx,
540 |v: chrono::NaiveTime| {
541 let t = i64::from(v.num_seconds_from_midnight()) * 1_000_000_000
542 + i64::from(v.nanosecond());
543 Ok::<_, DataFusionError>(t)
544 }
545 );
546 }
547 DataType::Utf8 => {
548 handle_primitive_type!(
549 builder,
550 field,
551 col,
552 StringBuilder,
553 String,
554 row,
555 idx,
556 just_return
557 );
558 }
559 DataType::LargeUtf8 => {
560 handle_primitive_type!(
561 builder,
562 field,
563 col,
564 LargeStringBuilder,
565 String,
566 row,
567 idx,
568 just_return
569 );
570 }
571 DataType::Binary => {
572 handle_primitive_type!(
573 builder,
574 field,
575 col,
576 BinaryBuilder,
577 Vec<u8>,
578 row,
579 idx,
580 just_return
581 );
582 }
583 DataType::LargeBinary => {
584 handle_primitive_type!(
585 builder,
586 field,
587 col,
588 LargeBinaryBuilder,
589 Vec<u8>,
590 row,
591 idx,
592 just_return
593 );
594 }
595 _ => {
596 return Err(DataFusionError::NotImplemented(format!(
597 "Unsupported data type {:?} for col: {:?}",
598 field.data_type(),
599 col
600 )));
601 }
602 }
603 }
604 }
605 let projected_columns = array_builders
606 .into_iter()
607 .enumerate()
608 .filter(|(idx, _)| projections_contains(projection, *idx))
609 .map(|(_, mut builder)| builder.finish())
610 .collect::<Vec<ArrayRef>>();
611 Ok(RecordBatch::try_new(projected_schema, projected_columns)?)
612}
613
614fn to_decimal_256(decimal: &BigDecimal) -> i256 {
615 let (bigint_value, _) = decimal.as_bigint_and_exponent();
616 let mut bigint_bytes = bigint_value.to_signed_bytes_le();
617
618 let is_negative = bigint_value.sign() == num_bigint::Sign::Minus;
619 let fill_byte = if is_negative { 0xFF } else { 0x00 };
620
621 if bigint_bytes.len() > 32 {
622 bigint_bytes.truncate(32);
623 } else {
624 bigint_bytes.resize(32, fill_byte);
625 };
626
627 let mut array = [0u8; 32];
628 array.copy_from_slice(&bigint_bytes);
629
630 i256::from_le_bytes(array)
631}