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

operator/statement/
database_copy.rs
