Repository: roapi/roapi (language: python)

Issue:
Unify native support for S3, GCS and Azure using object_store
### Background
roapi supports S3 before Datafusion adopts crate `object_store`. And it uses `rusoto_s3` to preload S3 files to memory.

https://github.com/roapi/roapi/blob/68ee99a92491379c1d90845d9a807b43e1bde79a/columnq/src/io/s3.rs#L54

### Problem
1. S3 with `use_memory_table: false` fails
2. Added native GCS/Azure support may duplicate work with S3, when different client SDK required

### Expect
1. support S3 with both memory table and datafusion table
2. Easily support native GCS and Azure

### Proposal
Each action item can be a separate PR
- [x] fix S3 with `use_memory_table: false` and use ObjectStoreProvider and global SessionContext: https://github.com/roapi/roapi/pull/241
- [x] replace `rusoto_s3` (in `use_memory_table: true` path) with S3 client from `object_store` (`s3://`)
    - [x] https://github.com/roapi/roapi/pull/244
    - [x] https://github.com/roapi/roapi/pull/245
    - [x] https://github.com/roapi/roapi/pull/247
    - [x] ~release `v0.8.2`~
- [x] extend ObjectStoreProvider to support GCS with `object_store` (`gs://`)
    - [x] https://github.com/roapi/roapi/pull/248
    - [x] https://github.com/roapi/docs/pull/16
- [x] extend ObjectStoreProvider to support Azure with `object_store` (`az://`, `adl://`, `abfs://`, `abfss://`, `azure://`)
    - [x] https://github.com/roapi/roapi/pull/250
    - [x] https://github.com/roapi/docs/pull/17
- [x] release `v0.9.0`: https://github.com/roapi/roapi/pull/251

### Related
There are several action items proposed. If all done, this issue can close
- https://github.com/roapi/roapi/issues/227
- https://github.com/roapi/roapi/issues/186
- https://github.com/roapi/roapi/issues/121


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 87aca0936..bf6ad31d8 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -113,7 +113,7 @@ jobs:
           which cargo-cache || cargo install cargo-cache
           cargo cache trim -l 1G
 
-  s3_test:
+  s3_memory_test:
     runs-on: ubuntu-latest
     env:
       RUSTFLAGS: "-C target-cpu=skylake"
@@ -160,7 +160,60 @@ jobs:
       with:
         command: /home/runner/work/roapi/roapi/test_end_to_end/query_blogs.sh
         build: cargo build
