Make inline metastore robust when local FS store with path is provided

splitgraph · Jul 26, 2024 · 26bb61e · 26bb61e
1 parent b6e6a42
commit 26bb61e
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 5 deletions.
diff --git a/object_store_factory/src/lib.rs b/object_store_factory/src/lib.rs
@@ -4,8 +4,10 @@ mod local;
 mod memory;
 
 use object_store::parse_url_opts;
+use object_store::prefix::PrefixStore;
 use object_store::ObjectStore;
 use object_store::ObjectStoreScheme;
+use object_store::{local::LocalFileSystem, memory::InMemory};
 use std::collections::HashMap;
 use std::sync::Arc;
 use tracing::warn;
@@ -79,6 +81,18 @@ pub async fn build_object_store_from_opts(
     let (scheme, _) = ObjectStoreScheme::parse(url).unwrap();
 
     match scheme {
+        // `parse_url_opts` will swallow the URL path for memory/local FS stores
+        ObjectStoreScheme::Memory => {
+            let mut store: Box<dyn ObjectStore> = Box::new(InMemory::new());
+            if !url.path().is_empty() {
+                store = Box::new(PrefixStore::new(store, url.path()));
+            }
+
+            Ok(store)
+        }
+        ObjectStoreScheme::Local => {
+            Ok(Box::new(LocalFileSystem::new_with_prefix(url.path())?))
+        }
         ObjectStoreScheme::AmazonS3 => {
             let mut s3_options = aws::map_options_into_amazon_s3_config_keys(options)?;
             aws::add_amazon_s3_specific_options(url, &mut s3_options).await;

diff --git a/tests/fixtures.rs b/tests/fixtures.rs
@@ -26,11 +26,18 @@ pub fn schemas() -> ListSchemaResponse {
         schemas: vec![
             SchemaObject {
                 name: "local".to_string(),
-                tables: vec![TableObject {
-                    name: "file".to_string(),
-                    path: "delta-0.8.0-partitioned".to_string(),
-                    store: None,
-                }],
+                tables: vec![
+                    TableObject {
+                        name: "file".to_string(),
+                        path: "delta-0.8.0-partitioned".to_string(),
+                        store: None,
+                    },
+                    TableObject {
+                        name: "file_with_store".to_string(),
+                        path: "delta-0.8.0-partitioned".to_string(),
+                        store: Some("local_fs".to_string()),
+                    },
+                ],
             },
             SchemaObject {
                 name: "s3".to_string(),
@@ -82,6 +89,14 @@ pub fn schemas() -> ListSchemaResponse {
                     fake_gcs_creds(),
                 )]),
             },
+            StorageLocation {
+                name: "local_fs".to_string(),
+                location: format!(
+                    "file://{}/tests/data/",
+                    std::env::current_dir().unwrap().display()
+                ),
+                options: HashMap::new(),
+            },
         ],
     }
 }
diff --git a/tests/flight/inline_metastore.rs b/tests/flight/inline_metastore.rs
@@ -3,6 +3,7 @@ use crate::flight::*;
 #[rstest]
 #[should_panic(expected = "External error: Not a Delta table: no log files")]
 #[case("local.file", false)]
+#[case("local.file_with_store", false)]
 #[case("local.file", true)]
 #[case("s3.minio", true)]
 #[case("gcs.fake", true)]