You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by al...@apache.org on 2023/06/07 11:39:18 UTC

[arrow-datafusion] branch main updated: Add tests for object_store builders of datafusion-cli (#6576)

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/main by this push:
     new 25b6556c64 Add tests for object_store builders of datafusion-cli (#6576)
25b6556c64 is described below

commit 25b6556c64d3afe9c14625bf0682c98d5b1359f6
Author: r.4ntix <r....@gmail.com>
AuthorDate: Wed Jun 7 19:39:10 2023 +0800

    Add tests for object_store builders of datafusion-cli (#6576)
---
 datafusion-cli/src/exec.rs           |   4 +-
 datafusion-cli/src/object_storage.rs | 104 +++++++++++++++++++++++++++++++----
 2 files changed, 95 insertions(+), 13 deletions(-)

diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs
index 15d6d81b6d..cec0fe0373 100644
--- a/datafusion-cli/src/exec.rs
+++ b/datafusion-cli/src/exec.rs
@@ -248,7 +248,7 @@ mod tests {
 
         match &plan {
             LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) => {
-                create_external_table(&ctx, cmd)?;
+                create_external_table(&ctx, cmd).await?;
             }
             _ => assert!(false),
         };
@@ -327,7 +327,7 @@ mod tests {
         let err = create_external_table_test(location, &sql)
             .await
             .unwrap_err();
-        assert!(err.to_string().contains("A configuration file was passed"));
+        assert!(err.to_string().contains("No such file or directory"));
 
         Ok(())
     }
diff --git a/datafusion-cli/src/object_storage.rs b/datafusion-cli/src/object_storage.rs
index 7b0968d402..46b03a0a36 100644
--- a/datafusion-cli/src/object_storage.rs
+++ b/datafusion-cli/src/object_storage.rs
@@ -165,28 +165,110 @@ mod tests {
         logical_expr::{DdlStatement, LogicalPlan},
         prelude::SessionContext,
     };
+    use object_store::{aws::AmazonS3ConfigKey, gcp::GoogleConfigKey};
 
     use super::*;
 
-    #[ignore] // https://github.com/apache/arrow-rs/issues/4021
     #[tokio::test]
-    async fn oss_object_store_builder() -> Result<()> {
-        let access_key_id = "access_key_id";
-        let secret_access_key = "secret_access_key";
-        let region = "us-east-2";
+    async fn s3_object_store_builder() -> Result<()> {
+        let access_key_id = "fake_access_key_id";
+        let secret_access_key = "fake_secret_access_key";
+        let region = "fake_us-east-2";
+        let session_token = "fake_session_token";
         let location = "s3://bucket/path/file.parquet";
+
+        let table_url = ListingTableUrl::parse(location)?;
+        let sql = format!("CREATE EXTERNAL TABLE test STORED AS PARQUET OPTIONS('access_key_id' '{access_key_id}', 'secret_access_key' '{secret_access_key}', 'region' '{region}', 'session_token' {session_token}) LOCATION '{location}'");
+
+        let ctx = SessionContext::new();
+        let plan = ctx.state().create_logical_plan(&sql).await?;
+
+        if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &plan {
+            let builder = get_s3_object_store_builder(table_url.as_ref(), cmd).await?;
+            // get the actual configuration information, then assert_eq!
+            let config = [
+                (AmazonS3ConfigKey::AccessKeyId, access_key_id),
+                (AmazonS3ConfigKey::SecretAccessKey, secret_access_key),
+                (AmazonS3ConfigKey::Region, region),
+                (AmazonS3ConfigKey::Token, session_token),
+            ];
+            for (key, value) in config {
+                assert_eq!(value, builder.get_config_value(&key).unwrap());
+            }
+        } else {
+            return Err(DataFusionError::Plan(
+                "LogicalPlan is not a CreateExternalTable".to_string(),
+            ));
+        }
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn oss_object_store_builder() -> Result<()> {
+        let access_key_id = "fake_access_key_id";
+        let secret_access_key = "fake_secret_access_key";
+        let endpoint = "fake_endpoint";
+        let location = "oss://bucket/path/file.parquet";
+
+        let table_url = ListingTableUrl::parse(location)?;
+        let sql = format!("CREATE EXTERNAL TABLE test STORED AS PARQUET OPTIONS('access_key_id' '{access_key_id}', 'secret_access_key' '{secret_access_key}', 'endpoint' '{endpoint}') LOCATION '{location}'");
+
+        let ctx = SessionContext::new();
+        let plan = ctx.state().create_logical_plan(&sql).await?;
+
+        if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &plan {
+            let builder = get_oss_object_store_builder(table_url.as_ref(), cmd)?;
+            // get the actual configuration information, then assert_eq!
+            let config = [
+                (AmazonS3ConfigKey::AccessKeyId, access_key_id),
+                (AmazonS3ConfigKey::SecretAccessKey, secret_access_key),
+                (AmazonS3ConfigKey::Endpoint, endpoint),
+            ];
+            for (key, value) in config {
+                assert_eq!(value, builder.get_config_value(&key).unwrap());
+            }
+        } else {
+            return Err(DataFusionError::Plan(
+                "LogicalPlan is not a CreateExternalTable".to_string(),
+            ));
+        }
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn gcs_object_store_builder() -> Result<()> {
+        let service_account_path = "fake_service_account_path";
+        let service_account_key =
+            "{\"private_key\": \"fake_private_key.pem\",\"client_email\":\"fake_client_email\"}";
+        let application_credentials_path = "fake_application_credentials_path";
+        let location = "gcs://bucket/path/file.parquet";
+
         let table_url = ListingTableUrl::parse(location)?;
-        let sql = format!("CREATE EXTERNAL TABLE test STORED AS PARQUET OPTIONS('access_key_id' '{access_key_id}', 'secret_access_key' '{secret_access_key}', 'region' '{region}') LOCATION '{location}'");
+        let sql = format!("CREATE EXTERNAL TABLE test STORED AS PARQUET OPTIONS('service_account_path' '{service_account_path}', 'service_account_key' '{service_account_key}', 'application_credentials_path' '{application_credentials_path}') LOCATION '{location}'");
 
         let ctx = SessionContext::new();
         let plan = ctx.state().create_logical_plan(&sql).await?;
 
-        match &plan {
-            LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) => {
-                let _builder = get_oss_object_store_builder(table_url.as_ref(), cmd)?;
-                // get the actual configuration information, then assert_eq!
+        if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &plan {
+            let builder = get_gcs_object_store_builder(table_url.as_ref(), cmd)?;
+            // get the actual configuration information, then assert_eq!
+            let config = [
+                (GoogleConfigKey::ServiceAccount, service_account_path),
+                (GoogleConfigKey::ServiceAccountKey, service_account_key),
+                (
+                    GoogleConfigKey::ApplicationCredentials,
+                    application_credentials_path,
+                ),
+            ];
+            for (key, value) in config {
+                assert_eq!(value, builder.get_config_value(&key).unwrap());
             }
-            _ => assert!(false),
+        } else {
+            return Err(DataFusionError::Plan(
+                "LogicalPlan is not a CreateExternalTable".to_string(),
+            ));
         }
 
         Ok(())