frontend/instance/
import_packed.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::{PermissionReq, PermissionTableTarget, PermissionTableTargets};
16use common_query::Output;
17use session::context::QueryContextRef;
18use snafu::ResultExt;
19use sql::statements::copy::CopyDatabaseArgument;
20use sql::statements::statement::Statement;
21use tokio_util::sync::CancellationToken;
22
23use crate::error::{PermissionSnafu, Result};
24use crate::instance::Instance;
25
26impl Instance {
27    pub(crate) async fn copy_packed_database(
28        &self,
29        arg: CopyDatabaseArgument,
30        stmt: &Statement,
31        ctx: QueryContextRef,
32    ) -> Result<Output> {
33        let req = operator::statement::to_copy_database_request(arg, &ctx)?;
34        let plan = self
35            .statement_executor
36            .prepare_packed_import(req, &ctx)
37            .await?;
38        let targets = PermissionTableTargets::resolved(
39            plan.tables
40                .iter()
41                .map(|t| {
42                    let info = t.table_info();
43                    PermissionTableTarget::new(&info.catalog_name, &info.schema_name, &info.name)
44                })
45                .collect(),
46        );
47        self.check_table_permission(&ctx, PermissionReq::SqlStatement(stmt), targets)
48            .context(PermissionSnafu)?;
49        let executor = self.statement_executor.clone();
50        run_packed_import(move |cancellation| async move {
51            executor.import_packed(plan, &cancellation, ctx).await
52        })
53        .await
54    }
55}
56
57async fn run_packed_import<F, Fut>(worker: F) -> Result<Output>
58where
59    F: FnOnce(CancellationToken) -> Fut,
60    Fut: std::future::Future<Output = operator::error::Result<Output>> + Send + 'static,
61{
62    let cancellation = CancellationToken::new();
63    let _guard = cancellation.clone().drop_guard();
64    // Transport timeout/drop signals cancellation; the worker retains ownership
65    // until its started inserts finish, even after the response is gone.
66    common_runtime::spawn_query(worker(cancellation))
67        .await
68        .context(operator::error::JoinTaskSnafu)?
69        .map_err(Into::into)
70}
71
72#[cfg(test)]
73mod tests {
74    use std::time::Duration;
75
76    use tokio::sync::oneshot;
77
78    use super::*;
79
80    #[tokio::test]
81    async fn packed_request_drop_keeps_worker_alive_until_cleanup() {
82        let (started_tx, started_rx) = oneshot::channel();
83        let (cancelled_tx, cancelled_rx) = oneshot::channel();
84        let (drain_tx, drain_rx) = oneshot::channel();
85        let (finished_tx, mut finished_rx) = oneshot::channel();
86        let mut request = Box::pin(run_packed_import(move |cancellation| async move {
87            started_tx.send(()).unwrap();
88            cancellation.cancelled().await;
89            cancelled_tx.send(()).unwrap();
90            // Model an already-started insert that must finish before worker exit.
91            drain_rx.await.unwrap();
92            finished_tx.send(()).unwrap();
93            Ok(Output::new_with_affected_rows(1))
94        }));
95        tokio::select! {
96            _ = started_rx => {}
97            _ = &mut request => panic!("worker must wait for cancellation"),
98        }
99        drop(request);
100        tokio::time::timeout(Duration::from_secs(5), cancelled_rx)
101            .await
102            .unwrap()
103            .unwrap();
104        assert!(matches!(
105            finished_rx.try_recv(),
106            Err(oneshot::error::TryRecvError::Empty)
107        ));
108        drain_tx.send(()).unwrap();
109        tokio::time::timeout(Duration::from_secs(5), finished_rx)
110            .await
111            .unwrap()
112            .unwrap();
113    }
114}
frontend/instance/import_packed.rs

frontend/instance/
import_packed.rs
