Fix delete and update tests

splitgraph · Feb 5, 2024 · e538ead · e538ead
1 parent 7e50702
commit e538ead
Show file tree

Hide file tree

Showing 11 changed files with 294 additions and 259 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -51,6 +51,21 @@ object-store-s3 = ["object_store/aws"]
 remote-tables = ["dep:datafusion-remote-tables"]
 
 [patch.crates-io]
+# Pick up https://github.com/apache/arrow-rs/pull/5282
+arrow-arith = { git = "https://github.com/apache/arrow-rs", rev = "72d8a783176219f0864022daba70e84ceab7e221" }
+arrow-array = { git = "https://github.com/apache/arrow-rs", rev = "72d8a783176219f0864022daba70e84ceab7e221" }
+arrow-buffer = { git = "https://github.com/apache/arrow-rs", rev = "72d8a783176219f0864022daba70e84ceab7e221" }
+arrow-cast = { git = "https://github.com/apache/arrow-rs", rev = "72d8a783176219f0864022daba70e84ceab7e221" }
+arrow-csv = { git = "https://github.com/apache/arrow-rs", rev = "72d8a783176219f0864022daba70e84ceab7e221" }
+arrow-data = { git = "https://github.com/apache/arrow-rs", rev = "72d8a783176219f0864022daba70e84ceab7e221" }
+arrow-ipc = { git = "https://github.com/apache/arrow-rs", rev = "72d8a783176219f0864022daba70e84ceab7e221" }
+arrow-json = { git = "https://github.com/apache/arrow-rs", rev = "72d8a783176219f0864022daba70e84ceab7e221" }
+arrow-ord = { git = "https://github.com/apache/arrow-rs", rev = "72d8a783176219f0864022daba70e84ceab7e221" }
+arrow-row = { git = "https://github.com/apache/arrow-rs", rev = "72d8a783176219f0864022daba70e84ceab7e221" }
+arrow-schema = { git = "https://github.com/apache/arrow-rs", rev = "72d8a783176219f0864022daba70e84ceab7e221" }
+arrow-select = { git = "https://github.com/apache/arrow-rs", rev = "72d8a783176219f0864022daba70e84ceab7e221" }
+arrow-string = { git = "https://github.com/apache/arrow-rs", rev = "72d8a783176219f0864022daba70e84ceab7e221" }
+
 # Pick up https://github.com/apache/arrow-datafusion/pull/8894 and https://github.com/apache/arrow-datafusion/pull/9007
 datafusion = { git = "https://github.com/apache/arrow-datafusion", rev = "a7a74fa522aaef07e6605f414308f3c99bd1ea06" }
 datafusion-common = { git = "https://github.com/apache/arrow-datafusion", rev = "a7a74fa522aaef07e6605f414308f3c99bd1ea06" }
@@ -93,7 +108,7 @@ datafusion-expr = { workspace = true }
 
 datafusion-remote-tables = { path = "./datafusion_remote_tables", optional = true }
 
-deltalake = { path = "../delta-rs/crates/core", features = ["datafusion"], package = "deltalake-core" }
+deltalake = { git = "https://github.com/splitgraph/delta-rs", branch = "sqlparser-0.43", features = ["datafusion"] }
 
 futures = "0.3"
 hex = ">=0.4.0"

diff --git a/datafusion_remote_tables/Cargo.toml b/datafusion_remote_tables/Cargo.toml
@@ -19,7 +19,7 @@ arrow-schema = { workspace = true }
 async-trait = { workspace = true }
 
 # Remote query execution for a variety of DBs
-connectorx = { path = "../../connector-x/connectorx", features = [ "dst_arrow", "src_postgres", "src_mysql", "src_sqlite" ] }
+connectorx = { git = "https://github.com/splitgraph/connector-x", branch = "datafusion-35-upgrade", features = [ "dst_arrow", "src_postgres", "src_mysql", "src_sqlite" ] }
 
 datafusion = { workspace = true }
 datafusion-common = { workspace = true }

diff --git a/src/context/logical.rs b/src/context/logical.rs
@@ -484,24 +484,6 @@ mod tests {
         );
     }
 
