use quote style for column and table names based on the sql dialect (#102)

* don't wrap Engine with Options in SqlTable and SqlExec

* sql/sql_provider: use quote style for column and table names based on the dialect

---------

Co-authored-by: Phillip LeBlanc <phillip@leblanc.tech>

Author: hozan23

src/sql/sql_provider_datafusion/federation.rs

@@ -1,6 +1,5 @@
 use crate::sql::db_connection_pool::{dbconnection::get_schema, JoinPushDown};
 use async_trait::async_trait;
-use datafusion::sql::unparser::dialect::DefaultDialect;
 use datafusion_federation::sql::{SQLExecutor, SQLFederationProvider, SQLTableSource};
 use datafusion_federation::{FederatedTableProviderAdaptor, FederatedTableSource};
 use futures::TryStreamExt;
@@ -58,9 +57,7 @@ impl<T, P> SQLExecutor for SqlTable<T, P> {
     }
 
     fn dialect(&self) -> Arc<dyn Dialect> {
-        self.engine
-            .map(|e| e.dialect())
-            .unwrap_or_else(|| Arc::new(DefaultDialect {}))
+        self.engine.dialect()
     }
 
     fn execute(

src/sql/sql_provider_datafusion/mod.rs

@@ -32,7 +32,7 @@ use datafusion::{
         stream::RecordBatchStreamAdapter, DisplayAs, DisplayFormatType, ExecutionMode,
         ExecutionPlan, Partitioning, PlanProperties, SendableRecordBatchStream,
     },
-    sql::{unparser::Unparser, TableReference},
+    sql::{sqlparser::ast, unparser::Unparser, TableReference},
 };
 
 pub mod federation;
@@ -59,6 +59,7 @@ pub enum Engine {
     ODBC,
     Postgres,
     MySQL,
+    Default,
 }
 
 impl Engine {
@@ -68,19 +69,22 @@ impl Engine {
             Engine::SQLite => Arc::new(SqliteDialect {}),
             Engine::Postgres => Arc::new(PostgreSqlDialect {}),
             Engine::MySQL => Arc::new(MySqlDialect {}),
-            Engine::Spark | Engine::DuckDB | Engine::ODBC => Arc::new(DefaultDialect {}),
+            Engine::Spark | Engine::DuckDB | Engine::ODBC | Engine::Default => {
+                Arc::new(DefaultDialect {})
+            }
         }
     }
 }
 
 pub type Result<T, E = Error> = std::result::Result<T, E>;
 
+#[derive(Clone)]
 pub struct SqlTable<T: 'static, P: 'static> {
     name: &'static str,
     pool: Arc<dyn DbConnectionPool<T, P> + Send + Sync>,
     schema: SchemaRef,
     pub table_reference: TableReference,
-    engine: Option<Engine>,
+    engine: Engine,
 }
 
 impl<T, P> SqlTable<T, P> {
@@ -116,6 +120,7 @@ impl<T, P> SqlTable<T, P> {
         table_reference: impl Into<TableReference>,
         engine: Option<Engine>,
     ) -> Self {
+        let engine = engine.unwrap_or(Engine::Default);
         Self {
             name,
             pool: Arc::clone(pool),
@@ -139,7 +144,7 @@ impl<T, P> SqlTable<T, P> {
             Arc::clone(&self.pool),
             filters,
             limit,
-            self.engine,
+            Some(self.engine),
         )?))
     }
 
