operator/statement/
copy_database.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 std::collections::HashMap;
16use std::path::Path;
17use std::str::FromStr;
18use std::sync::Arc;
19
20use client::{Output, OutputData, OutputMeta};
21use common_catalog::format_full_table_name;
22use common_datasource::file_format::Format;
23use common_datasource::lister::{Lister, Source};
24#[cfg(windows)]
25use common_datasource::object_store::{FS_SCHEMA, parse_url};
26use common_datasource::object_store::{LocalFileAccess, build_backend, build_backend_for_write};
27use common_stat::get_total_cpu_cores;
28use common_telemetry::{debug, error, info, tracing};
29use futures::future::try_join_all;
30use object_store::Entry;
31use regex::Regex;
32use session::context::QueryContextRef;
33use snafu::{OptionExt, ResultExt, ensure};
34use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME};
35use table::requests::{CopyDatabaseRequest, CopyDirection, CopyTableRequest};
36use table::table_reference::TableReference;
37use tokio::sync::Semaphore;
38
39use crate::error;
40use crate::error::{CatalogSnafu, InvalidCopyDatabasePathSnafu};
41use crate::statement::StatementExecutor;
42
43pub(crate) const COPY_DATABASE_TIME_START_KEY: &str = "start_time";
44pub(crate) const COPY_DATABASE_TIME_END_KEY: &str = "end_time";
45pub(crate) const CONTINUE_ON_ERROR_KEY: &str = "continue_on_error";
46pub(crate) const PARALLELISM_KEY: &str = "parallelism";
47
48fn is_directory_location(location: &str) -> bool {
49    if location.ends_with('/') {
50        return true;
51    }
52
53    #[cfg(windows)]
54    {
55        location.ends_with(std::path::MAIN_SEPARATOR)
56            && matches!(
57                parse_url(location),
58                Ok((schema, _, _)) if schema.eq_ignore_ascii_case(FS_SCHEMA)
59            )
60    }
61
62    #[cfg(not(windows))]
63    false
64}
65
66/// Get parallelism from options, default to total CPU cores.
67fn parse_parallelism_from_option_map(options: &HashMap<String, String>) -> usize {
68    options
69        .get(PARALLELISM_KEY)
70        .and_then(|v| v.parse::<usize>().ok())
71        .unwrap_or_else(get_total_cpu_cores)
72        .max(1)
73}
74
75impl StatementExecutor {
76    #[tracing::instrument(skip_all)]
77    pub(crate) async fn copy_database_to(
78        &self,
79        req: CopyDatabaseRequest,
80        ctx: QueryContextRef,
81    ) -> error::Result<Output> {
82        // Location must end with a separator so that every table is exported to a file.
83        ensure!(
84            is_directory_location(&req.location),
85            InvalidCopyDatabasePathSnafu {
86                value: req.location,
87            }
88        );
89        build_backend_for_write(&req.location, &req.connection, &self.local_file_access)
90            .await
91            .context(error::BuildBackendSnafu)?;
92
93        let parallelism = parse_parallelism_from_option_map(&req.with);
94        info!(
95            "Copy database {}.{} to dir: {}, time: {:?}, parallelism: {}",
96            req.catalog_name, req.schema_name, req.location, req.time_range, parallelism
97        );
98        let table_names = self
99            .catalog_manager
100            .table_names(&req.catalog_name, &req.schema_name, Some(&ctx))
101            .await
102            .context(CatalogSnafu)?;
103        let num_tables = table_names.len();
104
105        let suffix = Format::try_from(&req.with)
106            .context(error::ParseFileFormatSnafu)?
107            .suffix();
108
109        let mut tasks = Vec::with_capacity(num_tables);
110        let semaphore = Arc::new(Semaphore::new(parallelism));
111
112        for (i, table_name) in table_names.into_iter().enumerate() {
113            let table = self
114                .get_table(&TableReference {
115                    catalog: &req.catalog_name,
116                    schema: &req.schema_name,
117                    table: &table_name,
118                })
119                .await?;
120            // Only base tables, ignores views and temporary tables.
121            if table.table_type() != table::metadata::TableType::Base {
122                continue;
123            }
124            // Ignores physical tables of metric engine.
125            if table.table_info().meta.engine == METRIC_ENGINE_NAME
126                && !table
127                    .table_info()
128                    .meta
129                    .options
130                    .extra_options
131                    .contains_key(LOGICAL_TABLE_METADATA_KEY)
132            {
133                continue;
134            }
135
136            let semaphore_moved = semaphore.clone();
137            let mut table_file = req.location.clone();
138            table_file.push_str(&table_name);
139            table_file.push_str(suffix);
140            let table_no = i + 1;
141            let moved_ctx = ctx.clone();
142            let full_table_name =
143                format_full_table_name(&req.catalog_name, &req.schema_name, &table_name);
144            let copy_table_req = CopyTableRequest {
145                catalog_name: req.catalog_name.clone(),
146                schema_name: req.schema_name.clone(),
147                table_name,
148                location: table_file.clone(),
149                with: req.with.clone(),
150                connection: req.connection.clone(),
151                pattern: None,
152                direction: CopyDirection::Export,
153                timestamp_range: req.time_range,
154                limit: None,
155            };
156
157            tasks.push(async move {
158                let _permit = semaphore_moved.acquire().await.unwrap();
159                info!(
160                    "Copy table({}/{}): {} to {}",
161                    table_no, num_tables, full_table_name, table_file
162                );
163                self.copy_table_to(copy_table_req, moved_ctx).await
164            });
165        }
166
167        let results = try_join_all(tasks).await?;
168        let exported_rows = results.into_iter().sum();
169        Ok(Output::new_with_affected_rows(exported_rows))
170    }
171
172    /// Imports data to database from a given location and returns total rows imported.
173    #[tracing::instrument(skip_all)]
174    pub(crate) async fn copy_database_from(
175        &self,
176        req: CopyDatabaseRequest,
177        ctx: QueryContextRef,
178    ) -> error::Result<Output> {
179        // Location must end with a directory separator.
180        ensure!(
181            is_directory_location(&req.location),
182            InvalidCopyDatabasePathSnafu {
183                value: req.location,
184            }
185        );
186
187        let parallelism = parse_parallelism_from_option_map(&req.with);
188        info!(
189            "Copy database {}.{} from dir: {}, time: {:?}, parallelism: {}",
190            req.catalog_name, req.schema_name, req.location, req.time_range, parallelism
191        );
192        let suffix = Format::try_from(&req.with)
193            .context(error::ParseFileFormatSnafu)?
194            .suffix();
195
196        let entries = list_files_to_copy(&req, suffix, &self.local_file_access).await?;
197
198        let continue_on_error = req
199            .with
200            .get(CONTINUE_ON_ERROR_KEY)
201            .and_then(|v| bool::from_str(v).ok())
202            .unwrap_or(false);
203
204        let mut tasks = Vec::with_capacity(entries.len());
205        let semaphore = Arc::new(Semaphore::new(parallelism));
206
207        for e in entries {
208            let table_name = match parse_file_name_to_copy(&e) {
209                Ok(table_name) => table_name,
210                Err(err) => {
211                    if continue_on_error {
212                        error!(err; "Failed to import table from file: {:?}", e);
213                        continue;
214                    } else {
215                        return Err(err);
216                    }
217                }
218            };
219
220            let req = CopyTableRequest {
221                catalog_name: req.catalog_name.clone(),
222                schema_name: req.schema_name.clone(),
223                table_name: table_name.clone(),
224                location: format!("{}{}", req.location, e.path()),
225                with: req.with.clone(),
226                connection: req.connection.clone(),
227                pattern: None,
228                direction: CopyDirection::Import,
229                timestamp_range: None,
230                limit: None,
231            };
232            let moved_ctx = ctx.clone();
233            let moved_table_name = table_name.clone();
234            let moved_semaphore = semaphore.clone();
235            tasks.push(async move {
236                let _permit = moved_semaphore.acquire().await.unwrap();
237                debug!("Copy table, arg: {:?}", req);
238                match self.copy_table_from(req, moved_ctx).await {
239                    Ok(o) => {
240                        let (rows, cost) = o.extract_rows_and_cost();
241                        Ok((rows, cost))
242                    }
243                    Err(err) => {
244                        if continue_on_error {
245                            error!(err; "Failed to import file to table: {}", moved_table_name);
246                            Ok((0, 0))
247                        } else {
248                            Err(err)
249                        }
250                    }
251                }
252            });
253        }
254
255        let results = try_join_all(tasks).await?;
256        let (rows_inserted, insert_cost) = results
257            .into_iter()
258            .fold((0, 0), |(acc_rows, acc_cost), (rows, cost)| {
259                (acc_rows + rows, acc_cost + cost)
260            });
261
262        Ok(Output::new(
263            OutputData::AffectedRows(rows_inserted),
264            OutputMeta::new_with_cost(insert_cost),
265        ))
266    }
267}
268
269/// Parses table names from files' names.
270fn parse_file_name_to_copy(e: &Entry) -> error::Result<String> {
271    Path::new(e.name())
272        .file_stem()
273        .and_then(|os_str| os_str.to_str())
274        .map(|s| s.to_string())
275        .context(error::InvalidTableNameSnafu {
276            table_name: e.name().to_string(),
277        })
278}
279
280/// Lists all files with expected suffix that can be imported to database.
281async fn list_files_to_copy(
282    req: &CopyDatabaseRequest,
283    suffix: &str,
284    local_file_access: &LocalFileAccess,
285) -> error::Result<Vec<Entry>> {
286    let object_store = build_backend(&req.location, &req.connection, local_file_access)
287        .await
288        .context(error::BuildBackendSnafu)?;
289
290    let pattern = Regex::try_from(format!(".*{}", suffix)).context(error::BuildRegexSnafu)?;
291    let lister = Lister::new(
292        object_store.clone(),
293        Source::Dir,
294        "/".to_string(),
295        Some(pattern),
296    );
297    lister.list().await.context(error::ListObjectsSnafu)
298}
299
300#[cfg(test)]
301mod tests {
302    use std::collections::{HashMap, HashSet};
303
304    use common_datasource::object_store::LocalFileAccess;
305    use common_stat::get_total_cpu_cores;
306    use object_store::ObjectStore;
307    use object_store::services::Fs;
308    use object_store::util::normalize_dir;
309    #[cfg(not(windows))]
310    use path_slash::PathExt;
311    use table::requests::CopyDatabaseRequest;
312
313    use crate::statement::copy_database::{
314        list_files_to_copy, parse_file_name_to_copy, parse_parallelism_from_option_map,
315    };
316
317    #[tokio::test]
318    async fn test_list_files_and_parse_table_name() {
319        let dir = common_test_util::temp_dir::create_temp_dir("test_list_files_to_copy");
320        let store_dir = normalize_dir(dir.path().to_str().unwrap());
321        let builder = Fs::default().root(&store_dir);
322        let object_store = ObjectStore::new(builder).unwrap();
323        object_store.write("a.parquet", "").await.unwrap();
324        object_store.write("b.parquet", "").await.unwrap();
325        object_store.write("c.csv", "").await.unwrap();
326        object_store.write("d", "").await.unwrap();
327        object_store.write("e.f.parquet", "").await.unwrap();
328
329        #[cfg(not(windows))]
330        let location = normalize_dir(&dir.path().to_slash().unwrap());
331        #[cfg(windows)]
332        let location = format!("{}\\", dir.path().display());
333        let request = CopyDatabaseRequest {
334            catalog_name: "catalog_0".to_string(),
335            schema_name: "schema_0".to_string(),
336            location,
337            with: [("FORMAT".to_string(), "parquet".to_string())]
338                .into_iter()
339                .collect(),
340            connection: Default::default(),
341            time_range: None,
342        };
343        let local_file_access = LocalFileAccess::sandboxed(dir.path()).unwrap();
344        let listed = list_files_to_copy(&request, ".parquet", &local_file_access)
345            .await
346            .unwrap()
347            .into_iter()
348            .map(|e| parse_file_name_to_copy(&e).unwrap())
349            .collect::<HashSet<_>>();
350
351        assert_eq!(
352            ["a".to_string(), "b".to_string(), "e.f".to_string()]
353                .into_iter()
354                .collect::<HashSet<_>>(),
355            listed
356        );
357    }
358
359    #[test]
360    fn test_parse_parallelism_from_option_map() {
361        let options = HashMap::new();
362        assert_eq!(
363            parse_parallelism_from_option_map(&options),
364            get_total_cpu_cores()
365        );
366
367        let options = HashMap::from([("parallelism".to_string(), "0".to_string())]);
368        assert_eq!(parse_parallelism_from_option_map(&options), 1);
369    }
370}
operator/statement/copy_database.rs

operator/statement/
copy_database.rs
