Add MySQL write support (datafusion-contrib#134)

* Add MySQLTableWriter for insertion into MySQL

* Add test_arrow_mysql_roundtrip

* Refactor use_json_insert_for_type function to separate sqlite and mysql feature checks

* feat: add feature mysql-federation

* update datafusion-federation to 0.3.1

Author: gengteng

src/mysql.rs

@@ -0,0 +1,190 @@
+use crate::mysql::MySQL;
+use crate::util::on_conflict::OnConflict;
+use crate::util::retriable_error::check_and_mark_retriable_error;
+use crate::util::{constraints, to_datafusion_error};
+use arrow::datatypes::SchemaRef;
+use async_trait::async_trait;
+use datafusion::{
+    catalog::Session,
+    datasource::{TableProvider, TableType},
+    execution::{SendableRecordBatchStream, TaskContext},
+    logical_expr::{dml::InsertOp, Expr},
+    physical_plan::{
+        insert::{DataSink, DataSinkExec},
+        metrics::MetricsSet,
+        DisplayAs, DisplayFormatType, ExecutionPlan,
+    },
+};
+use futures::StreamExt;
+use mysql_async::TxOpts;
+use snafu::ResultExt;
+use std::any::Any;
+use std::fmt;
+use std::sync::Arc;
+
+#[derive(Debug, Clone)]
+pub struct MySQLTableWriter {
+    pub read_provider: Arc<dyn TableProvider>,
+    mysql: Arc<MySQL>,
+    on_conflict: Option<OnConflict>,
+}
+
+impl MySQLTableWriter {
+    pub fn create(
+        read_provider: Arc<dyn TableProvider>,
+        mysql: MySQL,
+        on_conflict: Option<OnConflict>,
+    ) -> Arc<Self> {
+        Arc::new(Self {
+            read_provider,
+            mysql: Arc::new(mysql),
+            on_conflict,
+        })
+    }
+
+    pub fn mysql(&self) -> Arc<MySQL> {
+        Arc::clone(&self.mysql)
+    }
+}
+
+#[async_trait]
+impl TableProvider for MySQLTableWriter {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn schema(&self) -> SchemaRef {
+        self.read_provider.schema()
+    }
+
+    fn table_type(&self) -> TableType {
+        TableType::Base
+    }
+
+    async fn scan(
+        &self,
+        state: &dyn Session,
+        projection: Option<&Vec<usize>>,
+        filters: &[Expr],
+        limit: Option<usize>,
+    ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+        self.read_provider
+            .scan(state, projection, filters, limit)
+            .await
+    }
+
+    async fn insert_into(
+        &self,
+        _state: &dyn Session,
+        input: Arc<dyn ExecutionPlan>,
+        op: InsertOp,
+    ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+        Ok(Arc::new(DataSinkExec::new(
+            input,
+            Arc::new(MySQLDataSink::new(
+                Arc::clone(&self.mysql),
+                op == InsertOp::Overwrite,
+                self.on_conflict.clone(),
+            )),
+            self.schema(),
+            None,
+        )))
+    }
+}
+
+pub struct MySQLDataSink {
+    pub mysql: Arc<MySQL>,
+    pub overwrite: bool,
+    pub on_conflict: Option<OnConflict>,
+}
+
+#[async_trait]
+impl DataSink for MySQLDataSink {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        None
+    }
+
+    async fn write_all(
+        &self,
+        mut data: SendableRecordBatchStream,
+        _context: &Arc<TaskContext>,
+    ) -> datafusion::common::Result<u64> {
+        let mut num_rows = 0u64;
+
+        let mut db_conn = self.mysql.connect().await.map_err(to_datafusion_error)?;
+        let mysql_conn = MySQL::mysql_conn(&mut db_conn).map_err(to_datafusion_error)?;
+
+        let mut conn_guard = mysql_conn.conn.lock().await;
+        let mut tx = conn_guard
+            .start_transaction(TxOpts::default())
+            .await
+            .context(super::UnableToBeginTransactionSnafu)
+            .map_err(to_datafusion_error)?;
+
+        if self.overwrite {
+            self.mysql
+                .delete_all_table_data(&mut tx)
+                .await
+                .map_err(to_datafusion_error)?;
+        }
+
+        while let Some(batch) = data.next().await {
+            let batch = batch.map_err(check_and_mark_retriable_error)?;
+            let batch_num_rows = batch.num_rows();
+
+            if batch_num_rows == 0 {
+                continue;
+            }
+
+            num_rows += batch_num_rows as u64;
+
+            constraints::validate_batch_with_constraints(
+                &[batch.clone()],
+                self.mysql.constraints(),
+            )
+            .await
+            .context(super::ConstraintViolationSnafu)
+            .map_err(to_datafusion_error)?;
+
+            self.mysql
+                .insert_batch(&mut tx, batch, self.on_conflict.clone())
+                .await
+                .map_err(to_datafusion_error)?;
+        }
+
+        tx.commit()
+            .await
+            .context(super::UnableToCommitMySQLTransactionSnafu)
+            .map_err(to_datafusion_error)?;
+
+        drop(conn_guard);
+
+        Ok(num_rows)
+    }
+}
+
+impl MySQLDataSink {
+    pub fn new(mysql: Arc<MySQL>, overwrite: bool, on_conflict: Option<OnConflict>) -> Self {
+        Self {
+            mysql,
+            overwrite,
+            on_conflict,
+        }
+    }
+}
+
+impl fmt::Debug for MySQLDataSink {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "MySQLDataSink")
+    }
+}
+
+impl DisplayAs for MySQLDataSink {
+    fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "MySQLDataSink")
+    }
+}