@@ -177,16 +182,14 @@ impl<T, P> TableProvider for SqlTable<T, P> {
         &self,
         filters: &[&Expr],
     ) -> DataFusionResult<Vec<TableProviderFilterPushDown>> {
-        let dialect = self
-            .engine
-            .map(|e| e.dialect())
-            .unwrap_or_else(|| Arc::new(DefaultDialect {}));
         let filter_push_down: Vec<TableProviderFilterPushDown> = filters
             .iter()
-            .map(|f| match Unparser::new(dialect.as_ref()).expr_to_sql(f) {
-                Ok(_) => TableProviderFilterPushDown::Exact,
-                Err(_) => TableProviderFilterPushDown::Unsupported,
-            })
+            .map(
+                |f| match Unparser::new(self.engine.dialect().as_ref()).expr_to_sql(f) {
+                    Ok(_) => TableProviderFilterPushDown::Exact,
+                    Err(_) => TableProviderFilterPushDown::Unsupported,
+                },
+            )
             .collect();
 
         Ok(filter_push_down)
@@ -217,7 +220,7 @@ pub struct SqlExec<T, P> {
     filters: Vec<Expr>,
     limit: Option<usize>,
     properties: PlanProperties,
-    engine: Option<Engine>,
+    engine: Engine,
 }
 
 pub fn project_schema_safe(
@@ -248,6 +251,7 @@ impl<T, P> SqlExec<T, P> {
         engine: Option<Engine>,
     ) -> DataFusionResult<Self> {
         let projected_schema = project_schema_safe(schema, projections)?;
+        let engine = engine.unwrap_or(Engine::Default);
 
         Ok(Self {
             projected_schema: Arc::clone(&projected_schema),
@@ -286,15 +290,11 @@ impl<T, P> SqlExec<T, P> {
         let where_expr = if self.filters.is_empty() {
             String::new()
         } else {
-            let dialect = self
-                .engine
-                .map(|e| e.dialect())
-                .unwrap_or_else(|| Arc::new(DefaultDialect {}));
             let filter_expr = self
                 .filters
                 .iter()
                 .map(|f| {
-                    Unparser::new(dialect.as_ref())
+                    Unparser::new(self.engine.dialect().as_ref())
                         .expr_to_sql(f)
                         .map(|e| e.to_string())
                 })
@@ -311,16 +311,20 @@ impl<T, P> SqlExec<T, P> {
     }
 
     fn table_name_escaped(&self) -> String {
-        self.table_reference.to_quoted_string()
+        self.ident_escaped(&self.table_reference.to_string())
     }
 
     fn column_name_escaped(&self, column_name: &str) -> String {
-        match self.engine {
-            Some(Engine::ODBC) => column_name.to_string(),
-            _ => {
-                format!("\"{}\"", column_name)
-            }
-        }
+        self.ident_escaped(column_name)
+    }
+
+    fn ident_escaped(&self, ident: &str) -> String {
+        let quote_style = self.engine.dialect().identifier_quote_style(ident);
+        ast::Expr::Identifier(ast::Ident {
+            value: ident.to_string(),
+            quote_style,
+        })
+        .to_string()
     }
 }
 
@@ -485,6 +489,7 @@ mod tests {
                 ),
                 Field::new("userId", DataType::LargeUtf8, false),
                 Field::new("active", DataType::Boolean, false),
+                Field::new("5e48", DataType::LargeUtf8, false),
             ];
             let schema = Arc::new(Schema::new(fields));
             let pool = Arc::new(MockDBPool {})
@@ -511,18 +516,18 @@ mod tests {
 
         #[tokio::test]
         async fn test_sql_to_string_with_limit() -> Result<(), Box<dyn Error + Send + Sync>> {
-            let sql_exec = new_sql_exec(Some(&vec![0]), "users", &[], Some(3), None)?;
-            assert_eq!(sql_exec.sql()?, r#"SELECT "name" FROM users  LIMIT 3"#);
+            let sql_exec = new_sql_exec(Some(&vec![0, 1]), "users", &[], Some(3), None)?;
+            assert_eq!(sql_exec.sql()?, r#"SELECT "name", age FROM users  LIMIT 3"#);
             Ok(())
         }
 
         #[tokio::test]
         async fn test_sql_to_string_with_filters() -> Result<(), Box<dyn Error + Send + Sync>> {
             let filters = vec![col("age").gt_eq(lit(30)).and(col("name").eq(lit("x")))];
-            let sql_exec = new_sql_exec(None, "users", &filters, None, None)?;
+            let sql_exec = new_sql_exec(Some(&vec![0, 1]), "users", &filters, None, None)?;
             assert_eq!(
                 sql_exec.sql()?,
-                r#"SELECT "name", "age", "createdDate", "userId", "active" FROM users WHERE ((age >= 30) AND ("name" = 'x')) "#
+                r#"SELECT "name", age FROM users WHERE ((age >= 30) AND ("name" = 'x')) "#
             );
             Ok(())
         }
@@ -531,21 +536,57 @@ mod tests {
         async fn test_sql_to_string_with_filters_and_limit(
         ) -> Result<(), Box<dyn Error + Send + Sync>> {
             let filters = vec![col("age").gt_eq(lit(30)).and(col("name").eq(lit("x")))];
-            let sql_exec = new_sql_exec(None, "users", &filters, Some(3), None)?;
+            let sql_exec = new_sql_exec(Some(&vec![0, 1]), "users", &filters, Some(3), None)?;
             assert_eq!(
                 sql_exec.sql()?,
-                r#"SELECT "name", "age", "createdDate", "userId", "active" FROM users WHERE ((age >= 30) AND ("name" = 'x')) LIMIT 3"#
+                r#"SELECT "name", age FROM users WHERE ((age >= 30) AND ("name" = 'x')) LIMIT 3"#
             );
             Ok(())
         }
 
         #[tokio::test]
         async fn test_sql_to_string_with_engine() -> Result<(), Box<dyn Error + Send + Sync>> {
             let filters = vec![col("age").gt_eq(lit(30)).and(col("name").eq(lit("x")))];
-            let sql_exec = new_sql_exec(None, "users", &filters, Some(3), Some(Engine::DuckDB))?;
+            let sql_exec = new_sql_exec(
+                Some(&vec![0, 1]),
+                "users",
+                &filters,
+                Some(3),
+                Some(Engine::DuckDB),
+            )?;
+            assert_eq!(
+                sql_exec.sql()?,
+                r#"SELECT "name", age FROM users WHERE ((age >= 30) AND ("name" = 'x')) LIMIT 3"#
+            );
+            Ok(())
+        }
+
+        #[tokio::test]
+        async fn test_sql_to_string_with_not_reasonable_name(
+        ) -> Result<(), Box<dyn Error + Send + Sync>> {
+            let filters = vec![col("5e48").eq(lit("test")).and(col("name").eq(lit("x")))];
+            let sql_exec = new_sql_exec(Some(&vec![0, 1, 5]), "users", &filters, Some(3), None)?;
+            assert_eq!(
+                sql_exec.sql()?,
+                r#"SELECT "name", age, "5e48" FROM users WHERE (("5e48" = 'test') AND ("name" = 'x')) LIMIT 3"#
+            );
+            Ok(())
+        }
+
+        #[tokio::test]
+        async fn test_sql_to_string_with_not_reasonable_name_mysql(
+        ) -> Result<(), Box<dyn Error + Send + Sync>> {
+            let filters = vec![col("5e48").eq(lit("test")).and(col("name").eq(lit("x")))];
+            let sql_exec = new_sql_exec(
+                Some(&vec![0, 1, 5]),
+                "users",
+                &filters,
+                Some(3),
+                Some(Engine::MySQL),
+            )?;
             assert_eq!(
                 sql_exec.sql()?,
-                r#"SELECT "name", "age", "createdDate", "userId", "active" FROM users WHERE ((age >= 30) AND ("name" = 'x')) LIMIT 3"#
+                r#"SELECT `name`, `age`, `5e48` FROM `users` WHERE ((`5e48` = 'test') AND (`name` = 'x')) LIMIT 3"#
             );
             Ok(())
         }