Propagate MySQL wrong table error

datafusion-contrib · Sep 20, 2024 · 45d2113 · 45d2113
1 parent c9356e3
commit 45d2113
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 14 deletions.
diff --git a/src/mysql.rs b/src/mysql.rs
@@ -55,7 +55,7 @@ impl MySQLTableFactory {
         let table_provider = Arc::new(
             SqlTable::new("mysql", &pool, table_reference, None)
                 .await
-                .context(UnableToConstructSQLTableSnafu)?
+                .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?
                 .with_dialect(Arc::new(MySqlDialect {})),
         );
 

diff --git a/src/postgres.rs b/src/postgres.rs
@@ -8,7 +8,6 @@ use crate::sql::db_connection_pool::{
     DbConnectionPool,
 };
 use crate::sql::sql_provider_datafusion::{
-    self,
     expr::{self, Engine},
     SqlTable,
 };
@@ -90,13 +89,6 @@ pub enum Error {
         source: tokio_postgres::error::Error,
     },
 
-    #[snafu(display(
-        "Unable to construct the DataFusion SQL Table Provider for Postgres: {source}"
-    ))]
-    UnableToConstructSqlTable {
-        source: sql_provider_datafusion::Error,
-    },
-
     #[snafu(display("Unable to generate SQL: {source}"))]
     UnableToGenerateSQL { source: expr::Error },
 

diff --git a/src/sql/db_connection_pool/dbconnection/mysqlconn.rs b/src/sql/db_connection_pool/dbconnection/mysqlconn.rs
@@ -78,11 +78,27 @@ impl<'a> AsyncDbConnection<Conn, &'a (dyn ToValue + Sync)> for MySQLConnection {
         let columns_meta_query =
             format!("SHOW COLUMNS FROM {}", table_reference.to_quoted_string());
 
-        let columns_meta: Vec<Row> = conn
-            .exec(&columns_meta_query, Params::Empty)
-            .await
-            .boxed()
-            .context(super::UnableToGetSchemaSnafu)?;
+        let columns_meta: Vec<Row> = match conn.exec(&columns_meta_query, Params::Empty).await {
+            Ok(columns_meta) => columns_meta,
+            Err(e) => match e {
+                mysql_async::Error::Server(server_error) => {
+                    if server_error.code == 1146 {
+                        return Err(super::Error::UndefinedTable {
+                            source: Box::new(server_error.clone()),
+                            table_name: table_reference.to_string(),
+                        });
+                    }
+                    return Err(super::Error::UnableToGetSchema {
+                        source: Box::new(mysql_async::Error::Server(server_error)),
+                    });
+                }
+                _ => {
+                    return Err(super::Error::UnableToGetSchema {
+                        source: Box::new(e),
+                    })
+                }
+            },
+        };
 
         Ok(columns_meta_to_schema(columns_meta).context(super::UnableToGetSchemaSnafu)?)
     }