Move to obj store for tpch gen (#323)

Author: matthewmturner

.github/actions/setup-rust/action.yml

@@ -37,4 +37,4 @@ runs:
       shell: bash
       run: |
         cargo install taplo-cli --locked
-        cargo install cargo-machete
+        cargo install cargo-machete just

.github/workflows/test.yml

@@ -20,7 +20,7 @@ on: [push, pull_request]
 
 jobs:
   test:
-    name: Core
+    name: DB
     runs-on: ubuntu-latest
     strategy:
       matrix:
@@ -43,7 +43,7 @@ jobs:
         uses: ./.github/actions/setup-rust
       - name: Run core tests
         run: |
-          cargo test
+          cargo test db
   test-flightsql:
     name: Extension / FlightSQL
     runs-on: ubuntu-latest
@@ -72,6 +72,57 @@ jobs:
           # test to only be run against the server spun up in that test.  With parallelism tests 
           # can connec to server in different test which breaks determinism.
           cargo test --features=flightsql extension_cases::flightsql -- --test-threads=1
+  test-cli:
+    name: App / CLI
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        arch: [amd64]
+    steps:
+      - uses: actions/checkout@v3
+      - name: Setup Rust Toolchain
+        uses: ./.github/actions/setup-rust
+      - name: Start LocalStack
+        uses: LocalStack/setup-localstack@v0.2.3
+        with:
+          image-tag: 'latest'
+          install-awslocal: 'true'
+          configuration: DEBUG=1
+      - name: Run Tests against LocalStack
+        run: |
+          awslocal s3 mb s3://test
+          awslocal s3 mv data/aggregate_test_100.csv s3://test/
+          awslocal s3 mb s3://tpch-db
+          echo "Test Execution complete!"
+      - name: Run CLI tests
+        run: |
+          cargo test cli_cases
+  test-tui:
+    name: App / TUI
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        arch: [amd64]
+    steps:
+      - uses: actions/checkout@v3
+      - name: Setup Rust Toolchain
+        uses: ./.github/actions/setup-rust
+      - name: Start LocalStack
+        uses: LocalStack/setup-localstack@v0.2.3
+        with:
+          image-tag: 'latest'
+          install-awslocal: 'true'
+          configuration: DEBUG=1
+      - name: Run Tests against LocalStack
+        run: |
+          awslocal s3 mb s3://test
+          awslocal s3 mv data/aggregate_test_100.csv s3://test/
+          awslocal s3 mb s3://tpch-db
+          echo "Test Execution complete!"
+      - name: Run CLI tests
+        run: |
+          cargo test tui_cases
+
   test-s3:
     name: Extension / S3
     runs-on: ubuntu-latest

Cargo.lock

@@ -34,6 +34,7 @@ lazy_static = "1.4.0"
 log = "0.4.22"
 metrics = { version = "0.24.0", optional = true }
 metrics-exporter-prometheus = { version = "0.16.0", optional = true }
+object_store = "0.11.2"
 parquet = "54"
 pin-project-lite = { version = "0.2.14" }
 prost = "0.13.1"
@@ -62,7 +63,7 @@ tracing = { version = "0.1.41", features = ["log"] }
 tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
 tui-logger = { version = "0.12", features = ["tracing-support"] }
 tui-textarea = { version = "0.6.1", features = ["search"] }
-url = { version = "2.5.2", optional = true }
+url = { version = "2.5.2", features = ["serde"] }
 uuid = { version = "1.10.0", optional = true }
 
 [dev-dependencies]
@@ -80,7 +81,7 @@ url = "2.5.2"
 
 # When addding a new feature, also add it to the features tested list in CI (`.github/workflows/rust.yml`)
 [features]
-default = ["functions-parquet"]
+default = ["functions-parquet", "s3"]
 deltalake = ["datafusion-app/deltalake"]
 flightsql = [
   "datafusion-app/flightsql",
@@ -105,7 +106,7 @@ http = [
   "dep:uuid",
 ]
 huggingface = ["datafusion-app/huggingface"]
-s3 = ["datafusion-app/s3", "url"]
+s3 = ["datafusion-app/s3"]
 udfs-wasm = ["datafusion-app/udfs-wasm"]
 
 [[bin]]

Cargo.toml

@@ -36,3 +36,7 @@ bench-http-custom file:
     echo "Running bench on $custom_bench_path"
     echo ""
     oha --urls-from-file "$custom_bench_path"
+
+setup-test-env:
+    localstack start -d
+    awslocal s3api create-bucket --bucket tmp --acl public-read

justfile

@@ -30,6 +30,7 @@ use serde::Deserialize;
 
 #[cfg(any(feature = "flightsql", feature = "http"))]
 use datafusion_app::config::AuthConfig;
+use url::Url;
 
 lazy_static! {
     pub static ref PROJECT_NAME: String = env!("CARGO_CRATE_NAME").to_uppercase().to_string();
@@ -194,7 +195,7 @@ fn default_interaction_config() -> InteractionConfig {
 #[derive(Debug, Clone, Deserialize)]
 pub struct DbConfig {
     #[serde(default = "default_db_path")]
-    pub path: PathBuf,
+    pub path: Url,
 }
 
 impl Default for DbConfig {
@@ -209,9 +210,18 @@ fn default_db_config() -> DbConfig {
     }
 }
 
-fn default_db_path() -> PathBuf {
+#[allow(unused)]
+fn default_db_path() -> Url {
     let base = directories::BaseDirs::new().expect("Base directories should be available");
-    base.data_dir().to_path_buf().join("dft")
+    let path = base
+        .data_dir()
+        .to_path_buf()
+        .join("dft/")
+        .to_str()
+        .unwrap()
+        .to_string();
+    let with_schema = format!("file://{path}");
+    Url::parse(&with_schema).unwrap()
 }
 
 #[derive(Clone, Debug, Deserialize)]

src/config.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::{fs::read_dir, sync::Arc};
+use std::sync::Arc;
 
 use color_eyre::{Report, Result};
 use datafusion::{
@@ -26,34 +26,26 @@ use datafusion::{
     },
     prelude::SessionContext,
 };
-use log::{error, info};
+use log::info;
 
 use crate::config::DbConfig;
 
 pub async fn register_db(ctx: &SessionContext, db_config: &DbConfig) -> Result<()> {
     info!("registering tables to database");
-    let tables_path = db_config.path.join("tables");
-    if !tables_path.exists() || !tables_path.is_dir() {
-        info!("no tables directory configured, skipping table registration");
-        return Ok(());
-    }
-    let catalogs = read_dir(tables_path)?;
-    info!("...reading catalogs");
-    for maybe_catalog in catalogs {
-        let catalog = maybe_catalog?;
-        let catalog_file_name = catalog.file_name();
-        let catalog_name = catalog_file_name.to_str().ok_or(Report::msg(format!(
-            "invalid catalog path {catalog_file_name:?}"
-        )))?;
-        // Every catalog should be a directory
-        if !catalog.path().is_dir() {
-            error!("catalog {catalog_name:?} is not a directory, skipping");
-            continue;
-        }
-        let catalog_path = catalog.path();
-        info!("...handling {:?} catalog", catalog_name);
+    let tables_url = db_config.path.join("tables")?;
+    let listing_tables_url = ListingTableUrl::parse(tables_url.clone())?;
+    let store_url = listing_tables_url.object_store();
+    let store = ctx.runtime_env().object_store(store_url)?;
+    let tables_path = object_store::path::Path::from_url_path(tables_url.path())?;
+    let catalogs = store.list_with_delimiter(Some(&tables_path)).await?;
+    for catalog in catalogs.common_prefixes {
+        let catalog_name = catalog
+            .filename()
+            .ok_or(Report::msg("missing catalog name"))?;
+        info!("...handling {catalog_name} catalog");
         let maybe_catalog = ctx.catalog(catalog_name);
         let catalog_provider = match maybe_catalog {
+            Some(catalog) => catalog,
             None => {
                 info!("...catalog does not exist, createing");
                 let mem_catalog_provider = Arc::new(MemoryCatalogProvider::new());
@@ -62,23 +54,16 @@ pub async fn register_db(ctx: &SessionContext, db_config: &DbConfig) -> Result<(
                     "missing catalog {catalog_name}, shouldnt be possible"
                 )))?
             }
-            Some(catalog) => catalog,
         };
-        for maybe_schema in read_dir(&catalog_path)? {
-            let schema = maybe_schema?;
-            let schema_file_name = schema.file_name();
-            let schema_name = schema_file_name.to_str().ok_or(Report::msg(format!(
-                "invalid schema path {schema_file_name:?}"
-            )))?;
-            // Every schema should be a directory
-            if !schema.path().is_dir() {
-                error!("schema {schema_name:?} is not a directory, skipping",);
-                continue;
-            }
-            let schema_path = schema.path();
-            info!("...handling {:?} schema", schema_name);
+        let schemas = store.list_with_delimiter(Some(&catalog)).await?;
+        for schema in schemas.common_prefixes {
+            let schema_name = schema
+                .filename()
+                .ok_or(Report::msg("missing schema name"))?;
+            info!("...handling {schema_name} schema");
             let maybe_schema = catalog_provider.schema(schema_name);
             let schema_provider = match maybe_schema {
+                Some(schema) => schema,
                 None => {
                     info!("...schema does not exist, creating");
                     let mem_schema_provider = Arc::new(MemorySchemaProvider::new());
@@ -89,24 +74,19 @@ pub async fn register_db(ctx: &SessionContext, db_config: &DbConfig) -> Result<(
                             "missing schema {schema_name}, shouldnt be possible"
                         )))?
                 }
-                Some(schema) => schema,
             };
-            for maybe_table in read_dir(schema_path)? {
-                let table = maybe_table?;
-                // Every table should be a directory even if there is a single data file
-                if !table.path().is_dir() {
-                    error!("table {:?} is not a directory, skipping", catalog.path());
-                    continue;
-                }
-                let table_path = table.path();
-                let table_file_name = table.file_name();
-                let table_name = table_file_name.to_str().ok_or(Report::msg(format!(
-                    "invalid table path {table_file_name:?}"
-                )))?;
-                info!("...handling table {table_name:?}");
-                let table_url = ListingTableUrl::parse(table_path.to_str().ok_or(Report::msg(
-                    format!("Invalid table path for {table_path:?}"),
-                ))?)?;
+            let tables = store.list_with_delimiter(Some(&schema)).await?;
+            for table_path in tables.common_prefixes {
+                let table_name = table_path
+                    .filename()
+                    .ok_or(Report::msg("missing table name"))?;
+                info!("...handling table \"{catalog_name}.{schema_name}.{table_name}\"");
+
+                let p = tables_url
+                    .join(&format!("{catalog_name}/"))?
+                    .join(&format!("{schema_name}/"))?
+                    .join(&format!("{table_name}/"))?;
+                let table_url = ListingTableUrl::parse(p)?;
                 let file_format = ParquetFormat::new();
                 let listing_options =
                     ListingOptions::new(Arc::new(file_format)).with_file_extension(".parquet");
@@ -119,11 +99,12 @@ pub async fn register_db(ctx: &SessionContext, db_config: &DbConfig) -> Result<(
                     .with_schema(resolved_schema);
                 // Create a new TableProvider
                 let provider = Arc::new(ListingTable::try_new(config)?);
-                info!("...registering {table_name}");
+                info!("...table registered");
                 schema_provider.register_table(table_name.to_string(), provider)?;
             }
         }
     }
+
     Ok(())
 }
 
@@ -179,9 +160,9 @@ mod test {
         let ctx = setup();
         let dir = tempfile::tempdir().unwrap();
         let db_path = dir.path().join("db");
-        let config = DbConfig {
-            path: db_path.clone(),
-        };
+        let path = format!("file://{}/", db_path.to_str().unwrap());
+        let db_url = url::Url::parse(&path).unwrap();
+        let config = DbConfig { path: db_url };
         let data_path = db_path.join("tables").join("dft").join("stuff").join("hi");
 
         let df = ctx.sql("SELECT 1").await.unwrap();
@@ -231,9 +212,9 @@ mod test {
         let ctx = setup();
         let dir = tempfile::tempdir().unwrap();
         let db_path = dir.path().join("db");
-        let config = DbConfig {
-            path: db_path.clone(),
-        };
+        let path = format!("file://{}/", db_path.to_str().unwrap());
+        let db_url = url::Url::parse(&path).unwrap();
+        let config = DbConfig { path: db_url };
         let data_1_path = db_path.join("tables").join("dft").join("stuff").join("hi");
         let data_2_path = db_path.join("tables").join("dft").join("stuff").join("bye");
 
@@ -290,9 +271,9 @@ mod test {
         let ctx = setup();
         let dir = tempfile::tempdir().unwrap();
         let db_path = dir.path().join("db");
-        let config = DbConfig {
-            path: db_path.clone(),
-        };
+        let path = format!("file://{}/", db_path.to_str().unwrap());
+        let db_url = url::Url::parse(&path).unwrap();
+        let config = DbConfig { path: db_url };
         let data_1_path = db_path.join("tables").join("dft").join("stuff").join("hi");
         let data_2_path = db_path
             .join("tables")
@@ -353,9 +334,9 @@ mod test {
         let ctx = setup();
         let dir = tempfile::tempdir().unwrap();
         let db_path = dir.path().join("db");
-        let config = DbConfig {
-            path: db_path.clone(),
-        };
+        let path = format!("file://{}/", db_path.to_str().unwrap());
+        let db_url = url::Url::parse(&path).unwrap();
+        let config = DbConfig { path: db_url };
         let data_1_path = db_path.join("tables").join("dft2").join("stuff").join("hi");
         let data_2_path = db_path
             .join("tables")

src/db.rs

@@ -62,7 +62,7 @@ async fn app_entry_point(cli: DftArgs) -> Result<()> {
     }
     let cfg = create_config(cli.config_path());
     if let Some(Command::GenerateTpch { scale_factor }) = cli.command {
-        tpch::generate(cfg.clone(), scale_factor)?;
+        tpch::generate(cfg.clone(), scale_factor).await?;
         return Ok(());
     }