-    #[tokio::test]
-    async fn test_plan_insert_type_mismatch() {
-        let ctx = in_memory_context_with_test_db().await;
-
-        // Try inserting a string into a date (note this will work fine for inserting
-        // e.g. Utf-8 into numbers at plan time but should fail at execution time if the value
-        // doesn't convert)
-        let plan = ctx
-            .create_logical_plan("INSERT INTO testcol.some_table SELECT 'abc', to_timestamp('2022-01-01T12:00:00')")
-            .await.unwrap();
-
-        let err = ctx.create_physical_plan(&plan).await.unwrap_err();
-        assert_eq!(
-            err.to_string(),
-            "Arrow error: Cast error: Cannot cast string 'abc' to value of Date32 type"
-        );
-    }
-
     #[tokio::test]
     async fn test_plan_insert_values_wrong_number() {
         let ctx = in_memory_context_with_test_db().await;

diff --git a/src/context/mod.rs b/src/context/mod.rs
@@ -259,6 +259,25 @@ mod tests {
     use super::test_utils::in_memory_context;
     use super::*;
 
+    #[tokio::test]
+    async fn test_timestamp_to_date_casting() -> Result<()> {
+        let ctx = in_memory_context().await;
+
+        let plan = ctx.plan_query("SELECT '1998-11-30 00:00:00'::date").await?;
+
+        let results = ctx.collect(plan).await?;
+        let expected = [
+            "+-----------------------------+",
+            "| Utf8(\"1998-11-30 00:00:00\") |",
+            "+-----------------------------+",
+            "| 1998-11-30                  |",
+            "+-----------------------------+",
+        ];
+        assert_batches_eq!(expected, &results);
+
+        Ok(())
+    }
+
     #[rstest]
     #[case::regular_type_names("float", "float")]
     #[case::legacy_type_names("f32", "f32")]

diff --git a/src/context/physical.rs b/src/context/physical.rs
@@ -899,6 +899,7 @@ impl SeafowlContext {
 
 #[cfg(test)]
 mod tests {
+    use crate::context::test_utils::in_memory_context_with_test_db;
     use datafusion::assert_batches_eq;
     use std::sync::Arc;
 
@@ -1021,4 +1022,22 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_plan_insert_type_mismatch() {
+        let ctx = in_memory_context_with_test_db().await;
+
+        // Try inserting a string into a date (note this will work fine for inserting
+        // e.g. Utf-8 into numbers at plan time but should fail at execution time if the value
+        // doesn't convert)
+        let plan = ctx
+            .create_logical_plan("INSERT INTO testcol.some_table SELECT 'abc', to_timestamp('2022-01-01T12:00:00')")
+            .await.unwrap();
+
+        let err = ctx.create_physical_plan(&plan).await.unwrap_err();
+        assert_eq!(
+            err.to_string(),
+            "Arrow error: Cast error: Cannot cast string 'abc' to value of Date32 type"
+        );
+    }
 }
diff --git a/src/provider.rs b/src/provider.rs
@@ -86,7 +86,7 @@ impl SchemaProvider for SeafowlSchema {
         // Ultimately though, since the map gets re-created for each query the only point in
         // updating the existing table is to optimize potential multi-lookups during processing of
         // a single query.
-        let table_log_store = match self.tables.read().get(name) {
+        let mut delta_table = match self.tables.read().get(name) {
             None => return None,
             Some(table) => match table.as_any().downcast_ref::<DeltaTable>() {
                 // This shouldn't happen since we store only DeltaTable's in the map
@@ -98,14 +98,12 @@ impl SchemaProvider for SeafowlSchema {
                     } else {
                         // A negative table version indicates that the table was never loaded; we need
                         // to do it before returning it.
-                        delta_table.log_store()
+                        delta_table.clone()
                     }
                 }
             },
         };
 
-        let mut delta_table = DeltaTable::new(table_log_store, Default::default());
-
         if let Err(err) = delta_table.load().await {
             warn!("Failed to load table {name}: {err}");
             return None;

diff --git a/src/wasm_udf/mod.rs b/src/wasm_udf/mod.rs
@@ -1,2 +1,3 @@
 pub mod data_types;
+#[allow(deprecated)]
 pub mod wasm;
diff --git a/tests/statements/ddl.rs b/tests/statements/ddl.rs
@@ -533,7 +533,7 @@ async fn test_create_external_table(
     #[case] location: Option<&str>,
     #[case] options: &str,
     #[case] object_store_type: ObjectStoreType,
-) {
+) -> Result<()> {
     let url = match location {
         None => {
             let (mock_server, _) = testutils::make_mock_parquet_server(true, true).await;
@@ -574,8 +574,7 @@ async fn test_create_external_table(
             )
             .as_str(),
         )
-        .await
-        .unwrap();
+        .await?;
 
     // Test we see the table in the information_schema
     let results = list_tables_query(&context).await;
@@ -594,11 +593,8 @@ async fn test_create_external_table(
     assert_batches_eq!(expected, &results);
 
     // Test standard query
-    let plan = context
-        .plan_query("SELECT * FROM staging.file")
-        .await
-        .unwrap();
-    let results = context.collect(plan).await.unwrap();
+    let plan = context.plan_query("SELECT * FROM staging.file").await?;
+    let results = context.collect(plan).await?;
     let expected = if location.is_none() {
         vec![
             "+-------+",
@@ -626,8 +622,7 @@ async fn test_create_external_table(
     // Test dropping the external table works
     context
         .collect(context.plan_query("DROP TABLE staging.file").await.unwrap())
-        .await
-        .unwrap();
+        .await?;
 
     let results = list_tables_query(&context).await;
 
@@ -642,4 +637,6 @@ async fn test_create_external_table(
         "+--------------------+-------------+",
     ];
     assert_batches_eq!(expected, &results);
+
+    Ok(())
 }