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, TimestampNanosecondBuilder, UInt8Builder, UInt16Builder,
14 UInt32Builder, UInt64Builder, 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 datafusion::prelude::Expr;
21use derive_getters::Getters;
22use derive_with::With;
23use futures::StreamExt;
24use futures::lock::Mutex;
25use mysql_async::consts::{ColumnFlags, ColumnType};
26use mysql_async::prelude::Queryable;
27use mysql_async::{Column, FromValueError, Row, Value};
28use std::sync::Arc;
29
30#[derive(Debug, Clone, With, Getters)]
31pub struct MysqlConnectionOptions {
32 pub(crate) host: String,
33 pub(crate) port: u16,
34 pub(crate) username: String,
35 pub(crate) password: String,
36 pub(crate) database: Option<String>,
37 pub(crate) pool_max_size: usize,
38 pub(crate) stream_chunk_size: usize,
39}
40
41impl MysqlConnectionOptions {
42 pub fn new(
43 host: impl Into<String>,
44 port: u16,
45 username: impl Into<String>,
46 password: impl Into<String>,
47 ) -> Self {
48 Self {
49 host: host.into(),
50 port,
51 username: username.into(),
52 password: password.into(),
53 database: None,
54 pool_max_size: 10,
55 stream_chunk_size: 2048,
56 }
57 }
58}
59
60#[derive(Debug)]
61pub struct MysqlPool {
62 pool: mysql_async::Pool,
63}
64
65pub(crate) fn connect_mysql(options: &MysqlConnectionOptions) -> DFResult<MysqlPool> {
66 let pool_opts = mysql_async::PoolOpts::new().with_constraints(
67 mysql_async::PoolConstraints::new(0, options.pool_max_size)
68 .expect("Failed to create pool constraints"),
69 );
70 let opts_builder = mysql_async::OptsBuilder::default()
71 .ip_or_hostname(options.host.clone())
72 .tcp_port(options.port)
73 .user(Some(options.username.clone()))
74 .pass(Some(options.password.clone()))
75 .db_name(options.database.clone())
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 filters: &[Expr],
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, 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, None) => {
490 handle_primitive_type!(
491 builder,
492 field,
493 col,
494 TimestampMicrosecondBuilder,
495 time::PrimitiveDateTime,
496 row,
497 idx,
498 |v: time::PrimitiveDateTime| {
499 let timestamp_micros =
500 (v.assume_utc().unix_timestamp_nanos() / 1_000) as i64;
501 Ok::<_, DataFusionError>(timestamp_micros)
502 }
503 );
504 }
505 DataType::Timestamp(TimeUnit::Nanosecond, None) => {
506 handle_primitive_type!(
507 builder,
508 field,
509 col,
510 TimestampNanosecondBuilder,
511 chrono::NaiveTime,
512 row,
513 idx,
514 |v: chrono::NaiveTime| {
515 let t = i64::from(v.num_seconds_from_midnight()) * 1_000_000_000
516 + i64::from(v.nanosecond());
517 Ok::<_, DataFusionError>(t)
518 }
519 );
520 }
521 DataType::Time32(TimeUnit::Second) => {
522 handle_primitive_type!(
523 builder,
524 field,
525 col,
526 Time32SecondBuilder,
527 chrono::NaiveTime,
528 row,
529 idx,
530 |v: chrono::NaiveTime| {
531 Ok::<_, DataFusionError>(v.num_seconds_from_midnight() as i32)
532 }
533 );
534 }
535 DataType::Time64(TimeUnit::Nanosecond) => {
536 handle_primitive_type!(
537 builder,
538 field,
539 col,
540 Time64NanosecondBuilder,
541 chrono::NaiveTime,
542 row,
543 idx,
544 |v: chrono::NaiveTime| {
545 let t = i64::from(v.num_seconds_from_midnight()) * 1_000_000_000
546 + i64::from(v.nanosecond());
547 Ok::<_, DataFusionError>(t)
548 }
549 );
550 }
551 DataType::Utf8 => {
552 handle_primitive_type!(
553 builder,
554 field,
555 col,
556 StringBuilder,
557 String,
558 row,
559 idx,
560 just_return
561 );
562 }
563 DataType::LargeUtf8 => {
564 handle_primitive_type!(
565 builder,
566 field,
567 col,
568 LargeStringBuilder,
569 String,
570 row,
571 idx,
572 just_return
573 );
574 }
575 DataType::Binary => {
576 handle_primitive_type!(
577 builder,
578 field,
579 col,
580 BinaryBuilder,
581 Vec<u8>,
582 row,
583 idx,
584 just_return
585 );
586 }
587 DataType::LargeBinary => {
588 handle_primitive_type!(
589 builder,
590 field,
591 col,
592 LargeBinaryBuilder,
593 Vec<u8>,
594 row,
595 idx,
596 just_return
597 );
598 }
599 _ => {
600 return Err(DataFusionError::NotImplemented(format!(
601 "Unsupported data type {:?} for col: {:?}",
602 field.data_type(),
603 col
604 )));
605 }
606 }
607 }
608 }
609 let projected_columns = array_builders
610 .into_iter()
611 .enumerate()
612 .filter(|(idx, _)| projections_contains(projection, *idx))
613 .map(|(_, mut builder)| builder.finish())
614 .collect::<Vec<ArrayRef>>();
615 Ok(RecordBatch::try_new(projected_schema, projected_columns)?)
616}
617
618fn to_decimal_256(decimal: &BigDecimal) -> i256 {
619 let (bigint_value, _) = decimal.as_bigint_and_exponent();
620 let mut bigint_bytes = bigint_value.to_signed_bytes_le();
621
622 let is_negative = bigint_value.sign() == num_bigint::Sign::Minus;
623 let fill_byte = if is_negative { 0xFF } else { 0x00 };
624
625 if bigint_bytes.len() > 32 {
626 bigint_bytes.truncate(32);
627 } else {
628 bigint_bytes.resize(32, fill_byte);
629 };
630
631 let mut array = [0u8; 32];
632 array.copy_from_slice(&bigint_bytes);
633
634 i256::from_le_bytes(array)
635}