cli/data/export_v2/
data.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use common_base::secrets::{ExposeSecret, SecretString};
16use common_telemetry::info;
17use object_store::util::{join_path, normalize_path};
18use snafu::ResultExt;
19use url::Url;
20
21use crate::common::ObjectStoreConfig;
22use crate::data::export_v2::error::{DatabaseSnafu, InvalidUriSnafu, Result, UrlParseSnafu};
23use crate::data::export_v2::manifest::{DataFormat, TimeRange};
24use crate::data::path::data_dir_for_schema_chunk;
25use crate::data::snapshot_storage::StorageScheme;
26use crate::data::sql::{escape_sql_identifier, escape_sql_literal};
27use crate::database::DatabaseClient;
28
29pub(super) struct CopyOptions {
30    pub(super) format: DataFormat,
31    pub(super) time_range: TimeRange,
32    pub(super) parallelism: usize,
33    pub(super) packed: bool,
34}
35
36pub(super) struct CopyTarget {
37    pub(super) location: String,
38    pub(super) connection: String,
39    secrets: Vec<Option<String>>,
40}
41
42pub(crate) struct CopySource {
43    pub(crate) location: String,
44    pub(crate) connection: String,
45    secrets: Vec<Option<String>>,
46}
47
48impl CopyTarget {
49    fn mask_sql(&self, sql: &str) -> String {
50        mask_secrets(sql, &self.secrets)
51    }
52}
53
54impl CopySource {
55    fn mask_sql(&self, sql: &str) -> String {
56        mask_secrets(sql, &self.secrets)
57    }
58}
59
60pub(super) fn build_copy_target(
61    snapshot_uri: &str,
62    storage: &ObjectStoreConfig,
63    schema: &str,
64    chunk_id: u32,
65) -> Result<CopyTarget> {
66    let location = build_copy_location(snapshot_uri, storage, schema, chunk_id)?;
67    Ok(CopyTarget {
68        location: location.location,
69        connection: location.connection,
70        secrets: location.secrets,
71    })
72}
73
74pub(crate) fn build_copy_source(
75    snapshot_uri: &str,
76    storage: &ObjectStoreConfig,
77    schema: &str,
78    chunk_id: u32,
79) -> Result<CopySource> {
80    let location = build_copy_location(snapshot_uri, storage, schema, chunk_id)?;
81    Ok(CopySource {
82        location: location.location,
83        connection: location.connection,
84        secrets: location.secrets,
85    })
86}
87
88struct CopyLocation {
89    location: String,
90    connection: String,
91    secrets: Vec<Option<String>>,
92}
93
94fn build_copy_location(
95    snapshot_uri: &str,
96    storage: &ObjectStoreConfig,
97    schema: &str,
98    chunk_id: u32,
99) -> Result<CopyLocation> {
100    let url = Url::parse(snapshot_uri).context(UrlParseSnafu)?;
101    let scheme = StorageScheme::from_uri(snapshot_uri)?;
102    let suffix = data_dir_for_schema_chunk(schema, chunk_id);
103
104    match scheme {
105        StorageScheme::File => {
106            let root = url.to_file_path().map_err(|_| {
107                InvalidUriSnafu {
108                    uri: snapshot_uri,
109                    reason: "file:// URI must use an absolute path like file:///tmp/backup",
110                }
111                .build()
112            })?;
113            let location = normalize_path(&format!("{}/{}", root.to_string_lossy(), suffix));
114            Ok(CopyLocation {
115                location,
116                connection: String::new(),
117                secrets: Vec::new(),
118            })
119        }
120        StorageScheme::S3 => {
121            let (bucket, root) = extract_bucket_root(&url, snapshot_uri)?;
122            let location = format!("s3://{}/{}", bucket, join_root(&root, &suffix));
123            let (connection, secrets) = build_s3_connection(storage);
124            Ok(CopyLocation {
125                location,
126                connection,
127                secrets,
128            })
129        }
130        StorageScheme::Oss => {
131            let (bucket, root) = extract_bucket_root(&url, snapshot_uri)?;
132            let location = format!("oss://{}/{}", bucket, join_root(&root, &suffix));
133            let (connection, secrets) = build_oss_connection(storage);
134            Ok(CopyLocation {
135                location,
136                connection,
137                secrets,
138            })
139        }
140        StorageScheme::Gcs => {
141            let (bucket, root) = extract_bucket_root(&url, snapshot_uri)?;
142            let location = format!("gcs://{}/{}", bucket, join_root(&root, &suffix));
143            let (connection, secrets) = build_gcs_connection(storage, snapshot_uri)?;
144            Ok(CopyLocation {
145                location,
146                connection,
147                secrets,
148            })
149        }
150        StorageScheme::Azblob => {
151            let (bucket, root) = extract_bucket_root(&url, snapshot_uri)?;
152            let location = format!("azblob://{}/{}", bucket, join_root(&root, &suffix));
153            let (connection, secrets) = build_azblob_connection(storage);
154            Ok(CopyLocation {
155                location,
156                connection,
157                secrets,
158            })
159        }
160    }
161}
162
163pub(super) async fn execute_copy_database(
164    database_client: &DatabaseClient,
165    catalog: &str,
166    schema: &str,
167    target: &CopyTarget,
168    options: &CopyOptions,
169) -> Result<()> {
170    let with_options = build_with_options(options);
171    let sql = format!(
172        r#"COPY DATABASE "{}"."{}" TO '{}' WITH ({}){};"#,
173        escape_sql_identifier(catalog),
174        escape_sql_identifier(schema),
175        escape_sql_literal(&target.location),
176        with_options,
177        target.connection
178    );
179    let safe_sql = target.mask_sql(&sql);
180    info!("Executing sql: {}", safe_sql);
181    database_client
182        .sql_in_public(&sql)
183        .await
184        .context(DatabaseSnafu)?;
185    Ok(())
186}
187
188pub(crate) async fn execute_copy_database_from(
189    database_client: &DatabaseClient,
190    catalog: &str,
191    schema: &str,
192    source: &CopySource,
193    format: DataFormat,
194    packed: bool,
195) -> Result<()> {
196    let sql = format!(
197        r#"COPY DATABASE "{}"."{}" FROM '{}' WITH (FORMAT='{}'{}){};"#,
198        escape_sql_identifier(catalog),
199        escape_sql_identifier(schema),
200        escape_sql_literal(&source.location),
201        format,
202        if packed {
203            ", metric_data_layout='packed'"
204        } else {
205            ""
206        },
207        source.connection
208    );
209    let safe_sql = source.mask_sql(&sql);
210    info!("Executing sql: {}", safe_sql);
211    database_client
212        .sql_in_public(&sql)
213        .await
214        .context(DatabaseSnafu)?;
215    Ok(())
216}
217
218fn build_with_options(options: &CopyOptions) -> String {
219    let mut parts = vec![format!("FORMAT='{}'", options.format)];
220    if options.packed {
221        parts.push("experimental_metric_export='true'".to_string());
222        parts.push("metric_data_layout='packed'".to_string());
223    }
224    if let Some(start) = options.time_range.start {
225        parts.push(format!(
226            "START_TIME='{}'",
227            escape_sql_literal(&start.to_rfc3339())
228        ));
229    }
230    if let Some(end) = options.time_range.end {
231        parts.push(format!(
232            "END_TIME='{}'",
233            escape_sql_literal(&end.to_rfc3339())
234        ));
235    }
236    parts.push(format!("PARALLELISM={}", options.parallelism));
237    parts.join(", ")
238}
239
240fn extract_bucket_root(url: &Url, snapshot_uri: &str) -> Result<(String, String)> {
241    let bucket = url.host_str().unwrap_or("").to_string();
242    if bucket.is_empty() {
243        return InvalidUriSnafu {
244            uri: snapshot_uri,
245            reason: "URI must include bucket/container in host",
246        }
247        .fail();
248    }
249    let root = url
250        .path()
251        .trim_start_matches('/')
252        .trim_end_matches('/')
253        .to_string();
254    Ok((bucket, root))
255}
256
257fn join_root(root: &str, suffix: &str) -> String {
258    join_path(root, suffix).trim_start_matches('/').to_string()
259}
260
261fn build_s3_connection(storage: &ObjectStoreConfig) -> (String, Vec<Option<String>>) {
262    let access_key_id = expose_optional_secret(&storage.s3.s3_access_key_id);
263    let secret_access_key = expose_optional_secret(&storage.s3.s3_secret_access_key);
264
265    let mut options = Vec::new();
266    if let Some(access_key_id) = &access_key_id {
267        options.push(format!(
268            "ACCESS_KEY_ID='{}'",
269            escape_sql_literal(access_key_id)
270        ));
271    }
272    if let Some(secret_access_key) = &secret_access_key {
273        options.push(format!(
274            "SECRET_ACCESS_KEY='{}'",
275            escape_sql_literal(secret_access_key)
276        ));
277    }
278    if let Some(region) = &storage.s3.s3_region {
279        options.push(format!("REGION='{}'", escape_sql_literal(region)));
280    }
281    if let Some(endpoint) = &storage.s3.s3_endpoint {
282        options.push(format!("ENDPOINT='{}'", escape_sql_literal(endpoint)));
283    }
284
285    let secrets = vec![access_key_id, secret_access_key];
286    let connection = if options.is_empty() {
287        String::new()
288    } else {
289        format!(" CONNECTION ({})", options.join(", "))
290    };
291    (connection, secrets)
292}
293
294fn build_oss_connection(storage: &ObjectStoreConfig) -> (String, Vec<Option<String>>) {
295    let access_key_id = expose_optional_secret(&storage.oss.oss_access_key_id);
296    let access_key_secret = expose_optional_secret(&storage.oss.oss_access_key_secret);
297
298    let mut options = Vec::new();
299    if let Some(access_key_id) = &access_key_id {
300        options.push(format!(
301            "ACCESS_KEY_ID='{}'",
302            escape_sql_literal(access_key_id)
303        ));
304    }
305    if let Some(access_key_secret) = &access_key_secret {
306        options.push(format!(
307            "ACCESS_KEY_SECRET='{}'",
308            escape_sql_literal(access_key_secret)
309        ));
310    }
311    if !storage.oss.oss_endpoint.is_empty() {
312        options.push(format!(
313            "ENDPOINT='{}'",
314            escape_sql_literal(&storage.oss.oss_endpoint)
315        ));
316    }
317
318    let secrets = vec![access_key_id, access_key_secret];
319    let connection = if options.is_empty() {
320        String::new()
321    } else {
322        format!(" CONNECTION ({})", options.join(", "))
323    };
324    (connection, secrets)
325}
326
327fn build_gcs_connection(
328    storage: &ObjectStoreConfig,
329    snapshot_uri: &str,
330) -> Result<(String, Vec<Option<String>>)> {
331    let credential_path = expose_optional_secret(&storage.gcs.gcs_credential_path);
332    let credential = expose_optional_secret(&storage.gcs.gcs_credential);
333
334    if credential.is_none() && credential_path.is_some() {
335        return InvalidUriSnafu {
336            uri: snapshot_uri,
337            reason: "gcs_credential_path is not supported for server-side COPY; provide gcs_credential or rely on server-side ADC",
338        }
339        .fail();
340    }
341
342    let mut options = Vec::new();
343    if let Some(credential) = &credential {
344        options.push(format!("CREDENTIAL='{}'", escape_sql_literal(credential)));
345    }
346    if !storage.gcs.gcs_scope.is_empty() {
347        options.push(format!(
348            "SCOPE='{}'",
349            escape_sql_literal(&storage.gcs.gcs_scope)
350        ));
351    }
352    if !storage.gcs.gcs_endpoint.is_empty() {
353        options.push(format!(
354            "ENDPOINT='{}'",
355            escape_sql_literal(&storage.gcs.gcs_endpoint)
356        ));
357    }
358
359    let connection = if options.is_empty() {
360        String::new()
361    } else {
362        format!(" CONNECTION ({})", options.join(", "))
363    };
364    let secrets = vec![credential_path, credential];
365    Ok((connection, secrets))
366}
367
368fn build_azblob_connection(storage: &ObjectStoreConfig) -> (String, Vec<Option<String>>) {
369    let account_name = expose_optional_secret(&storage.azblob.azblob_account_name);
370    let account_key = expose_optional_secret(&storage.azblob.azblob_account_key);
371    let sas_token = storage.azblob.azblob_sas_token.clone();
372
373    let mut options = Vec::new();
374    if let Some(account_name) = &account_name {
375        options.push(format!(
376            "ACCOUNT_NAME='{}'",
377            escape_sql_literal(account_name)
378        ));
379    }
380    if let Some(account_key) = &account_key {
381        options.push(format!("ACCOUNT_KEY='{}'", escape_sql_literal(account_key)));
382    }
383    if let Some(sas_token) = &sas_token {
384        options.push(format!("SAS_TOKEN='{}'", escape_sql_literal(sas_token)));
385    }
386    if !storage.azblob.azblob_endpoint.is_empty() {
387        options.push(format!(
388            "ENDPOINT='{}'",
389            escape_sql_literal(&storage.azblob.azblob_endpoint)
390        ));
391    }
392
393    let secrets = vec![account_name, account_key, sas_token];
394    let connection = if options.is_empty() {
395        String::new()
396    } else {
397        format!(" CONNECTION ({})", options.join(", "))
398    };
399    (connection, secrets)
400}
401
402fn expose_optional_secret(secret: &Option<SecretString>) -> Option<String> {
403    secret.as_ref().map(|s| s.expose_secret().to_owned())
404}
405
406fn mask_secrets(sql: &str, secrets: &[Option<String>]) -> String {
407    let mut masked = sql.to_string();
408    for secret in secrets {
409        if let Some(secret) = secret
410            && !secret.is_empty()
411        {
412            let escaped = escape_sql_literal(secret);
413            if escaped != *secret {
414                masked = masked.replace(&escaped, "[REDACTED]");
415            }
416            masked = masked.replace(secret, "[REDACTED]");
417        }
418    }
419    masked
420}
421
422#[cfg(test)]
423mod tests {
424    use common_base::secrets::SecretString;
425    use common_test_util::temp_dir::create_temp_dir;
426
427    use super::*;
428    use crate::common::{PrefixedAzblobConnection, PrefixedGcsConnection, PrefixedOssConnection};
429
430    #[test]
431    fn test_packed_copy_options() {
432        let mut options = CopyOptions {
433            format: DataFormat::Parquet,
434            time_range: TimeRange::unbounded(),
435            parallelism: 2,
436            packed: false,
437        };
438        assert_eq!(
439            build_with_options(&options),
440            "FORMAT='parquet', PARALLELISM=2"
441        );
442        options.packed = true;
443        assert_eq!(
444            build_with_options(&options),
445            "FORMAT='parquet', experimental_metric_export='true', metric_data_layout='packed', PARALLELISM=2"
446        );
447    }
448
449    #[test]
450    fn test_build_oss_connection_includes_endpoint() {
451        let storage = ObjectStoreConfig {
452            oss: PrefixedOssConnection {
453                oss_endpoint: "https://oss.example.com".to_string(),
454                oss_access_key_id: Some(SecretString::from("key_id".to_string())),
455                oss_access_key_secret: Some(SecretString::from("key_secret".to_string())),
456                ..Default::default()
457            },
458            ..Default::default()
459        };
460
461        let (connection, _) = build_oss_connection(&storage);
462        assert!(connection.contains("ENDPOINT='https://oss.example.com'"));
463    }
464
465    #[test]
466    fn test_build_gcs_connection_uses_scope_and_inline_credential() {
467        let storage = ObjectStoreConfig {
468            gcs: PrefixedGcsConnection {
469                gcs_scope: "scope-a".to_string(),
470                gcs_endpoint: "https://storage.googleapis.com".to_string(),
471                gcs_credential: Some(SecretString::from("credential-json".to_string())),
472                ..Default::default()
473            },
474            ..Default::default()
475        };
476
477        let (connection, _) = build_gcs_connection(&storage, "gcs://bucket/root").unwrap();
478        assert!(connection.contains("CREDENTIAL='credential-json'"));
479        assert!(connection.contains("SCOPE='scope-a'"));
480        assert!(connection.contains("ENDPOINT='https://storage.googleapis.com'"));
481        assert!(!connection.contains("CREDENTIAL_PATH"));
482    }
483
484    #[test]
485    fn test_build_gcs_connection_rejects_credential_path_only() {
486        let storage = ObjectStoreConfig {
487            gcs: PrefixedGcsConnection {
488                gcs_scope: "scope-a".to_string(),
489                gcs_credential_path: Some(SecretString::from("/tmp/creds.json".to_string())),
490                ..Default::default()
491            },
492            ..Default::default()
493        };
494
495        let error = build_gcs_connection(&storage, "gcs://bucket/root")
496            .expect_err("credential_path-only should be rejected")
497            .to_string();
498        assert!(error.contains("gcs_credential_path is not supported"));
499    }
500
501    #[test]
502    fn test_build_azblob_connection_includes_endpoint() {
503        let storage = ObjectStoreConfig {
504            azblob: PrefixedAzblobConnection {
505                azblob_account_name: Some(SecretString::from("account".to_string())),
506                azblob_account_key: Some(SecretString::from("key".to_string())),
507                azblob_endpoint: "https://blob.example.com".to_string(),
508                ..Default::default()
509            },
510            ..Default::default()
511        };
512
513        let (connection, _) = build_azblob_connection(&storage);
514        assert!(connection.contains("ENDPOINT='https://blob.example.com'"));
515    }
516
517    #[test]
518    fn test_build_azblob_connection_redacts_sas_token() {
519        let storage = ObjectStoreConfig {
520            azblob: PrefixedAzblobConnection {
521                azblob_account_name: Some(SecretString::from("account".to_string())),
522                azblob_account_key: Some(SecretString::from("key".to_string())),
523                azblob_sas_token: Some("sig=secret-token".to_string()),
524                ..Default::default()
525            },
526            ..Default::default()
527        };
528
529        let (connection, secrets) = build_azblob_connection(&storage);
530        let masked = mask_secrets(&connection, &secrets);
531
532        assert!(connection.contains("SAS_TOKEN='sig=secret-token'"));
533        assert!(masked.contains("SAS_TOKEN='[REDACTED]'"));
534        assert!(!masked.contains("sig=secret-token"));
535    }
536
537    #[test]
538    fn test_mask_secrets_redacts_sql_escaped_literals() {
539        let sql =
540            "COPY DATABASE \"greptime\".\"public\" TO 's3://bucket' CONNECTION (SECRET='ab''cd');";
541        let masked = mask_secrets(sql, &[Some("ab'cd".to_string())]);
542
543        assert!(!masked.contains("ab'cd"));
544        assert!(!masked.contains("ab''cd"));
545        assert!(masked.contains("SECRET='[REDACTED]'"));
546    }
547
548    #[test]
549    fn test_build_copy_target_decodes_file_uri_path() {
550        let storage = ObjectStoreConfig::default();
551        let snapshot_root = create_temp_dir("my backup");
552        let snapshot_uri = Url::from_file_path(snapshot_root.path())
553            .expect("absolute platform path should convert to file:// URI")
554            .to_string();
555        let expected = normalize_path(&format!(
556            "{}/{}",
557            snapshot_root.path().to_string_lossy(),
558            data_dir_for_schema_chunk("public", 7)
559        ));
560        let target = build_copy_target(&snapshot_uri, &storage, "public", 7)
561            .expect("file:// copy target should be built");
562
563        assert!(snapshot_uri.contains("%20"));
564        assert!(!target.location.contains("%20"));
565        assert!(target.location.contains("my backup"));
566        assert_eq!(target.location, expected);
567    }
568}
cli/data/export_v2/data.rs

cli/data/export_v2/
data.rs
