From b17947b5d8027ed2b731320cf8cb1c72d7c5895b Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Wed, 4 Mar 2026 12:21:21 +0800 Subject: [PATCH] feat: add COPY support for GCS/Azblob with OSS/GCS/Azblob integration coverage (#7743) * feat: add gcs and azblob backends to copy datasource Signed-off-by: WenyXu * test: add copy to/from oss integration coverage Signed-off-by: WenyXu * test: add copy to/from gcs/azure integration coverage Signed-off-by: WenyXu * refactor: use existing layers Signed-off-by: WenyXu * feat: hide `credential_path` in cli Signed-off-by: WenyXu * chore: remove logs Signed-off-by: WenyXu * chore: apply suggestions Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/cli/src/common/object_store.rs | 4 +- src/cli/src/data/export.rs | 8 +- src/cli/src/data/storage_export.rs | 6 +- src/common/datasource/src/object_store.rs | 18 + .../datasource/src/object_store/azblob.rs | 77 +++ src/common/datasource/src/object_store/fs.rs | 9 +- src/common/datasource/src/object_store/gcs.rs | 71 +++ src/common/datasource/src/object_store/oss.rs | 17 +- src/common/datasource/src/object_store/s3.rs | 21 +- tests-integration/src/tests/instance_test.rs | 462 +++++++++++++++++- 10 files changed, 646 insertions(+), 47 deletions(-) create mode 100644 src/common/datasource/src/object_store/azblob.rs create mode 100644 src/common/datasource/src/object_store/gcs.rs diff --git a/src/cli/src/common/object_store.rs b/src/cli/src/common/object_store.rs index e1d3a189c9a..0b8372e509f 100644 --- a/src/cli/src/common/object_store.rs +++ b/src/cli/src/common/object_store.rs @@ -85,7 +85,7 @@ impl FieldValidator for Option { macro_rules! wrap_with_clap_prefix { ( $new_name:ident, $prefix:literal, $enable_flag:literal, $base:ty, { - $( $( #[doc = $doc:expr] )? $( #[alias = $alias:literal] )? $field:ident : $type:ty $( = $default:expr )? ),* $(,)? + $( $( #[doc = $doc:expr] )? $( #[alias = $alias:literal] )? $( #[hide = $hide:literal] )? $field:ident : $type:ty $( = $default:expr )? ),* $(,)? } ) => { paste!{ @@ -94,6 +94,7 @@ macro_rules! wrap_with_clap_prefix { $( $( #[doc = $doc] )? $( #[clap(alias = $alias)] )? + $( #[clap(hide = $hide)] )? #[clap(long, requires = $enable_flag $(, default_value_t = $default )? )] pub [<$prefix $field>]: $type, )* @@ -320,6 +321,7 @@ wrap_with_clap_prefix! { #[doc = "The scope of the object store."] scope: String = Default::default(), #[doc = "The credential path of the object store."] + #[hide = true] credential_path: Option, #[doc = "The credential of the object store."] credential: Option, diff --git a/src/cli/src/data/export.rs b/src/cli/src/data/export.rs index a5e6ee82eb3..1cdb159336a 100644 --- a/src/cli/src/data/export.rs +++ b/src/cli/src/data/export.rs @@ -724,8 +724,6 @@ mod tests { "test-root", "--gcs-scope", "test-scope", - "--gcs-credential-path", - "/path/to/credential", "--gcs-credential", "test-credential-content", "--gcs-endpoint", @@ -750,7 +748,7 @@ mod tests { "test-root", "--gcs-scope", "test-scope", - // No credential_path or credential + // No credential // No endpoint (optional) ]); @@ -1061,8 +1059,6 @@ mod tests { "", // Empty root "--gcs-scope", "test-scope", - "--gcs-credential-path", - "/path/to/credential", "--gcs-credential", "test-credential", "--gcs-endpoint", @@ -1219,7 +1215,7 @@ mod tests { "test-root", "--gcs-scope", "test-scope", - // No credential_path, credential, or endpoint + // No credential, or endpoint ]); let result = cmd.build().await; diff --git a/src/cli/src/data/storage_export.rs b/src/cli/src/data/storage_export.rs index acec2355981..01bf252bbf5 100644 --- a/src/cli/src/data/storage_export.rs +++ b/src/cli/src/data/storage_export.rs @@ -230,7 +230,6 @@ impl StorageExport for GcsBackend { ); let mut connection_options = Vec::new(); - let credential_path = expose_optional_secret(&self.config.gcs_credential_path); if !credential_path.is_empty() { connection_options.push(format!("CREDENTIAL_PATH='{}'", credential_path)); @@ -266,10 +265,7 @@ impl StorageExport for GcsBackend { fn mask_sensitive_info(&self, sql: &str) -> String { mask_secrets( sql.to_string(), - &[ - expose_optional_secret(&self.config.gcs_credential_path), - expose_optional_secret(&self.config.gcs_credential), - ], + &[expose_optional_secret(&self.config.gcs_credential)], ) } } diff --git a/src/common/datasource/src/object_store.rs b/src/common/datasource/src/object_store.rs index 942df20336c..56eaf3968bf 100644 --- a/src/common/datasource/src/object_store.rs +++ b/src/common/datasource/src/object_store.rs @@ -12,7 +12,9 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub mod azblob; pub mod fs; +pub mod gcs; pub mod oss; pub mod s3; @@ -24,7 +26,9 @@ use regex::Regex; use snafu::{OptionExt, ResultExt}; use url::{ParseError, Url}; +use self::azblob::build_azblob_backend; use self::fs::build_fs_backend; +use self::gcs::build_gcs_backend; use self::s3::build_s3_backend; use crate::error::{self, Result}; use crate::object_store::oss::build_oss_backend; @@ -33,6 +37,8 @@ use crate::util::find_dir_and_filename; pub const FS_SCHEMA: &str = "FS"; pub const S3_SCHEMA: &str = "S3"; pub const OSS_SCHEMA: &str = "OSS"; +pub const GCS_SCHEMA: &str = "GCS"; +pub const AZBLOB_SCHEMA: &str = "AZBLOB"; /// Returns `(schema, Option, path)` pub fn parse_url(url: &str) -> Result<(String, Option, String)> { @@ -74,6 +80,18 @@ pub fn build_backend(url: &str, connection: &HashMap) -> Result< })?; Ok(build_oss_backend(&host, &root, connection)?) } + GCS_SCHEMA => { + let host = host.context(error::EmptyHostPathSnafu { + url: url.to_string(), + })?; + Ok(build_gcs_backend(&host, &root, connection)?) + } + AZBLOB_SCHEMA => { + let host = host.context(error::EmptyHostPathSnafu { + url: url.to_string(), + })?; + Ok(build_azblob_backend(&host, &root, connection)?) + } FS_SCHEMA => Ok(build_fs_backend(&root)?), _ => error::UnsupportedBackendProtocolSnafu { diff --git a/src/common/datasource/src/object_store/azblob.rs b/src/common/datasource/src/object_store/azblob.rs new file mode 100644 index 00000000000..ea6c19ca28e --- /dev/null +++ b/src/common/datasource/src/object_store/azblob.rs @@ -0,0 +1,77 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; + +use object_store::ObjectStore; +use object_store::services::Azblob; +use object_store::util::{with_instrument_layers, with_retry_layers}; +use snafu::ResultExt; + +use crate::error::{self, Result}; + +const ENDPOINT: &str = "endpoint"; +const ACCOUNT_NAME: &str = "account_name"; +const ACCOUNT_KEY: &str = "account_key"; +const SAS_TOKEN: &str = "sas_token"; + +pub fn is_supported_in_azblob(key: &str) -> bool { + [ENDPOINT, ACCOUNT_NAME, ACCOUNT_KEY, SAS_TOKEN].contains(&key) +} + +pub fn build_azblob_backend( + container: &str, + root: &str, + connection: &HashMap, +) -> Result { + let mut builder = Azblob::default().root(root).container(container); + + if let Some(account_name) = connection.get(ACCOUNT_NAME) { + builder = builder.account_name(account_name); + } + + if let Some(account_key) = connection.get(ACCOUNT_KEY) { + builder = builder.account_key(account_key); + } + + if let Some(endpoint) = connection.get(ENDPOINT) { + builder = builder.endpoint(endpoint); + } + + if let Some(sas_token) = connection.get(SAS_TOKEN) { + builder = builder.sas_token(sas_token); + } + + let object_store = ObjectStore::new(builder) + .context(error::BuildBackendSnafu)? + .finish(); + Ok(with_instrument_layers( + with_retry_layers(object_store), + true, + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_is_supported_in_azblob() { + assert!(is_supported_in_azblob(ENDPOINT)); + assert!(is_supported_in_azblob(ACCOUNT_NAME)); + assert!(is_supported_in_azblob(ACCOUNT_KEY)); + assert!(is_supported_in_azblob(SAS_TOKEN)); + assert!(!is_supported_in_azblob("foo")); + } +} diff --git a/src/common/datasource/src/object_store/fs.rs b/src/common/datasource/src/object_store/fs.rs index 95e5a691442..0f537b1e693 100644 --- a/src/common/datasource/src/object_store/fs.rs +++ b/src/common/datasource/src/object_store/fs.rs @@ -14,7 +14,7 @@ use object_store::ObjectStore; use object_store::services::Fs; -use object_store::util::DefaultLoggingInterceptor; +use object_store::util::with_instrument_layers; use snafu::ResultExt; use crate::error::{BuildBackendSnafu, Result}; @@ -23,11 +23,6 @@ pub fn build_fs_backend(root: &str) -> Result { let builder = Fs::default(); let object_store = ObjectStore::new(builder.root(root)) .context(BuildBackendSnafu)? - .layer(object_store::layers::LoggingLayer::new( - DefaultLoggingInterceptor, - )) - .layer(object_store::layers::TracingLayer) - .layer(object_store::layers::build_prometheus_metrics_layer(true)) .finish(); - Ok(object_store) + Ok(with_instrument_layers(object_store, true)) } diff --git a/src/common/datasource/src/object_store/gcs.rs b/src/common/datasource/src/object_store/gcs.rs new file mode 100644 index 00000000000..db812851941 --- /dev/null +++ b/src/common/datasource/src/object_store/gcs.rs @@ -0,0 +1,71 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; + +use object_store::ObjectStore; +use object_store::services::Gcs; +use object_store::util::{with_instrument_layers, with_retry_layers}; +use snafu::ResultExt; + +use crate::error::{self, Result}; + +const ENDPOINT: &str = "endpoint"; +const CREDENTIAL: &str = "credential"; +const SCOPE: &str = "scope"; + +pub fn is_supported_in_gcs(key: &str) -> bool { + [ENDPOINT, CREDENTIAL, SCOPE].contains(&key) +} + +pub fn build_gcs_backend( + bucket: &str, + root: &str, + connection: &HashMap, +) -> Result { + let mut builder = Gcs::default().root(root).bucket(bucket); + + if let Some(scope) = connection.get(SCOPE) { + builder = builder.scope(scope); + } + + if let Some(credential) = connection.get(CREDENTIAL) { + builder = builder.credential(credential); + } + + if let Some(endpoint) = connection.get(ENDPOINT) { + builder = builder.endpoint(endpoint); + } + + let object_store = ObjectStore::new(builder) + .context(error::BuildBackendSnafu)? + .finish(); + Ok(with_instrument_layers( + with_retry_layers(object_store), + true, + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_is_supported_in_gcs() { + assert!(is_supported_in_gcs(ENDPOINT)); + assert!(is_supported_in_gcs(CREDENTIAL)); + assert!(is_supported_in_gcs(SCOPE)); + assert!(!is_supported_in_gcs("foo")); + } +} diff --git a/src/common/datasource/src/object_store/oss.rs b/src/common/datasource/src/object_store/oss.rs index 47cd9b489f0..aded3eca2cd 100644 --- a/src/common/datasource/src/object_store/oss.rs +++ b/src/common/datasource/src/object_store/oss.rs @@ -16,6 +16,7 @@ use std::collections::HashMap; use object_store::ObjectStore; use object_store::services::Oss; +use object_store::util::{with_instrument_layers, with_retry_layers}; use snafu::ResultExt; use crate::error::{self, Result}; @@ -75,19 +76,13 @@ pub fn build_oss_backend( } } - let op = ObjectStore::new(builder) + let object_store = ObjectStore::new(builder) .context(error::BuildBackendSnafu)? - .layer( - object_store::layers::RetryLayer::new() - .with_jitter() - .with_notify(object_store::util::PrintDetailedError), - ) - .layer(object_store::layers::LoggingLayer::default()) - .layer(object_store::layers::TracingLayer) - .layer(object_store::layers::build_prometheus_metrics_layer(true)) .finish(); - - Ok(op) + Ok(with_instrument_layers( + with_retry_layers(object_store), + true, + )) } #[cfg(test)] diff --git a/src/common/datasource/src/object_store/s3.rs b/src/common/datasource/src/object_store/s3.rs index 36762ff0acb..5570185a003 100644 --- a/src/common/datasource/src/object_store/s3.rs +++ b/src/common/datasource/src/object_store/s3.rs @@ -16,7 +16,7 @@ use std::collections::HashMap; use object_store::ObjectStore; use object_store::services::S3; -use object_store::util::DefaultLoggingInterceptor; +use object_store::util::{with_instrument_layers, with_retry_layers}; use snafu::ResultExt; use crate::error::{self, Result}; @@ -99,20 +99,13 @@ pub fn build_s3_backend( } } - // TODO(weny): Consider finding a better way to eliminate duplicate code. - Ok(ObjectStore::new(builder) + let object_store = ObjectStore::new(builder) .context(error::BuildBackendSnafu)? - .layer( - object_store::layers::RetryLayer::new() - .with_jitter() - .with_notify(object_store::util::PrintDetailedError), - ) - .layer(object_store::layers::LoggingLayer::new( - DefaultLoggingInterceptor, - )) - .layer(object_store::layers::TracingLayer) - .layer(object_store::layers::build_prometheus_metrics_layer(true)) - .finish()) + .finish(); + Ok(with_instrument_layers( + with_retry_layers(object_store), + true, + )) } #[cfg(test)] diff --git a/tests-integration/src/tests/instance_test.rs b/tests-integration/src/tests/instance_test.rs index 8c2967bc2ab..5879ae9dfab 100644 --- a/tests-integration/src/tests/instance_test.rs +++ b/tests-integration/src/tests/instance_test.rs @@ -1672,8 +1672,6 @@ async fn test_execute_copy_to_s3(instance: Arc) { #[apply(both_instances_cases)] async fn test_execute_copy_from_s3(instance: Arc) { - use common_telemetry::info; - common_telemetry::init_default_ut_logging(); if let Ok(bucket) = env::var("GT_S3_BUCKET") && !bucket.is_empty() @@ -1755,8 +1753,466 @@ async fn test_execute_copy_from_s3(instance: Arc) { "{} CONNECTION (ACCESS_KEY_ID='{}',SECRET_ACCESS_KEY='{}',REGION='{}')", test.sql, key_id, key, region, ); - info!("Running sql: {}", sql); + let output = execute_sql(&instance, &sql).await.data; + assert!(matches!(output, OutputData::AffectedRows(2))); + let output = execute_sql( + &instance, + &format!("select * from {} order by ts", test.table_name), + ) + .await + .data; + let expected = "\ ++-------+------+--------+---------------------+ +| host | cpu | memory | ts | ++-------+------+--------+---------------------+ +| host1 | 66.6 | 1024.0 | 2022-06-15T07:02:37 | +| host2 | 88.8 | 333.3 | 2022-06-15T07:02:38 | ++-------+------+--------+---------------------+"; + check_output_stream(output, expected).await; + } + } +} + +#[apply(both_instances_cases)] +async fn test_execute_copy_to_oss(instance: Arc) { + if let Ok(bucket) = env::var("GT_OSS_BUCKET") + && !bucket.is_empty() + { + let instance = instance.frontend(); + + assert!(matches!( + execute_sql( + &instance, + "create table demo(host string, cpu double, memory double, ts timestamp time index);", + ) + .await + .data, + OutputData::AffectedRows(0) + )); + + let output = execute_sql( + &instance, + r#"insert into demo(host, cpu, memory, ts) values + ('host1', 66.6, 1024, 1655276557000), + ('host2', 88.8, 333.3, 1655276558000) + "#, + ) + .await + .data; + assert!(matches!(output, OutputData::AffectedRows(2))); + + let key_id = env::var("GT_OSS_ACCESS_KEY_ID").unwrap(); + let key = env::var("GT_OSS_ACCESS_KEY").unwrap(); + let endpoint = env::var("GT_OSS_ENDPOINT").unwrap(); + + let root = uuid::Uuid::new_v4().to_string(); + + let copy_to_stmt = format!( + "Copy demo TO 'oss://{}/{}/export/demo.parquet' CONNECTION (ACCESS_KEY_ID='{}',ACCESS_KEY_SECRET='{}',ENDPOINT='{}')", + bucket, root, key_id, key, endpoint + ); + + let output = execute_sql(&instance, ©_to_stmt).await.data; + assert!(matches!(output, OutputData::AffectedRows(2))); + } +} + +#[apply(both_instances_cases)] +async fn test_execute_copy_from_oss(instance: Arc) { + common_telemetry::init_default_ut_logging(); + if let Ok(bucket) = env::var("GT_OSS_BUCKET") + && !bucket.is_empty() + { + let instance = instance.frontend(); + + assert!(matches!( + execute_sql( + &instance, + "create table demo(host string, cpu double, memory double, ts timestamp time index);", + ) + .await + .data, + OutputData::AffectedRows(0) + )); + + let output = execute_sql( + &instance, + r#"insert into demo(host, cpu, memory, ts) values + ('host1', 66.6, 1024, 1655276557000), + ('host2', 88.8, 333.3, 1655276558000) + "#, + ) + .await + .data; + assert!(matches!(output, OutputData::AffectedRows(2))); + + let root = uuid::Uuid::new_v4().to_string(); + let key_id = env::var("GT_OSS_ACCESS_KEY_ID").unwrap(); + let key = env::var("GT_OSS_ACCESS_KEY").unwrap(); + let endpoint = env::var("GT_OSS_ENDPOINT").unwrap(); + + let copy_to_stmt = format!( + "Copy demo TO 'oss://{}/{}/export/demo.parquet' CONNECTION (ACCESS_KEY_ID='{}',ACCESS_KEY_SECRET='{}',ENDPOINT='{}')", + bucket, root, key_id, key, endpoint + ); + + let output = execute_sql(&instance, ©_to_stmt).await.data; + assert!(matches!(output, OutputData::AffectedRows(2))); + + struct Test<'a> { + sql: &'a str, + table_name: &'a str, + } + let tests = [ + Test { + sql: &format!( + "Copy with_filename FROM 'oss://{}/{}/export/demo.parquet'", + bucket, root + ), + table_name: "with_filename", + }, + Test { + sql: &format!("Copy with_path FROM 'oss://{}/{}/export/'", bucket, root), + table_name: "with_path", + }, + Test { + sql: &format!( + "Copy with_pattern FROM 'oss://{}/{}/export/' WITH (PATTERN = 'demo.*')", + bucket, root + ), + table_name: "with_pattern", + }, + ]; + + for test in tests { + assert!(matches!( + execute_sql( + &instance, + &format!( + "create table {}(host string, cpu double, memory double, ts timestamp time index);", + test.table_name + ), + ) + .await + .data, + OutputData::AffectedRows(0) + )); + + let sql = format!( + "{} CONNECTION (ACCESS_KEY_ID='{}',ACCESS_KEY_SECRET='{}',ENDPOINT='{}')", + test.sql, key_id, key, endpoint, + ); + let output = execute_sql(&instance, &sql).await.data; + assert!(matches!(output, OutputData::AffectedRows(2))); + + let output = execute_sql( + &instance, + &format!("select * from {} order by ts", test.table_name), + ) + .await + .data; + let expected = "\ ++-------+------+--------+---------------------+ +| host | cpu | memory | ts | ++-------+------+--------+---------------------+ +| host1 | 66.6 | 1024.0 | 2022-06-15T07:02:37 | +| host2 | 88.8 | 333.3 | 2022-06-15T07:02:38 | ++-------+------+--------+---------------------+"; + check_output_stream(output, expected).await; + } + } +} + +#[apply(both_instances_cases)] +async fn test_execute_copy_to_gcs(instance: Arc) { + if let (Ok(bucket), Ok(scope), Ok(credential)) = ( + env::var("GT_GCS_BUCKET"), + env::var("GT_GCS_SCOPE"), + env::var("GT_GCS_CREDENTIAL"), + ) && !bucket.is_empty() + { + let endpoint = env::var("GT_GCS_ENDPOINT").unwrap_or_default(); + let instance = instance.frontend(); + + assert!(matches!( + execute_sql( + &instance, + "create table demo(host string, cpu double, memory double, ts timestamp time index);", + ) + .await + .data, + OutputData::AffectedRows(0) + )); + + let output = execute_sql( + &instance, + r#"insert into demo(host, cpu, memory, ts) values + ('host1', 66.6, 1024, 1655276557000), + ('host2', 88.8, 333.3, 1655276558000) + "#, + ) + .await + .data; + assert!(matches!(output, OutputData::AffectedRows(2))); + + let root = uuid::Uuid::new_v4().to_string(); + let mut connection = format!("SCOPE='{}',CREDENTIAL='{}'", scope, credential); + if !endpoint.is_empty() { + connection = format!("{},ENDPOINT='{}'", connection, endpoint); + } + + let copy_to_stmt = format!( + "Copy demo TO 'gcs://{}/{}/export/demo.parquet' CONNECTION ({})", + bucket, root, connection + ); + + let output = execute_sql(&instance, ©_to_stmt).await.data; + assert!(matches!(output, OutputData::AffectedRows(2))); + } +} + +#[apply(both_instances_cases)] +async fn test_execute_copy_from_gcs(instance: Arc) { + common_telemetry::init_default_ut_logging(); + if let (Ok(bucket), Ok(scope), Ok(credential)) = ( + env::var("GT_GCS_BUCKET"), + env::var("GT_GCS_SCOPE"), + env::var("GT_GCS_CREDENTIAL"), + ) && !bucket.is_empty() + { + let endpoint = env::var("GT_GCS_ENDPOINT").unwrap_or_default(); + let instance = instance.frontend(); + + assert!(matches!( + execute_sql( + &instance, + "create table demo(host string, cpu double, memory double, ts timestamp time index);", + ) + .await + .data, + OutputData::AffectedRows(0) + )); + + let output = execute_sql( + &instance, + r#"insert into demo(host, cpu, memory, ts) values + ('host1', 66.6, 1024, 1655276557000), + ('host2', 88.8, 333.3, 1655276558000) + "#, + ) + .await + .data; + assert!(matches!(output, OutputData::AffectedRows(2))); + + let root = uuid::Uuid::new_v4().to_string(); + let mut connection = format!("SCOPE='{}',CREDENTIAL='{}'", scope, credential); + if !endpoint.is_empty() { + connection = format!("{},ENDPOINT='{}'", connection, endpoint); + } + + let copy_to_stmt = format!( + "Copy demo TO 'gcs://{}/{}/export/demo.parquet' CONNECTION ({})", + bucket, root, connection + ); + + let output = execute_sql(&instance, ©_to_stmt).await.data; + assert!(matches!(output, OutputData::AffectedRows(2))); + + struct Test<'a> { + sql: &'a str, + table_name: &'a str, + } + let tests = [ + Test { + sql: &format!( + "Copy with_filename FROM 'gcs://{}/{}/export/demo.parquet'", + bucket, root + ), + table_name: "with_filename", + }, + Test { + sql: &format!("Copy with_path FROM 'gcs://{}/{}/export/'", bucket, root), + table_name: "with_path", + }, + Test { + sql: &format!( + "Copy with_pattern FROM 'gcs://{}/{}/export/' WITH (PATTERN = 'demo.*')", + bucket, root + ), + table_name: "with_pattern", + }, + ]; + + for test in tests { + assert!(matches!( + execute_sql( + &instance, + &format!( + "create table {}(host string, cpu double, memory double, ts timestamp time index);", + test.table_name + ), + ) + .await + .data, + OutputData::AffectedRows(0) + )); + + let sql = format!("{} CONNECTION ({})", test.sql, connection); + let output = execute_sql(&instance, &sql).await.data; + assert!(matches!(output, OutputData::AffectedRows(2))); + + let output = execute_sql( + &instance, + &format!("select * from {} order by ts", test.table_name), + ) + .await + .data; + let expected = "\ ++-------+------+--------+---------------------+ +| host | cpu | memory | ts | ++-------+------+--------+---------------------+ +| host1 | 66.6 | 1024.0 | 2022-06-15T07:02:37 | +| host2 | 88.8 | 333.3 | 2022-06-15T07:02:38 | ++-------+------+--------+---------------------+"; + check_output_stream(output, expected).await; + } + } +} + +#[apply(both_instances_cases)] +async fn test_execute_copy_to_azblob(instance: Arc) { + if let (Ok(container), Ok(account_name), Ok(account_key), Ok(endpoint)) = ( + env::var("GT_AZBLOB_CONTAINER"), + env::var("GT_AZBLOB_ACCOUNT_NAME"), + env::var("GT_AZBLOB_ACCOUNT_KEY"), + env::var("GT_AZBLOB_ENDPOINT"), + ) && !container.is_empty() + { + let instance = instance.frontend(); + + assert!(matches!( + execute_sql( + &instance, + "create table demo(host string, cpu double, memory double, ts timestamp time index);", + ) + .await + .data, + OutputData::AffectedRows(0) + )); + + let output = execute_sql( + &instance, + r#"insert into demo(host, cpu, memory, ts) values + ('host1', 66.6, 1024, 1655276557000), + ('host2', 88.8, 333.3, 1655276558000) + "#, + ) + .await + .data; + assert!(matches!(output, OutputData::AffectedRows(2))); + + let root = uuid::Uuid::new_v4().to_string(); + + let copy_to_stmt = format!( + "Copy demo TO 'azblob://{}/{}/export/demo.parquet' CONNECTION (ACCOUNT_NAME='{}',ACCOUNT_KEY='{}',ENDPOINT='{}')", + container, root, account_name, account_key, endpoint + ); + + let output = execute_sql(&instance, ©_to_stmt).await.data; + assert!(matches!(output, OutputData::AffectedRows(2))); + } +} + +#[apply(both_instances_cases)] +async fn test_execute_copy_from_azblob(instance: Arc) { + common_telemetry::init_default_ut_logging(); + if let (Ok(container), Ok(account_name), Ok(account_key), Ok(endpoint)) = ( + env::var("GT_AZBLOB_CONTAINER"), + env::var("GT_AZBLOB_ACCOUNT_NAME"), + env::var("GT_AZBLOB_ACCOUNT_KEY"), + env::var("GT_AZBLOB_ENDPOINT"), + ) && !container.is_empty() + { + let instance = instance.frontend(); + + assert!(matches!( + execute_sql( + &instance, + "create table demo(host string, cpu double, memory double, ts timestamp time index);", + ) + .await + .data, + OutputData::AffectedRows(0) + )); + + let output = execute_sql( + &instance, + r#"insert into demo(host, cpu, memory, ts) values + ('host1', 66.6, 1024, 1655276557000), + ('host2', 88.8, 333.3, 1655276558000) + "#, + ) + .await + .data; + assert!(matches!(output, OutputData::AffectedRows(2))); + + let root = uuid::Uuid::new_v4().to_string(); + + let copy_to_stmt = format!( + "Copy demo TO 'azblob://{}/{}/export/demo.parquet' CONNECTION (ACCOUNT_NAME='{}',ACCOUNT_KEY='{}',ENDPOINT='{}')", + container, root, account_name, account_key, endpoint + ); + + let output = execute_sql(&instance, ©_to_stmt).await.data; + assert!(matches!(output, OutputData::AffectedRows(2))); + + struct Test<'a> { + sql: &'a str, + table_name: &'a str, + } + let tests = [ + Test { + sql: &format!( + "Copy with_filename FROM 'azblob://{}/{}/export/demo.parquet'", + container, root + ), + table_name: "with_filename", + }, + Test { + sql: &format!( + "Copy with_path FROM 'azblob://{}/{}/export/'", + container, root + ), + table_name: "with_path", + }, + Test { + sql: &format!( + "Copy with_pattern FROM 'azblob://{}/{}/export/' WITH (PATTERN = 'demo.*')", + container, root + ), + table_name: "with_pattern", + }, + ]; + + for test in tests { + assert!(matches!( + execute_sql( + &instance, + &format!( + "create table {}(host string, cpu double, memory double, ts timestamp time index);", + test.table_name + ), + ) + .await + .data, + OutputData::AffectedRows(0) + )); + + let sql = format!( + "{} CONNECTION (ACCOUNT_NAME='{}',ACCOUNT_KEY='{}',ENDPOINT='{}')", + test.sql, account_name, account_key, endpoint + ); let output = execute_sql(&instance, &sql).await.data; assert!(matches!(output, OutputData::AffectedRows(2)));