operator/statement/
database_copy.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
15//! Shared selection and destination rules for database COPY and prepared exports.
16
17use std::collections::HashMap;
18
19use common_datasource::object_store::FILE_SCHEMA;
20#[cfg(windows)]
21use common_datasource::object_store::{FS_SCHEMA, parse_url};
22use common_stat::get_total_cpu_cores;
23use session::context::QueryContextRef;
24use snafu::{OptionExt, ResultExt, ensure};
25use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME};
26use table::TableRef;
27use table::metadata::TableType;
28use table::requests::CopyDatabaseRequest;
29use table::table_reference::TableReference;
30use tokio::sync::Semaphore;
31use url::Url;
32
33use crate::error::{self, Result};
34use crate::statement::StatementExecutor;
35
36fn is_directory_location(location: &str) -> bool {
37    if location.ends_with('/') {
38        return true;
39    }
40
41    #[cfg(windows)]
42    {
43        location.ends_with(std::path::MAIN_SEPARATOR)
44            && matches!(
45                parse_url(location),
46                Ok((schema, _, _)) if schema.eq_ignore_ascii_case(FS_SCHEMA)
47            )
48    }
49
50    #[cfg(not(windows))]
51    false
52}
53
54/// Get parallelism from options, default to total CPU cores.
55pub(crate) fn parse_parallelism_from_option_map(options: &HashMap<String, String>) -> usize {
56    options
57        .get("parallelism")
58        .and_then(|v| v.parse::<usize>().ok())
59        .unwrap_or_else(get_total_cpu_cores)
60        .clamp(1, Semaphore::MAX_PERMITS)
61}
62
63/// Rejects layouts unsupported by the per-table export path.
64pub(crate) fn validate_database_export_layout(options: &HashMap<String, String>) -> Result<()> {
65    if let Some(layout) = options.get("metric_data_layout") {
66        return error::InvalidCopyParameterSnafu {
67            key: "metric_data_layout",
68            value: layout,
69        }
70        .fail();
71    }
72    Ok(())
73}
74
75pub(crate) fn validate_database_directory(location: &str) -> Result<()> {
76    ensure!(
77        is_directory_location(location),
78        error::InvalidCopyDatabasePathSnafu { value: location }
79    );
80    #[cfg(windows)]
81    if common_datasource::object_store::handle_windows_path(location).is_some() {
82        return Ok(());
83    }
84    let parsed_directory = match Url::parse(location) {
85        Ok(url) => {
86            url.query().is_none() && url.fragment().is_none() && is_directory_location(url.path())
87        }
88        Err(_) => true,
89    };
90    ensure!(
91        parsed_directory,
92        error::InvalidCopyDatabasePathSnafu { value: location }
93    );
94    Ok(())
95}
96
97/// The writer key and its externally reported location, resolved together.
98pub(crate) struct DatabaseExportFile {
99    pub(crate) path: String,
100    pub(crate) location: String,
101}
102
103impl DatabaseExportFile {
104    pub(crate) fn new(directory: &str, name: &str, suffix: &str) -> Result<Self> {
105        let filename = format!("{name}{suffix}");
106        #[cfg(windows)]
107        if common_datasource::object_store::handle_windows_path(directory).is_some() {
108            return Ok(Self {
109                location: format!("{directory}{filename}"),
110                path: filename,
111            });
112        }
113        match Url::parse(directory) {
114            Ok(mut url) => {
115                url.path_segments_mut()
116                    .map_err(|_| error::InvalidCopyDatabasePathSnafu { value: directory }.build())?
117                    .pop_if_empty()
118                    .push(&filename);
119                // File URLs are decoded by the filesystem backend; object-store
120                // backends use the encoded URL path as their key.
121                let path = if url.scheme().eq_ignore_ascii_case(FILE_SCHEMA) {
122                    filename
123                } else {
124                    url.path()
125                        .rsplit('/')
126                        .next()
127                        .unwrap_or_default()
128                        .to_string()
129                };
130                Ok(Self {
131                    path,
132                    location: url.into(),
133                })
134            }
135            Err(url::ParseError::RelativeUrlWithoutBase) => Ok(Self {
136                location: format!("{directory}{filename}"),
137                path: filename,
138            }),
139            Err(source) => Err(source)
140                .context(common_datasource::error::InvalidUrlSnafu { url: directory })
141                .context(error::BuildBackendSnafu),
142        }
143    }
144}
145
146/// Resolve a listed writer key back to its table name and COPY input location.
147pub(crate) fn database_import_source(directory: &str, path: &str) -> Result<(String, String)> {
148    let mut filename = path.rsplit('/').next().unwrap_or(path).to_string();
149    let mut location = format!("{directory}{path}");
150    #[cfg(windows)]
151    let literal_path = common_datasource::object_store::handle_windows_path(directory).is_some();
152    #[cfg(not(windows))]
153    let literal_path = false;
154    if !literal_path && let Ok(mut url) = Url::parse(directory) {
155        if url.scheme().eq_ignore_ascii_case(FILE_SCHEMA) {
156            url.path_segments_mut()
157                .map_err(|_| error::InvalidCopyDatabasePathSnafu { value: directory }.build())?
158                .pop_if_empty()
159                .extend(path.split('/'));
160        } else {
161            // Listed object keys already contain the export URL's escaping.
162            url.set_path(&format!("{}{path}", url.path()));
163            filename = percent_encoding::percent_decode_str(&filename)
164                .decode_utf8()
165                .ok()
166                .context(error::InvalidTableNameSnafu { table_name: path })?
167                .into_owned();
168        }
169        location = url.into();
170    }
171    let table_name = filename
172        .rsplit_once('.')
173        .map(|(stem, _)| stem)
174        .filter(|stem| !stem.is_empty())
175        .context(error::InvalidTableNameSnafu { table_name: path })?
176        .to_string();
177    Ok((table_name, location))
178}
179
180impl StatementExecutor {
181    /// Capture each selected data table once. Views, temporary and Metric physical
182    /// tables do not have data outputs. `None` selects the whole schema.
183    pub async fn capture_database_export_tables(
184        &self,
185        req: &CopyDatabaseRequest,
186        names: Option<&[String]>,
187        ctx: &QueryContextRef,
188    ) -> Result<Vec<TableRef>> {
189        let mut names = match names {
190            Some(names) => names.to_vec(),
191            None => self
192                .catalog_manager
193                .table_names(&req.catalog_name, &req.schema_name, Some(ctx))
194                .await
195                .context(error::CatalogSnafu)?,
196        };
197        names.sort();
198        let mut tables = Vec::new();
199        for name in names {
200            let table = self
201                .get_table(&TableReference::full(
202                    &req.catalog_name,
203                    &req.schema_name,
204                    &name,
205                ))
206                .await?;
207            let info = table.table_info();
208            if table.table_type() == TableType::Base
209                && (info.meta.engine != METRIC_ENGINE_NAME
210                    || info
211                        .meta
212                        .options
213                        .extra_options
214                        .contains_key(LOGICAL_TABLE_METADATA_KEY))
215            {
216                tables.push(table);
217            }
218        }
219        Ok(tables)
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226
227    #[test]
228    fn directory_url_components_cannot_capture_output_names() {
229        for location in [
230            "file:///copy/fresh?attempt=/",
231            "file:///copy/fresh#attempt/",
232            "s3://bucket/fresh?attempt=/",
233            "s3://bucket/fresh#attempt/",
234            "file:///copy/fresh",
235        ] {
236            assert!(validate_database_directory(location).is_err(), "{location}");
237        }
238        for location in ["/copy/fresh/", "file:///copy/fresh/", "s3://bucket/fresh/"] {
239            validate_database_directory(location).unwrap();
240        }
241    }
242
243    #[cfg(windows)]
244    #[test]
245    fn windows_directory_names_are_literal_paths() {
246        for location in ["C:/copy/fresh#1/", r"C:\copy\fresh#1\"] {
247            validate_database_directory(location).unwrap();
248        }
249        assert!(validate_database_directory("C:/copy/fresh#1").is_err());
250    }
251
252    #[tokio::test]
253    async fn output_locations_resolve_to_writer_keys() {
254        use common_datasource::object_store::{LocalFileAccess, build_backend_for_write_with_path};
255
256        let dir = common_test_util::temp_dir::create_temp_dir("database_export_paths");
257        let access = LocalFileAccess::sandboxed(dir.path()).unwrap();
258        let file_url = Url::from_directory_path(dir.path()).unwrap().to_string();
259        let connection = HashMap::from([
260            ("region".into(), "us-east-1".into()),
261            ("access_key_id".into(), "test-key".into()),
262            ("secret_access_key".into(), "test-secret".into()),
263        ]);
264        for directory in [
265            format!("{}/", dir.path().display()),
266            file_url,
267            "s3://export-bucket/data/".into(),
268        ] {
269            for name in ["a#b", "a:b"] {
270                if cfg!(windows) && name.contains(':') && !directory.starts_with("s3:") {
271                    continue;
272                }
273                let file = DatabaseExportFile::new(&directory, name, ".parquet").unwrap();
274                let backend =
275                    build_backend_for_write_with_path(&file.location, &connection, &access)
276                        .await
277                        .unwrap();
278                assert_eq!(backend.object_path.as_deref(), Some(file.path.as_str()));
279                let (table_name, input_location) =
280                    database_import_source(&directory, &file.path).unwrap();
281                assert_eq!(table_name, name);
282                assert_eq!(input_location, file.location);
283                let (table_name, nested_location) =
284                    database_import_source(&directory, &format!("nested/{}", file.path)).unwrap();
285                assert_eq!(table_name, name);
286                assert_eq!(
287                    nested_location,
288                    file.location
289                        .replace(&directory, &format!("{directory}nested/"))
290                );
291                if !directory.starts_with("s3:") {
292                    backend
293                        .object_store
294                        .write(&file.path, "test")
295                        .await
296                        .unwrap();
297                    assert_eq!(
298                        std::fs::read(dir.path().join(format!("{name}.parquet"))).unwrap(),
299                        b"test"
300                    );
301                } else {
302                    assert_eq!(
303                        file.path,
304                        if name == "a#b" {
305                            "a%23b.parquet"
306                        } else {
307                            "a:b.parquet"
308                        }
309                    );
310                }
311            }
312        }
313    }
314
315    #[test]
316    fn test_parse_parallelism_from_option_map() {
317        let options = HashMap::new();
318        assert_eq!(
319            parse_parallelism_from_option_map(&options),
320            get_total_cpu_cores()
321        );
322
323        let options = HashMap::from([("parallelism".to_string(), "0".to_string())]);
324        assert_eq!(parse_parallelism_from_option_map(&options), 1);
325
326        let options = HashMap::from([("parallelism".to_string(), usize::MAX.to_string())]);
327        assert_eq!(
328            parse_parallelism_from_option_map(&options),
329            Semaphore::MAX_PERMITS
330        );
331    }
332}
operator/statement/database_copy.rs

operator/statement/
database_copy.rs
