operator/req_convert/insert/
stmt_to_region.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 api::helper::{ColumnDataTypeWrapper, to_grpc_value};
16use api::v1::column_def::options_from_column_schema;
17use api::v1::region::InsertRequests as RegionInsertRequests;
18use api::v1::{ColumnSchema as GrpcColumnSchema, Row, Rows};
19use catalog::CatalogManager;
20use common_time::Timezone;
21use datatypes::schema::{ColumnSchema, SchemaRef};
22use datatypes::value::Value;
23use partition::manager::PartitionRuleManager;
24use session::context::{QueryContext, QueryContextRef};
25use snafu::{OptionExt, ResultExt, ensure};
26use sql::ast::ObjectNamePartExt;
27use sql::statements::insert::Insert;
28use sqlparser::ast::{ObjectName, Value as SqlValue};
29use table::TableRef;
30use table::metadata::TableInfoRef;
31
32use crate::error::{
33    CatalogSnafu, ColumnDataTypeSnafu, ColumnDefaultValueSnafu, ColumnNoneDefaultValueSnafu,
34    ColumnNotFoundSnafu, InvalidSqlSnafu, MissingInsertBodySnafu, ParseSqlSnafu, Result,
35    SchemaReadOnlySnafu, TableNotFoundSnafu, TableReadOnlySnafu,
36};
37use crate::insert::InstantAndNormalInsertRequests;
38use crate::req_convert::common::partitioner::Partitioner;
39use crate::req_convert::insert::semantic_type;
40
41const DEFAULT_PLACEHOLDER_VALUE: &str = "default";
42
43pub struct StatementToRegion<'a> {
44    catalog_manager: &'a dyn CatalogManager,
45    partition_manager: &'a PartitionRuleManager,
46    ctx: &'a QueryContext,
47}
48
49impl<'a> StatementToRegion<'a> {
50    pub fn new(
51        catalog_manager: &'a dyn CatalogManager,
52        partition_manager: &'a PartitionRuleManager,
53        ctx: &'a QueryContext,
54    ) -> Self {
55        Self {
56            catalog_manager,
57            partition_manager,
58            ctx,
59        }
60    }
61
62    pub async fn convert(
63        &self,
64        stmt: &Insert,
65        query_ctx: &QueryContextRef,
66    ) -> Result<(InstantAndNormalInsertRequests, TableInfoRef)> {
67        let (rows, table_info) = self.prepare(stmt, query_ctx).await?;
68        let requests = self.partition(rows, table_info.clone(), query_ctx).await?;
69        Ok((requests, table_info))
70    }
71
72    /// Routes already prepared rows while retaining TTL and WAL behavior.
73    pub async fn partition(
74        &self,
75        rows: Rows,
76        table_info: TableInfoRef,
77        query_ctx: &QueryContextRef,
78    ) -> Result<InstantAndNormalInsertRequests> {
79        let requests = Partitioner::new(self.partition_manager)
80            .partition_insert_requests(&table_info, rows, query_ctx.skip_wal())
81            .await?;
82        let requests = RegionInsertRequests { requests };
83        if table_info.is_ttl_instant_table() {
84            Ok(InstantAndNormalInsertRequests {
85                normal_requests: Default::default(),
86                instant_requests: requests,
87            })
88        } else {
89            Ok(InstantAndNormalInsertRequests {
90                normal_requests: requests,
91                instant_requests: Default::default(),
92            })
93        }
94    }
95
96    /// Resolves SQL values and their table schema without partition routing.
97    pub async fn prepare(
98        &self,
99        stmt: &Insert,
100        query_ctx: &QueryContextRef,
101    ) -> Result<(Rows, TableInfoRef)> {
102        let name = stmt.table_name().context(ParseSqlSnafu)?;
103        let (catalog, schema, table_name) = self.get_full_name(name)?;
104        let table = self.get_table(&catalog, &schema, &table_name).await?;
105        let table_schema = table.schema();
106
107        ensure!(
108            !common_catalog::consts::is_readonly_schema(&schema),
109            SchemaReadOnlySnafu { name: schema }
110        );
111        // The computed entity-graph tables overlay a writable schema, so the
112        // schema-level check above does not cover them.
113        ensure!(
114            !common_catalog::consts::is_readonly_table(&schema, &table_name),
115            TableReadOnlySnafu { name: table_name }
116        );
117
118        let column_names = column_names(stmt, &table_schema);
119        let column_count = column_names.len();
120
121        let sql_rows = stmt.values_body().context(MissingInsertBodySnafu)?;
122        let row_count = sql_rows.len();
123
124        sql_rows.iter().try_for_each(|r| {
125            ensure!(
126                r.len() == column_count,
127                InvalidSqlSnafu {
128                    err_msg: format!(
129                        "column count mismatch, columns: {}, values: {}",
130                        column_count,
131                        r.len()
132                    )
133                }
134            );
135            Ok(())
136        })?;
137
138        let mut rows = vec![
139            Row {
140                values: Vec::with_capacity(column_count)
141            };
142            row_count
143        ];
144
145        fn find_insert_columns<'a>(
146            table: &'a TableRef,
147            column_names: &[&String],
148        ) -> Result<Vec<&'a ColumnSchema>> {
149            let schema = table.schema_ref();
150            column_names
151                .iter()
152                .map(|name| {
153                    schema
154                        .column_schema_by_name(name)
155                        .context(ColumnNotFoundSnafu { msg: *name })
156                })
157                .collect::<Result<Vec<_>>>()
158        }
159
160        let insert_columns = find_insert_columns(&table, &column_names)?;
161        let converter = SqlRowConverter::new(&insert_columns, query_ctx);
162        let value_rows = converter.convert(&sql_rows)?;
163        for (i, row) in value_rows.into_iter().enumerate() {
164            for value in row {
165                let grpc_value = to_grpc_value(value);
166                rows[i].values.push(grpc_value);
167            }
168        }
169
170        let table_info = table.table_info();
171        let mut schema = Vec::with_capacity(column_count);
172        for column_schema in insert_columns {
173            let (datatype, datatype_extension) =
174                ColumnDataTypeWrapper::try_from(column_schema.data_type.clone())
175                    .context(ColumnDataTypeSnafu)?
176                    .to_parts();
177
178            let column_name = &column_schema.name;
179            let semantic_type = semantic_type(&table_info, column_name)?;
180
181            let grpc_column_schema = GrpcColumnSchema {
182                column_name: column_name.clone(),
183                datatype: datatype.into(),
184                semantic_type: semantic_type.into(),
185                datatype_extension,
186                options: options_from_column_schema(column_schema),
187            };
188            schema.push(grpc_column_schema);
189        }
190
191        Ok((Rows { schema, rows }, table_info))
192    }
193
194    async fn get_table(&self, catalog: &str, schema: &str, table: &str) -> Result<TableRef> {
195        self.catalog_manager
196            .table(catalog, schema, table, None)
197            .await
198            .context(CatalogSnafu)?
199            .with_context(|| TableNotFoundSnafu {
200                table_name: common_catalog::format_full_table_name(catalog, schema, table),
201            })
202    }
203
204    fn get_full_name(&self, obj_name: &ObjectName) -> Result<(String, String, String)> {
205        match &obj_name.0[..] {
206            [table] => Ok((
207                self.ctx.current_catalog().to_owned(),
208                self.ctx.current_schema(),
209                table.to_string_unquoted(),
210            )),
211            [schema, table] => Ok((
212                self.ctx.current_catalog().to_owned(),
213                schema.to_string_unquoted(),
214                table.to_string_unquoted(),
215            )),
216            [catalog, schema, table] => Ok((
217                catalog.to_string_unquoted(),
218                schema.to_string_unquoted(),
219                table.to_string_unquoted(),
220            )),
221            _ => InvalidSqlSnafu {
222                err_msg: format!(
223                    "expect table name to be <catalog>.<schema>.<table>, <schema>.<table> or <table>, actual: {obj_name}",
224                ),
225            }.fail(),
226        }
227    }
228}
229
230struct SqlRowConverter<'a, 'b> {
231    insert_columns: &'a [&'a ColumnSchema],
232    query_context: &'b QueryContextRef,
233}
234
235impl<'a, 'b> SqlRowConverter<'a, 'b> {
236    fn new(insert_columns: &'a [&'a ColumnSchema], query_context: &'b QueryContextRef) -> Self {
237        Self {
238            insert_columns,
239            query_context,
240        }
241    }
242
243    fn convert(&self, sql_rows: &[Vec<SqlValue>]) -> Result<Vec<Vec<Value>>> {
244        let timezone = Some(&self.query_context.timezone());
245        let auto_string_to_numeric = self.query_context.auto_string_to_numeric();
246
247        let mut value_rows = Vec::with_capacity(sql_rows.len());
248        for sql_row in sql_rows {
249            let mut value_row = Vec::with_capacity(self.insert_columns.len());
250
251            for (insert_column, sql_value) in self.insert_columns.iter().zip(sql_row) {
252                let value =
253                    sql_value_to_value(insert_column, sql_value, timezone, auto_string_to_numeric)?;
254                value_row.push(value);
255            }
256            value_rows.push(value_row);
257        }
258        Ok(value_rows)
259    }
260}
261
262fn column_names<'a>(stmt: &'a Insert, table_schema: &'a SchemaRef) -> Vec<&'a String> {
263    if !stmt.columns().is_empty() {
264        stmt.columns()
265    } else {
266        table_schema
267            .column_schemas()
268            .iter()
269            .map(|column| &column.name)
270            .collect()
271    }
272}
273
274/// Converts SQL value to gRPC value according to the column schema.
275/// If `auto_string_to_numeric` is true, tries to cast the string value to numeric values,
276/// and fills the default value if the cast fails.
277fn sql_value_to_value(
278    column_schema: &ColumnSchema,
279    sql_val: &SqlValue,
280    timezone: Option<&Timezone>,
281    auto_string_to_numeric: bool,
282) -> Result<Value> {
283    let column = &column_schema.name;
284    let value = if replace_default(sql_val) {
285        let default_value = column_schema
286            .create_default()
287            .context(ColumnDefaultValueSnafu {
288                column: column.clone(),
289            })?;
290
291        default_value.context(ColumnNoneDefaultValueSnafu {
292            column: column.clone(),
293        })?
294    } else {
295        common_sql::convert::sql_value_to_value(
296            column_schema,
297            sql_val,
298            timezone,
299            None,
300            auto_string_to_numeric,
301        )
302        .context(crate::error::SqlCommonSnafu)?
303    };
304    Ok(value)
305}
306
307fn replace_default(sql_val: &SqlValue) -> bool {
308    matches!(sql_val, SqlValue::Placeholder(s) if s.to_lowercase() == DEFAULT_PLACEHOLDER_VALUE)
309}
operator/req_convert/insert/stmt_to_region.rs

operator/req_convert/insert/
stmt_to_region.rs
