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

cli/data/export_v2/
data.rs
