frontend/instance/
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
15use auth::{
16    PermissionAction, PermissionChecker, PermissionCheckerRef, PermissionReq,
17    PermissionTableTarget, PermissionTableTargets,
18};
19use common_query::Output;
20use common_telemetry::error;
21use operator::statement::export_database::{DatabaseExportSummary, PreparedDatabaseExport};
22use session::context::QueryContextRef;
23use snafu::ResultExt;
24use sql::ast::{Ident, ObjectName};
25use sql::statements::OptionMap;
26use sql::statements::copy::{Copy, CopyDatabase, CopyDatabaseArgument};
27use sql::statements::statement::Statement;
28use table::requests::CopyDatabaseRequest;
29use tokio_util::sync::CancellationToken;
30
31use crate::error::{PermissionSnafu, Result};
32use crate::instance::Instance;
33
34pub(crate) fn parse_metric_export_requested(options: &OptionMap) -> Result<bool> {
35    match options.get("experimental_metric_export") {
36        None | Some("false") => Ok(false),
37        Some("true") => Ok(true),
38        Some(_) => Err(operator::error::InvalidDatabaseExportSnafu {
39            reason: "experimental_metric_export must be true or false",
40        }
41        .build()
42        .into()),
43    }
44}
45
46impl Instance {
47    pub(crate) async fn copy_metric_database(
48        &self,
49        arg: CopyDatabaseArgument,
50        ctx: QueryContextRef,
51    ) -> Result<Output> {
52        let req = operator::statement::to_copy_database_request(arg, &ctx)?;
53        let plan = self.prepare_database_export(req, None, &ctx).await?;
54        let cancellation = CancellationToken::new();
55        let _guard = cancellation.clone().drop_guard();
56        let executor = self.statement_executor.clone();
57        // Dropping the request only cancels; the runtime task retains and drains started I/O.
58        let task = common_runtime::spawn_query(async move {
59            let result = executor.export_database(plan, &cancellation, ctx).await;
60            if let Err(err) = &result {
61                error!(err; "Experimental database export failed after draining started work");
62            }
63            result
64        });
65        let summary = task.await.context(operator::error::JoinTaskSnafu)??;
66        Ok(Output::new_with_affected_rows(summary.rows))
67    }
68
69    #[allow(dead_code)]
70    async fn export_database(
71        &self,
72        req: CopyDatabaseRequest,
73        names: Option<&[String]>,
74        cancellation: &CancellationToken,
75        ctx: QueryContextRef,
76    ) -> Result<DatabaseExportSummary> {
77        if cancellation.is_cancelled() {
78            return Err(operator::error::DatabaseExportCancelledSnafu.build().into());
79        }
80        let plan = self.prepare_database_export(req, names, &ctx).await?;
81        Ok(self
82            .statement_executor
83            .export_database(plan, cancellation, ctx)
84            .await?)
85    }
86
87    /// Exercise the internal authenticated entry without exposing SQL/CLI activation.
88    #[cfg(any(test, feature = "testing"))]
89    pub async fn export_database_for_test(
90        &self,
91        req: CopyDatabaseRequest,
92        names: Option<&[String]>,
93        cancellation: &CancellationToken,
94        ctx: QueryContextRef,
95    ) -> Result<DatabaseExportSummary> {
96        self.export_database(req, names, cancellation, ctx).await
97    }
98
99    async fn prepare_database_export(
100        &self,
101        req: CopyDatabaseRequest,
102        names: Option<&[String]>,
103        ctx: &QueryContextRef,
104    ) -> Result<PreparedDatabaseExport> {
105        let stmt = Statement::Copy(Copy::CopyDatabase(CopyDatabase::To(CopyDatabaseArgument {
106            database_name: ObjectName::from(vec![
107                Ident::new(&req.catalog_name),
108                Ident::new(&req.schema_name),
109            ]),
110            with: req.with.clone().into(),
111            connection: req.connection.clone().into(),
112            location: req.location.clone(),
113        })));
114        self.plugins
115            .get::<PermissionCheckerRef>()
116            .as_ref()
117            .check_permission_with_context(
118                ctx.current_user(),
119                PermissionReq::SqlStatement(&stmt),
120                Some(&ctx.current_schema()),
121            )
122            .context(PermissionSnafu)?;
123        let tables = self
124            .statement_executor
125            .capture_database_export_tables(&req, names, ctx)
126            .await?;
127        let targets = PermissionTableTargets::resolved(
128            tables
129                .iter()
130                .map(|table| {
131                    let info = table.table_info();
132                    PermissionTableTarget::new(&info.catalog_name, &info.schema_name, &info.name)
133                })
134                .collect(),
135        );
136        self.check_table_permission(ctx, PermissionReq::SqlStatement(&stmt), targets.clone())
137            .context(PermissionSnafu)?;
138        // COPY is classified as a write by existing permission checkers.
139        self.check_table_permission(
140            ctx,
141            PermissionReq::Action(PermissionAction::read("database.export")),
142            targets,
143        )
144        .context(PermissionSnafu)?;
145        Ok(self
146            .statement_executor
147            .prepare_database_export(req, tables)
148            .await?)
149    }
150}
frontend/instance/export_database.rs

frontend/instance/
export_database.rs