src/mysql/write.rs

@@ -39,6 +39,7 @@ use crate::util::{
     indexes::IndexType,
     on_conflict::{self, OnConflict},
     secrets::to_secret_map,
+    to_datafusion_error,
 };
 
 use self::write::PostgresTableWriter;
@@ -310,10 +311,6 @@ impl TableProviderFactory for PostgresTableProviderFactory {
     }
 }
 
-fn to_datafusion_error(error: Error) -> DataFusionError {
-    DataFusionError::External(Box::new(error))
-}
-
 #[derive(Debug, Clone)]
 pub struct Postgres {
     table_name: String,

src/postgres.rs

@@ -13,10 +13,10 @@ use num_bigint::BigInt;
 use sea_query::{
     Alias, ColumnDef, ColumnType, Expr, GenericBuilder, Index, InsertStatement, IntoIden,
     IntoIndexColumn, Keyword, MysqlQueryBuilder, OnConflict, PostgresQueryBuilder, Query,
-    QueryBuilder, SimpleExpr, SqliteQueryBuilder, StringLen, Table,
+    QueryBuilder, SimpleExpr, SqliteQueryBuilder, Table,
 };
 use snafu::Snafu;
-use std::{any::Any, sync::Arc};
+use std::sync::Arc;
 use time::{OffsetDateTime, PrimitiveDateTime};
 
 #[derive(Debug, Snafu)]
@@ -105,6 +105,11 @@ impl CreateTableBuilder {
     #[must_use]
     pub fn build_mysql(self) -> String {
         self.build(MysqlQueryBuilder, &|f: &Arc<Field>| -> ColumnType {
+            // MySQL does not natively support Arrays, Structs, etc
+            // so we use JSON column type for List, FixedSizeList, LargeList, Struct, etc
+            if f.data_type().is_nested() {
+                return ColumnType::JsonBinary;
+            }
             map_data_type_to_column_type(f.data_type())
         })
     }
@@ -181,11 +186,20 @@ pub fn use_json_insert_for_type<T: QueryBuilder + 'static>(
     query_builder: &T,
 ) -> bool {
     #[cfg(feature = "sqlite")]
-    if (query_builder as &dyn Any)
-        .downcast_ref::<SqliteQueryBuilder>()
-        .is_some()
     {
-        return data_type.is_nested();
+        use std::any::Any;
+        let any_builder = query_builder as &dyn Any;
+        if any_builder.is::<SqliteQueryBuilder>() {
+            return data_type.is_nested();
+        }
+    }
+    #[cfg(feature = "mysql")]
+    {
+        use std::any::Any;
+        let any_builder = query_builder as &dyn Any;
+        if any_builder.is::<MysqlQueryBuilder>() {
+            return data_type.is_nested();
+        }
     }
     false
 }
@@ -1264,7 +1278,11 @@ pub(crate) fn map_data_type_to_column_type(data_type: &DataType) -> ColumnType {
         | DataType::FixedSizeList(list_type, _) => {
             ColumnType::Array(map_data_type_to_column_type(list_type.data_type()).into())
         }
-        DataType::Binary | DataType::LargeBinary => ColumnType::VarBinary(StringLen::Max),
+        // Originally mapped to VarBinary type, corresponding to MySQL's varbinary, which has a maximum length of 65535.
+        // This caused the error: "Row size too large. The maximum row size for the used table type, not counting BLOBs, is 65535.
+        // This includes storage overhead, check the manual. You have to change some columns to TEXT or BLOBs."
+        // Changing to Blob fixes this issue. This change does not affect Postgres, and for Sqlite, the mapping type changes from varbinary_blob to blob.
+        DataType::Binary | DataType::LargeBinary => ColumnType::Blob,
         DataType::FixedSizeBinary(num_bytes) => ColumnType::Binary(num_bytes.to_owned() as u32),
         DataType::Interval(_) => ColumnType::Interval(None, None),
         // Add more mappings here as needed

src/sql/arrow_sql_gen/statement.rs

@@ -50,6 +50,7 @@ pub enum Error {
     UnknownMySQLDatabase { message: String },
 }
 
+#[derive(Debug)]
 pub struct MySQLConnectionPool {
     pool: Arc<mysql_async::Pool>,
     join_push_down: JoinPushDown,

src/sql/db_connection_pool/mysqlpool.rs

@@ -4,13 +4,13 @@ use std::collections::HashMap;
 use std::hash::Hash;
 use std::sync::Arc;
 
+use crate::sql::sql_provider_datafusion::Engine;
+use datafusion::common::DataFusionError;
 use datafusion::{
     error::Result as DataFusionResult,
     sql::unparser::{dialect::DefaultDialect, Unparser},
 };
 
-use crate::sql::sql_provider_datafusion::Engine;
-
 pub mod column_reference;
 pub mod constraints;
 pub mod indexes;
@@ -79,6 +79,14 @@ pub fn remove_prefix_from_hashmap_keys<V>(
         .collect()
 }
 
+#[must_use]
+pub fn to_datafusion_error<E>(error: E) -> DataFusionError
+where
+    E: std::error::Error + Send + Sync + 'static,
+{
+    DataFusionError::External(Box::new(error))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;