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 <wenymedia@gmail.com>

* test: add copy to/from oss integration coverage

Signed-off-by: WenyXu <wenymedia@gmail.com>

* test: add copy to/from gcs/azure integration coverage

Signed-off-by: WenyXu <wenymedia@gmail.com>

* refactor: use existing layers

Signed-off-by: WenyXu <wenymedia@gmail.com>

* feat: hide `credential_path` in cli

Signed-off-by: WenyXu <wenymedia@gmail.com>

* chore: remove logs

Signed-off-by: WenyXu <wenymedia@gmail.com>

* chore: apply suggestions

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-03-04 04:21:21 +00:00
committed by GitHub
parent b1b76fde0b
commit b17947b5d8
10 changed files with 646 additions and 47 deletions
+3 -1
View File
@@ -85,7 +85,7 @@ impl FieldValidator for Option<SecretString> {
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<SecretString>,
#[doc = "The credential of the object store."]
credential: Option<SecretString>,
+2 -6
View File
@@ -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;
+1 -5
View File
@@ -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)],
)
}
}
+18
View File
@@ -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<host>, path)`
pub fn parse_url(url: &str) -> Result<(String, Option<String>, String)> {
@@ -74,6 +80,18 @@ pub fn build_backend(url: &str, connection: &HashMap<String, String>) -> 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 {
@@ -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<String, String>,
) -> Result<ObjectStore> {
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"));
}
}
+2 -7
View File
@@ -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<ObjectStore> {
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))
}
@@ -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<String, String>,
) -> Result<ObjectStore> {
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"));
}
}
+6 -11
View File
@@ -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)]
+7 -14
View File
@@ -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)]
+459 -3
View File
@@ -1672,8 +1672,6 @@ async fn test_execute_copy_to_s3(instance: Arc<dyn MockInstance>) {
#[apply(both_instances_cases)]
async fn test_execute_copy_from_s3(instance: Arc<dyn MockInstance>) {
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<dyn MockInstance>) {
"{} 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<dyn MockInstance>) {
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, &copy_to_stmt).await.data;
assert!(matches!(output, OutputData::AffectedRows(2)));
}
}
#[apply(both_instances_cases)]
async fn test_execute_copy_from_oss(instance: Arc<dyn MockInstance>) {
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, &copy_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<dyn MockInstance>) {
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, &copy_to_stmt).await.data;
assert!(matches!(output, OutputData::AffectedRows(2)));
}
}
#[apply(both_instances_cases)]
async fn test_execute_copy_from_gcs(instance: Arc<dyn MockInstance>) {
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, &copy_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<dyn MockInstance>) {
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, &copy_to_stmt).await.data;
assert!(matches!(output, OutputData::AffectedRows(2)));
}
}
#[apply(both_instances_cases)]
async fn test_execute_copy_from_azblob(instance: Arc<dyn MockInstance>) {
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, &copy_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)));