operator/
delete.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::collections::HashSet;
16use std::sync::Arc;
17use std::{iter, mem};
18
19use api::v1::region::{DeleteRequests as RegionDeleteRequests, RegionRequestHeader};
20use api::v1::{DeleteRequests, RowDeleteRequests};
21use catalog::CatalogManagerRef;
22use common_meta::node_manager::{AffectedRows, NodeManagerRef};
23use common_query::Output;
24use common_telemetry::tracing_context::TracingContext;
25use partition::manager::PartitionRuleManagerRef;
26use session::context::QueryContextRef;
27use snafu::{OptionExt, ResultExt, ensure};
28use table::TableRef;
29use table::requests::DeleteRequest as TableDeleteRequest;
30use tokio_util::sync::CancellationToken;
31
32use crate::error::{
33    CatalogSnafu, InvalidDeleteRequestSnafu, JoinTaskSnafu, MissingTimeIndexColumnSnafu,
34    RequestDeletesSnafu, Result, TableNotFoundSnafu,
35};
36use crate::region_req_factory::RegionRequestFactory;
37use crate::req_convert::common::preprocess_row_delete_requests;
38use crate::req_convert::delete::{ColumnToRow, RowToRegion, TableToRegion};
39use crate::write_dispatcher::{WriteDispatcher, WriteDispatcherRef};
40
41pub struct Deleter {
42    catalog_manager: CatalogManagerRef,
43    partition_manager: PartitionRuleManagerRef,
44    write_dispatcher: WriteDispatcherRef,
45}
46
47pub type DeleterRef = Arc<Deleter>;
48
49impl Deleter {
50    /// Creates a row deleter with the default write dispatcher.
51    pub fn new(
52        catalog_manager: CatalogManagerRef,
53        partition_manager: PartitionRuleManagerRef,
54        node_manager: NodeManagerRef,
55    ) -> Self {
56        Self {
57            write_dispatcher: Arc::new(WriteDispatcher::new(
58                node_manager,
59                partition_manager.clone(),
60            )),
61            catalog_manager,
62            partition_manager,
63        }
64    }
65
66    /// Installs the dispatcher shared with the other write paths.
67    pub fn with_write_dispatcher(mut self, dispatcher: WriteDispatcherRef) -> Self {
68        self.write_dispatcher = dispatcher;
69        self
70    }
71
72    pub async fn handle_column_deletes(
73        &self,
74        requests: DeleteRequests,
75        ctx: QueryContextRef,
76    ) -> Result<Output> {
77        let row_deletes = ColumnToRow::convert(requests)?;
78        self.handle_row_deletes(row_deletes, ctx).await
79    }
80
81    pub async fn handle_row_deletes(
82        &self,
83        mut requests: RowDeleteRequests,
84        ctx: QueryContextRef,
85    ) -> Result<Output> {
86        preprocess_row_delete_requests(&mut requests.deletes)?;
87        // remove empty requests
88        requests.deletes.retain(|req| {
89            req.rows
90                .as_ref()
91                .map(|r| !r.rows.is_empty())
92                .unwrap_or_default()
93        });
94        validate_column_count_match(&requests)?;
95
96        let requests = self.trim_columns(requests, &ctx).await?;
97        let deletes = RowToRegion::new(
98            self.catalog_manager.as_ref(),
99            self.partition_manager.as_ref(),
100            &ctx,
101        )
102        .convert(requests)
103        .await?;
104
105        let affected_rows = self.do_request(deletes, &ctx).await?;
106
107        Ok(Output::new_with_affected_rows(affected_rows))
108    }
109
110    pub async fn handle_table_delete(
111        &self,
112        request: TableDeleteRequest,
113        ctx: QueryContextRef,
114    ) -> Result<AffectedRows> {
115        let catalog = request.catalog_name.as_str();
116        let schema = request.schema_name.as_str();
117        let table = request.table_name.as_str();
118        let table = self.get_table(catalog, schema, table).await?;
119        let table_info = table.table_info();
120
121        let deletes = TableToRegion::new(&table_info, &self.partition_manager)
122            .convert(request)
123            .await?;
124
125        let affected_rows = self.do_request(deletes, &ctx).await?;
126        Ok(affected_rows as _)
127    }
128}
129
130impl Deleter {
131    async fn do_request(
132        &self,
133        requests: RegionDeleteRequests,
134        ctx: &QueryContextRef,
135    ) -> Result<AffectedRows> {
136        let request_factory = RegionRequestFactory::new(RegionRequestHeader {
137            tracing_context: TracingContext::from_current_span().to_w3c(),
138            dbname: ctx.get_db_string(),
139            ..Default::default()
140        });
141
142        let dispatcher = self.write_dispatcher.clone();
143        let request = request_factory.build_delete(requests);
144        let writes = dispatcher.prepare(vec![request]).await?;
145        // Dropping the caller stops retries without abandoning in-flight writes.
146        let cancellation = CancellationToken::new();
147        let _guard = cancellation.clone().drop_guard();
148        let affected_rows = common_runtime::spawn_global(async move {
149            dispatcher.execute(writes, None, Some(&cancellation)).await
150        })
151        .await
152        .context(JoinTaskSnafu)?
153        .context(RequestDeletesSnafu)?
154        .affected_rows;
155        crate::metrics::DIST_DELETE_ROW_COUNT
156            .with_label_values(&[ctx.get_db_string().as_str()])
157            .inc_by(affected_rows as u64);
158        Ok(affected_rows)
159    }
160
161    async fn trim_columns(
162        &self,
163        mut requests: RowDeleteRequests,
164        ctx: &QueryContextRef,
165    ) -> Result<RowDeleteRequests> {
166        for req in &mut requests.deletes {
167            let catalog = ctx.current_catalog();
168            let schema = ctx.current_schema();
169            let table = self.get_table(catalog, &schema, &req.table_name).await?;
170            let key_column_names = self.key_column_names(&table)?;
171
172            let rows = req.rows.as_mut().unwrap();
173            let all_key = rows
174                .schema
175                .iter()
176                .all(|column| key_column_names.contains(&column.column_name));
177            if all_key {
178                continue;
179            }
180
181            for row in &mut rows.rows {
182                let source_values = mem::take(&mut row.values);
183                row.values = rows
184                    .schema
185                    .iter()
186                    .zip(source_values)
187                    .filter_map(|(column, v)| {
188                        key_column_names.contains(&column.column_name).then_some(v)
189                    })
190                    .collect();
191            }
192            rows.schema
193                .retain(|column| key_column_names.contains(&column.column_name));
194        }
195        Ok(requests)
196    }
197
198    fn key_column_names(&self, table: &TableRef) -> Result<HashSet<String>> {
199        let time_index = table
200            .schema()
201            .timestamp_column()
202            .with_context(|| table::error::MissingTimeIndexColumnSnafu {
203                table_name: table.table_info().name.clone(),
204            })
205            .context(MissingTimeIndexColumnSnafu)?
206            .name
207            .clone();
208
209        let key_column_names = table
210            .table_info()
211            .meta
212            .row_key_column_names()
213            .cloned()
214            .chain(iter::once(time_index))
215            .collect();
216
217        Ok(key_column_names)
218    }
219
220    async fn get_table(&self, catalog: &str, schema: &str, table: &str) -> Result<TableRef> {
221        self.catalog_manager
222            .table(catalog, schema, table, None)
223            .await
224            .context(CatalogSnafu)?
225            .with_context(|| TableNotFoundSnafu {
226                table_name: common_catalog::format_full_table_name(catalog, schema, table),
227            })
228    }
229}
230
231fn validate_column_count_match(requests: &RowDeleteRequests) -> Result<()> {
232    for request in &requests.deletes {
233        let rows = request.rows.as_ref().unwrap();
234        let column_count = rows.schema.len();
235        rows.rows.iter().try_for_each(|r| {
236            ensure!(
237                r.values.len() == column_count,
238                InvalidDeleteRequestSnafu {
239                    reason: format!(
240                        "column count mismatch, columns: {}, values: {}",
241                        column_count,
242                        r.values.len()
243                    )
244                }
245            );
246            Ok(())
247        })?;
248    }
249    Ok(())
250}
operator/delete.rs

operator/
delete.rs
