operator/statement/
export_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
15//! Request-scoped database export preparation. The frontend authorizes the entire
16//! captured selection before preparation; these trusted methods do not check ACLs.
17
18use std::collections::{BTreeMap, HashSet};
19use std::future::Future;
20
21use common_datasource::file_format::Format;
22use common_datasource::object_store::{FILE_SCHEMA, FS_SCHEMA, build_backend_for_write, parse_url};
23use common_datasource::packed_writer::{PackedTableWriter, PackedWriter};
24use common_meta::key::table_route::TableRouteValue;
25use futures::StreamExt;
26use futures::stream::FuturesUnordered;
27use session::context::QueryContextRef;
28use snafu::{OptionExt, ResultExt, ensure};
29use store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME};
30use table::TableRef;
31use table::metadata::TableType;
32use table::requests::{CopyDatabaseRequest, CopyDirection, CopyTableRequest};
33use tokio_util::sync::CancellationToken;
34
35use crate::error::{self, InvalidDatabaseExportSnafu, Result};
36use crate::statement::StatementExecutor;
37use crate::statement::database_copy::{
38    DatabaseExportFile, parse_parallelism_from_option_map, validate_database_directory,
39    validate_database_export_layout,
40};
41use crate::statement::export_logical_tables::writers::{ExportWriteBudget, retain_error};
42use crate::statement::export_logical_tables::{LogicalTableExport, LogicalTableExportLimits};
43
44/// A validated request-scoped selection, not a metadata snapshot or an ACL token.
45pub struct PreparedDatabaseExport {
46    request: CopyDatabaseRequest,
47    jobs: Vec<DatabaseExportJob>,
48}
49
50impl PreparedDatabaseExport {
51    /// Inspect physical grouping through the integration testing adapter.
52    #[cfg(feature = "testing")]
53    pub fn job_count_for_test(&self) -> usize {
54        self.jobs.len()
55    }
56}
57
58enum DatabaseExportJob {
59    Ordinary {
60        table: TableRef,
61        output: DatabaseExportFile,
62    },
63    Metric(LogicalTableExport),
64}
65
66impl StatementExecutor {
67    /// Validate all jobs and destinations before executing any query or writer.
68    /// Callers must authorize every captured table before calling this method.
69    pub async fn prepare_database_export(
70        &self,
71        req: CopyDatabaseRequest,
72        tables: Vec<TableRef>,
73    ) -> Result<PreparedDatabaseExport> {
74        if req
75            .with
76            .get("metric_data_layout")
77            .is_none_or(|v| v != "packed")
78        {
79            validate_database_export_layout(&req.with)?;
80        }
81        validate_database_directory(&req.location)?;
82        let format = Format::try_from(&req.with).context(error::ParseFileFormatSnafu)?;
83        ensure!(
84            matches!(format, Format::Parquet(_)),
85            error::UnsupportedFormatSnafu { format }
86        );
87        let (scheme, _, _) = parse_url(&req.location).context(error::BuildBackendSnafu)?;
88        let local =
89            scheme.eq_ignore_ascii_case(FS_SCHEMA) || scheme.eq_ignore_ascii_case(FILE_SCHEMA);
90        let mut filenames = HashSet::new();
91        let mut logical = Vec::new();
92        let mut jobs = Vec::new();
93        for table in tables {
94            let info = table.table_info();
95            let name = &info.name;
96            let output = DatabaseExportFile::new(&req.location, name, ".parquet")?;
97            ensure!(
98                filenames.insert(if local {
99                    output.path.to_ascii_lowercase()
100                } else {
101                    output.path.clone()
102                }),
103                InvalidDatabaseExportSnafu {
104                    reason: format!("duplicate output name: {name}")
105                }
106            );
107            ensure!(
108                info.catalog_name == req.catalog_name
109                    && info.schema_name == req.schema_name
110                    && table.table_type() == TableType::Base,
111                InvalidDatabaseExportSnafu {
112                    reason: "expected base tables in the selected schema"
113                }
114            );
115            if info.meta.engine == METRIC_ENGINE_NAME {
116                ensure!(
117                    info.meta
118                        .options
119                        .extra_options
120                        .contains_key(LOGICAL_TABLE_METADATA_KEY),
121                    InvalidDatabaseExportSnafu {
122                        reason: "expected a Metric logical table"
123                    }
124                );
125                logical.push(table);
126            } else {
127                jobs.push(DatabaseExportJob::Ordinary { table, output });
128            }
129        }
130        let ids = logical
131            .iter()
132            .map(|t| t.table_info().table_id())
133            .collect::<Vec<_>>();
134        let routes = self
135            .table_metadata_manager
136            .table_route_manager()
137            .table_route_storage()
138            .batch_get(&ids)
139            .await
140            .context(error::TableMetadataManagerSnafu)?;
141        let mut groups = BTreeMap::<_, Vec<TableRef>>::new();
142        for (table, route) in logical.into_iter().zip(routes) {
143            let Some(TableRouteValue::Logical(route)) = route else {
144                return InvalidDatabaseExportSnafu {
145                    reason: format!(
146                        "missing or non-logical route for {}",
147                        table.table_info().table_id()
148                    ),
149                }
150                .fail();
151            };
152            groups
153                .entry(route.physical_table_id())
154                .or_default()
155                .push(table);
156        }
157        let physical_ids = groups.keys().copied().collect::<Vec<_>>();
158        let mut physical = self
159            .catalog_manager
160            .tables_by_ids(&req.catalog_name, &req.schema_name, &physical_ids)
161            .await
162            .context(error::CatalogSnafu)?
163            .into_iter()
164            .map(|t| (t.table_info().table_id(), t))
165            .collect::<BTreeMap<_, _>>();
166        for (id, tables) in groups {
167            let table = physical
168                .remove(&id)
169                .with_context(|| InvalidDatabaseExportSnafu {
170                    reason: format!("missing physical table {id} in the selected schema"),
171                })?;
172            jobs.push(DatabaseExportJob::Metric(
173                LogicalTableExport::try_new_in_directory(table, &tables, &req.location)?,
174            ));
175        }
176        build_backend_for_write(&req.location, &req.connection, &self.local_file_access)
177            .await
178            .context(error::BuildBackendSnafu)?;
179        Ok(PreparedDatabaseExport { request: req, jobs })
180    }
181}
182
183/// Returned only after every started job has closed or drained its owned I/O.
184#[derive(Debug)]
185pub struct DatabaseExportSummary {
186    pub rows: usize,
187    pub output_files: Vec<String>,
188}
189
190impl StatementExecutor {
191    /// Execute with one job budget. The caller must keep this future alive until
192    /// it returns, including after cooperative cancellation. Closed files remain
193    /// owned by the caller's attempt; this method does not publish completion.
194    pub async fn export_database(
195        &self,
196        plan: PreparedDatabaseExport,
197        cancellation: &CancellationToken,
198        ctx: QueryContextRef,
199    ) -> Result<DatabaseExportSummary> {
200        let mut output_files = Vec::new();
201        for job in &plan.jobs {
202            match job {
203                DatabaseExportJob::Ordinary { output, .. } => {
204                    output_files.push(output.location.clone())
205                }
206                DatabaseExportJob::Metric(unit) => {
207                    output_files.extend(unit.output_files().map(|file| file.location.clone()))
208                }
209            }
210        }
211        output_files.sort();
212        let req = &plan.request;
213        let parallelism = parse_parallelism_from_option_map(&req.with);
214        let budget = ExportWriteBudget::new(parallelism);
215        let packed = if req
216            .with
217            .get("metric_data_layout")
218            .is_some_and(|v| v == "packed")
219        {
220            let store =
221                build_backend_for_write(&req.location, &req.connection, &self.local_file_access)
222                    .await
223                    .context(error::BuildBackendSnafu)?;
224            Some(
225                PackedWriter::new(store).context(error::WriteStreamToFileSnafu {
226                    path: &req.location,
227                })?,
228            )
229        } else {
230            None
231        };
232        let rows = run_database_export_jobs(plan.jobs, parallelism, cancellation, |job, token| {
233            let ctx = ctx.clone();
234            let budget = budget.clone();
235            let packed = packed.clone();
236            async move {
237                match job {
238                    DatabaseExportJob::Metric(unit) => self
239                        .export_logical_tables_managed(
240                            &unit,
241                            &req.location,
242                            &req.connection,
243                            req.time_range.as_ref(),
244                            LogicalTableExportLimits::default(),
245                            &token,
246                            ctx,
247                            budget,
248                            packed,
249                        )
250                        .await
251                        .map(|summary| summary.rows),
252                    DatabaseExportJob::Ordinary { table, output } => {
253                        let info = table.table_info();
254                        let copy = CopyTableRequest {
255                            catalog_name: info.catalog_name.clone(),
256                            schema_name: info.schema_name.clone(),
257                            table_name: info.name.clone(),
258                            location: output.location,
259                            with: req.with.clone(),
260                            connection: req.connection.clone(),
261                            pattern: None,
262                            direction: CopyDirection::Export,
263                            timestamp_range: req.time_range,
264                            limit: None,
265                        };
266                        let _permit = budget.writer(&token).await?;
267                        let destination = packed.map(|shared| {
268                            PackedTableWriter::new(shared, info.name.clone(), info.table_id(), true)
269                        });
270                        self.copy_captured_table_to_managed(
271                            table,
272                            copy,
273                            ctx,
274                            Some((&budget, &token)),
275                            destination,
276                        )
277                        .await
278                    }
279                }
280            }
281        })
282        .await;
283        if let Some(packed) = packed {
284            let mut packed = packed.lock().await;
285            let result = match rows {
286                Ok(rows) => packed
287                    .finish(cancellation)
288                    .await
289                    .context(error::WriteStreamToFileSnafu {
290                        path: &req.location,
291                    })
292                    .and_then(|files| {
293                        Ok(DatabaseExportSummary {
294                            rows,
295                            output_files: files
296                                .into_iter()
297                                .map(|file| {
298                                    DatabaseExportFile::new(&req.location, &file, "")
299                                        .map(|f| f.location)
300                                })
301                                .collect::<Result<_>>()?,
302                        })
303                    }),
304                Err(error) => Err(error),
305            };
306            if result.is_err()
307                && let Err(error) = packed.abort().await
308            {
309                common_telemetry::warn!(error; "Failed to abort Metric pack");
310            }
311            result
312        } else {
313            Ok(DatabaseExportSummary {
314                rows: rows?,
315                output_files,
316            })
317        }
318    }
319}
320
321async fn run_database_export_jobs<J, F: Future<Output = Result<usize>>>(
322    jobs: impl IntoIterator<Item = J>,
323    parallelism: usize,
324    cancellation: &CancellationToken,
325    mut run: impl FnMut(J, CancellationToken) -> F,
326) -> Result<usize> {
327    let token = CancellationToken::new();
328    let mut jobs = jobs.into_iter();
329    let mut active = FuturesUnordered::new();
330    let mut first_error = None;
331    let mut rows = 0;
332    loop {
333        while first_error.is_none()
334            && !cancellation.is_cancelled()
335            && !token.is_cancelled()
336            && active.len() < parallelism.max(1)
337        {
338            let Some(job) = jobs.next() else { break };
339            active.push(run(job, token.clone()));
340        }
341        if first_error.is_none() && cancellation.is_cancelled() {
342            first_error = Some(error::DatabaseExportCancelledSnafu.build());
343            token.cancel();
344        }
345        if active.is_empty() {
346            break;
347        }
348        let result = tokio::select! {
349            biased;
350            _ = cancellation.cancelled(), if first_error.is_none() => {
351                first_error = Some(error::DatabaseExportCancelledSnafu.build());
352                token.cancel();
353                continue;
354            }
355            result = active.next() => result,
356        };
357        match result {
358            Some(Ok(count)) => rows += count,
359            Some(Err(err)) => {
360                if !cancellation.is_cancelled() || first_error.is_none() {
361                    retain_error(&mut first_error, err);
362                }
363                token.cancel();
364            }
365            None => break,
366        }
367    }
368    match first_error {
369        Some(err) => Err(err),
370        None => Ok(rows),
371    }
372}
373
374#[cfg(test)]
375mod tests {
376    use std::sync::Arc;
377    use std::sync::atomic::{AtomicUsize, Ordering};
378
379    use tokio::sync::{Semaphore, mpsc};
380
381    use super::*;
382
383    #[tokio::test]
384    async fn bounded_admission_and_drain() {
385        for cancel in [false, true] {
386            let cancellation = CancellationToken::new();
387            let start_error = Arc::new(Semaphore::new(0));
388            let finish_io = Arc::new(Semaphore::new(0));
389            let (started, mut receiver) = mpsc::unbounded_channel();
390            let finished = Arc::new(AtomicUsize::new(0));
391            let task = tokio::spawn({
392                let cancellation = cancellation.clone();
393                let start_error = start_error.clone();
394                let finish_io = finish_io.clone();
395                let finished = finished.clone();
396                async move {
397                    run_database_export_jobs(0..4, 2, &cancellation, |job, token| {
398                        let started = started.clone();
399                        let start_error = start_error.clone();
400                        let finish_io = finish_io.clone();
401                        let finished = finished.clone();
402                        async move {
403                            started.send(job).unwrap();
404                            if job == 0 {
405                                if cancel {
406                                    token.cancelled().await;
407                                } else {
408                                    start_error.acquire().await.unwrap().forget();
409                                }
410                                return InvalidDatabaseExportSnafu {
411                                    reason: "first error",
412                                }
413                                .fail();
414                            }
415                            // Already-started ordinary I/O drains after cancellation.
416                            token.cancelled().await;
417                            started.send(10).unwrap();
418                            finish_io.acquire().await.unwrap().forget();
419                            finished.fetch_add(1, Ordering::SeqCst);
420                            InvalidDatabaseExportSnafu {
421                                reason: "drain error",
422                            }
423                            .fail()
424                        }
425                    })
426                    .await
427                }
428            });
429            assert_eq!(receiver.recv().await, Some(0));
430            assert_eq!(receiver.recv().await, Some(1));
431            assert!(receiver.try_recv().is_err());
432            if cancel {
433                cancellation.cancel();
434            } else {
435                start_error.add_permits(1);
436            }
437            assert_eq!(receiver.recv().await, Some(10));
438            assert!(!task.is_finished());
439            finish_io.add_permits(1);
440            let err = task.await.unwrap().unwrap_err();
441            if cancel {
442                assert!(matches!(err, error::Error::DatabaseExportCancelled { .. }));
443            } else {
444                assert!(
445                    matches!(err, error::Error::InvalidDatabaseExport { reason } if reason == "first error")
446                );
447            }
448            assert_eq!(finished.load(Ordering::SeqCst), 1);
449            assert_eq!(receiver.recv().await, None);
450        }
451    }
452
453    #[tokio::test]
454    async fn cancellation_before_admission_and_successful_refill() {
455        let token = CancellationToken::new();
456        token.cancel();
457        let result = run_database_export_jobs(0..4, 2, &token, |_, _| async {
458            panic!("cancelled job admitted")
459        })
460        .await;
461        assert!(matches!(
462            result,
463            Err(error::Error::DatabaseExportCancelled { .. })
464        ));
465        let started = AtomicUsize::new(0);
466        let result = run_database_export_jobs(0..7, 2, &CancellationToken::new(), |job, _| {
467            started.fetch_add(1, Ordering::SeqCst);
468            async move { Ok(job) }
469        })
470        .await
471        .unwrap();
472        assert_eq!(result, 21);
473        assert_eq!(started.load(Ordering::SeqCst), 7);
474    }
475}
operator/statement/export_database.rs

operator/statement/
export_database.rs