-        start: cargo run --bin roapi -- -c test_end_to_end/test_s3.yml
+        start: cargo run --bin roapi -- -c test_end_to_end/test_s3_memory.yml
+        wait-on: "http://127.0.0.1:8000/api/schema"
+        # By default, wait-on will retry for 60 seconds. You can pass a custom timeout in seconds using wait-on-timeout.
+        # 10 minutes = 600 seconds
+        wait-on-timeout: 600
+
+  s3_direct_test:
+    runs-on: ubuntu-latest
+    env:
+      RUSTFLAGS: "-C target-cpu=skylake"
+    steps:
+    - uses: actions/checkout@v2
+    - uses: actions/cache@v2
+      with:
+        path: |
+          ~/.cargo/bin/
+          ~/.cargo/registry/index/
+          ~/.cargo/registry/cache/
+          ~/.cargo/git/db/
+          target/
+        key: s3-1-${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock', '**/Cargo.toml') }}
+        restore-keys: |
+          s3-1-${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock', '**/Cargo.toml') }}
+          s3-1-${{ runner.os }}-cargo-
+    - name: Install nightly rust
+      uses: actions-rs/toolchain@v1
+      with:
+        profile: default
+        # toolchain: nightly
+        toolchain: nightly-${{ env.RUST_TC_NIGHTLY_VER }}
+        override: true
+    - name: Install minio server
+      run: |
+        docker run -d -p 9000:9000 quay.io/minio/minio server /data
+    - name: Install minio client
+      run: |
+        curl https://dl.min.io/client/mc/release/linux-amd64/mc --create-dirs -o $HOME/minio-binaries/mc
+        chmod +x $HOME/minio-binaries/mc
+    - name: Create bucket and copy test data
+      run: |
+        $HOME/minio-binaries/mc alias set local http://127.0.0.1:9000 minioadmin minioadmin
+        $HOME/minio-binaries/mc mb local/test-data
+        $HOME/minio-binaries/mc cp test_data/blogs.parquet local/test-data
+    - name: Start roapi and Query
+      env:
+        AWS_DEFAULT_REGION: "us-east-1"
+        AWS_ACCESS_KEY_ID: minioadmin
+        AWS_SECRET_ACCESS_KEY: minioadmin
+        AWS_ENDPOINT_URL: http://127.0.0.1:9000
+      uses: BerniWittmann/background-server-action@v1.0.4
+      with:
+        command: /home/runner/work/roapi/roapi/test_end_to_end/query_blogs.sh
+        build: cargo build
+        start: cargo run --bin roapi -- -c test_end_to_end/test_s3_direct.yml
         wait-on: "http://127.0.0.1:8000/api/schema"
         # By default, wait-on will retry for 60 seconds. You can pass a custom timeout in seconds using wait-on-timeout.
         # 10 minutes = 600 seconds
diff --git a/Cargo.lock b/Cargo.lock
index 991a55768..a2181719c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -950,6 +950,7 @@ dependencies = [
  "hyper-tls",
  "lazy_static",
  "log",
+ "object_store",
  "pretty_assertions",
  "regex",
  "reqwest",
@@ -966,6 +967,7 @@ dependencies = [
  "thiserror",
  "tokio",
  "uriparse",
+ "url",
  "yup-oauth2",
 ]
 
diff --git a/columnq/Cargo.toml b/columnq/Cargo.toml
index 1d0073a5f..4edcd7ec8 100644
--- a/columnq/Cargo.toml
+++ b/columnq/Cargo.toml
@@ -16,6 +16,8 @@ path = "src/lib.rs"
 arrow-schema = { version ="26", features = ["serde"] }
 
 datafusion = "14"
+object_store = "0.5.0"
+url = "2.2"
 
 log = "0"
 regex = "1"
diff --git a/columnq/src/columnq.rs b/columnq/src/columnq.rs
index 5f6c605cd..ec3f5faef 100644
--- a/columnq/src/columnq.rs
+++ b/columnq/src/columnq.rs
@@ -5,13 +5,48 @@ use std::sync::Arc;
 use datafusion::arrow;
 use datafusion::arrow::array::as_string_array;
 use datafusion::arrow::array::StringArray;
+use datafusion::datasource::object_store::{ObjectStoreRegistry, ObjectStoreProvider};
+use datafusion::error::{DataFusionError, Result as DatafusionResult};
 pub use datafusion::execution::context::SessionConfig;
 use datafusion::execution::context::SessionContext;
+use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv};
 use datafusion::physical_plan::collect;
 
+use object_store::aws::AmazonS3Builder;
 use crate::error::{ColumnQError, QueryError};
 use crate::query;
 use crate::table::{self, KeyValueSource, TableSource};
+use url::Url;
+
+pub struct ColumnQObjectStoreProvider {}
+impl ObjectStoreProvider for ColumnQObjectStoreProvider {
+    fn get_by_url(&self, url: &Url) -> DatafusionResult<Arc<dyn object_store::ObjectStore>> {
+        let url_schema = url.scheme();
+        match url_schema {
+            "s3" => {
+                let host = url.host_str().unwrap();
+                let mut s3_builder = AmazonS3Builder::from_env().with_bucket_name(host);
+                // for minio in CI
+                s3_builder = s3_builder.with_allow_http(true);
+
+                // TODO: can remove "with_endpoint" after object_store upgrade to 0.5.3
+                // https://github.com/apache/arrow-datafusion/pull/4929
+                if let Ok(endpoint) = std::env::var("AWS_ENDPOINT_URL") {
+                    s3_builder = s3_builder.with_endpoint(endpoint);
+                }
+
+                match s3_builder.build() {
+                    Ok(s3) => Ok(Arc::new(s3)),
+                    Err(err) => Err(DataFusionError::External(Box::new(err))),
+                }
+            },
+            _ => Err(DataFusionError::Execution(format!(
+                "Unsupported object store scheme {}",
+                url_schema
+            ))),
+        }
+    }
+}
 
 pub struct ColumnQ {
     pub dfctx: SessionContext,
@@ -25,7 +60,12 @@ impl ColumnQ {
     }
 
     pub fn new_with_config(config: SessionConfig) -> Self {
-        let dfctx = SessionContext::with_config(config);
+        let object_store_provider = ColumnQObjectStoreProvider {};
+        let object_store_registry = ObjectStoreRegistry::new_with_provider(Some(Arc::new(object_store_provider)));
+        let rn_config = RuntimeConfig::new().with_object_store_registry(Arc::new(object_store_registry));
+        let runtime_env = RuntimeEnv::new(rn_config).unwrap();
+        let dfctx = SessionContext::with_config_rt(config, Arc::new(runtime_env));
+
         let schema_map = HashMap::<String, arrow::datatypes::SchemaRef>::new();
         Self {
             dfctx,
@@ -35,7 +75,7 @@ impl ColumnQ {
     }
 
     pub async fn load_table(&mut self, t: &TableSource) -> Result<(), ColumnQError> {
-        let table = table::load(t).await?;
+        let table = table::load(t, &self.dfctx).await?;
         self.schema_map.insert(t.name.clone(), table.schema());
         self.dfctx.deregister_table(t.name.as_str())?;
         self.dfctx.register_table(t.name.as_str(), table)?;
@@ -48,7 +88,7 @@ impl ColumnQ {
 
         let kv_entry = self.kv_catalog.entry(kv.name.clone());
         let (key, value) = (kv.key.clone(), kv.value.clone());
-        let table = table::load(&kv.into()).await?;
+        let table = table::load(&kv.into(), &self.dfctx).await?;
         let schema = table.schema();
         let key_schema_idx = schema.index_of(&key)?;
         if schema.field(key_schema_idx).data_type() != &DataType::Utf8 {
@@ -141,3 +181,58 @@ impl Default for ColumnQ {
         Self::new()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::{env, str::FromStr};
+
+    use datafusion::datasource::object_store::ObjectStoreProvider;
+    use url::Url;
+
+    use super::ColumnQObjectStoreProvider;
+
+    #[test]
+    fn s3_object_store_type() {
+        let host_url = "s3://bucket_name/path";
+        let provider = ColumnQObjectStoreProvider {};
+
+        let err = provider
+            .get_by_url(&Url::from_str(host_url).unwrap())
+            .unwrap_err();
+        assert!(err.to_string().contains("Generic S3 error: Missing region"));
+
+        env::set_var("AWS_REGION", "us-east-1");
+        let res = provider
+            .get_by_url(&Url::from_str(host_url).unwrap());
+        let msg = match res {
+            Err(e) => format!("{}", e),
+            Ok(_) => "".to_string(),
+        };
+        assert_eq!("".to_string(), msg);
+        env::remove_var("AWS_REGION");
+    }
+
+    #[test]
+    fn gcs_object_store_type() {
+        let host_url = "gs://bucket_name/path";
+        let provider = ColumnQObjectStoreProvider {};
+        let err = provider
+            .get_by_url(&Url::from_str(host_url).unwrap())
+            .unwrap_err();
+        assert!(err
+            .to_string()
+            .contains("Unsupported object store scheme gs"))
+    }
+
+    #[test]
+    fn azure_object_store_type() {
+        let unknown = "az://bucket_name/path";
+        let provider = ColumnQObjectStoreProvider {};
+        let err = provider
+            .get_by_url(&Url::from_str(unknown).unwrap())
+            .unwrap_err();
+        assert!(err
+            .to_string()
+            .contains("Unsupported object store scheme az"))
+    }
+}
diff --git a/columnq/src/table/mod.rs b/columnq/src/table/mod.rs
index 1c43ae902..91f4713a4 100644
--- a/columnq/src/table/mod.rs
+++ b/columnq/src/table/mod.rs
@@ -473,7 +473,7 @@ impl TableSource {
     }
 }
 
-pub async fn load(t: &TableSource) -> Result<Arc<dyn TableProvider>, ColumnQError> {
+pub async fn load(t: &TableSource, dfctx: &datafusion::execution::context::SessionContext) -> Result<Arc<dyn TableProvider>, ColumnQError> {
     if let Some(opt) = &t.option {
         Ok(match opt {
             TableLoadOption::json { .. } => Arc::new(json::to_mem_table(t).await?),
@@ -481,7 +481,7 @@ pub async fn load(t: &TableSource) -> Result<Arc<dyn TableProvider>, ColumnQErro
                 Arc::new(ndjson::to_mem_table(t).await?)
             }
             TableLoadOption::csv { .. } => csv::to_datafusion_table(t).await?,
-            TableLoadOption::parquet { .. } => parquet::to_datafusion_table(t).await?,
+            TableLoadOption::parquet { .. } => parquet::to_datafusion_table(t, dfctx).await?,
             TableLoadOption::google_spreadsheet(_) => {
                 Arc::new(google_spreadsheets::to_mem_table(t).await?)
             }
@@ -504,7 +504,7 @@ pub async fn load(t: &TableSource) -> Result<Arc<dyn TableProvider>, ColumnQErro
             "csv" => csv::to_datafusion_table(t).await?,
             "json" => Arc::new(json::to_mem_table(t).await?),
             "ndjson" | "jsonl" => Arc::new(ndjson::to_mem_table(t).await?),
-            "parquet" => parquet::to_datafusion_table(t).await?,
+            "parquet" => parquet::to_datafusion_table(t, dfctx).await?,
             "arrow" => Arc::new(arrow_ipc_file::to_mem_table(t).await?),
             "arrows" => Arc::new(arrow_ipc_stream::to_mem_table(t).await?),
             "mysql" => Arc::new(database::DatabaseLoader::MySQL.to_mem_table(t)?),
@@ -704,8 +704,8 @@ batch_size: 512
     #[tokio::test]
     async fn test_load_sqlite_table() -> anyhow::Result<()> {
         let t = TableSource::new("uk_cities", "sqlite://../test_data/sqlite/sample.db");
-        let table = load(&t).await?;
         let ctx = datafusion::prelude::SessionContext::new();
+        let table = load(&t, &ctx).await?;
         let stats = table
             .scan(&ctx.state(), &None, &[], None)
             .await?
@@ -726,8 +726,8 @@ uri: "sqlite://../test_data/sqlite/sample.{}"
 "#,
                 ext
             ))?;
-            let table = load(&t).await?;
             let ctx = datafusion::prelude::SessionContext::new();
+            let table = load(&t, &ctx).await?;
             let stats = table
                 .scan(&ctx.state(), &None, &[], None)
                 .await?
diff --git a/columnq/src/table/parquet.rs b/columnq/src/table/parquet.rs
index 078851df5..729d07ee4 100644
--- a/columnq/src/table/parquet.rs
+++ b/columnq/src/table/parquet.rs
@@ -15,9 +15,9 @@ use datafusion::datasource::listing::{
 use datafusion::datasource::TableProvider;
 use datafusion::parquet::arrow::arrow_reader::ArrowReaderOptions;
 use datafusion::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
-use datafusion::prelude::SessionContext;
 
-pub async fn to_datafusion_table(t: &TableSource) -> Result<Arc<dyn TableProvider>, ColumnQError> {
+
+pub async fn to_datafusion_table(t: &TableSource, dfctx: &datafusion::execution::context::SessionContext) -> Result<Arc<dyn TableProvider>, ColumnQError> {
     let opt = t
         .option
         .clone()
@@ -32,8 +32,7 @@ pub async fn to_datafusion_table(t: &TableSource) -> Result<Arc<dyn TableProvide
         let schemaref = match &t.schema {
             Some(s) => Arc::new(s.into()),
             None => {
-                let ctx = SessionContext::new();
-                options.infer_schema(&ctx.state(), &table_url).await?
+                options.infer_schema(&dfctx.state(), &table_url).await?
             }
         };
 
@@ -107,6 +106,7 @@ mod tests {
 
     #[tokio::test]
     async fn load_flattened_parquet() {
+        let ctx = SessionContext::new();
         let t = to_datafusion_table(
             &TableSource::new(
                 "blogs".to_string(),
@@ -115,11 +115,11 @@ mod tests {
             .with_option(TableLoadOption::parquet(TableOptionParquet {
                 use_memory_table: false,
             })),
+            &ctx
         )
         .await
         .unwrap();
 
-        let ctx = SessionContext::new();
         let stats = t
             .scan(&ctx.state(), &None, &[], None)
             .await


Test patch:
diff --git a/columnq/src/test_util.rs b/columnq/src/test_util.rs
index 95fd2591b..78faa0e92 100644
--- a/columnq/src/test_util.rs
+++ b/columnq/src/test_util.rs
@@ -121,8 +121,8 @@ schema:
 
     // patch uri path with the correct test data path
     table_source.io_source = table::TableIoSource::Uri(test_data_path("ubuntu-ami.json"));
-
-    Ok(table::load(&table_source).await?)
+    let ctx = SessionContext::new();
+    Ok(table::load(&table_source, &ctx).await?)
 }
 
 pub fn register_table_properties(dfctx: &mut SessionContext) -> anyhow::Result<()> {
diff --git a/test_end_to_end/test_s3_direct.yml b/test_end_to_end/test_s3_direct.yml
new file mode 100644
index 000000000..3f19c7dc4
--- /dev/null
+++ b/test_end_to_end/test_s3_direct.yml
@@ -0,0 +1,12 @@
+addr:
+  # binding address for TCP port that speaks HTTP protocol
+  http: 0.0.0.0:8000
+  # binding address for TCP port that speaks Postgres wire protocol
+  postgres: 0.0.0.0:5432
+tables:
+  - name: "blogs"
+    uri: "s3://test-data/blogs.parquet"
+    option:
+      format: "parquet"
+      # following line is different from test_end_to_end/test_s3_memory.yml
+      use_memory_table: false
diff --git a/test_end_to_end/test_s3.yml b/test_end_to_end/test_s3_memory.yml
similarity index 100%
rename from test_end_to_end/test_s3.yml
rename to test_end_to_end/test_s3_memory.yml