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

frontend/instance/
export_database.rs
