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

operator/statement/
export_database.rs
