flow/adapter/
util.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
15//! Util functions for adapter
16
17use std::sync::Arc;
18
19use api::helper::ColumnDataTypeWrapper;
20use api::v1::column_def::options_from_column_schema;
21use api::v1::{ColumnDataType, ColumnDataTypeExtension, CreateTableExpr, SemanticType};
22use common_error::ext::BoxedError;
23use common_meta::key::table_info::TableInfoValue;
24use common_meta::rpc::ddl::TriggerReason;
25use datatypes::prelude::ConcreteDataType;
26use datatypes::schema::{ColumnDefaultConstraint, ColumnSchema};
27use itertools::Itertools;
28use operator::expr_helper;
29use session::context::QueryContextBuilder;
30use snafu::{OptionExt, ResultExt};
31use table::table_reference::TableReference;
32
33use crate::StreamingEngine;
34use crate::adapter::table_source::TableDesc;
35use crate::adapter::{AUTO_CREATED_PLACEHOLDER_TS_COL, TableName, WorkerHandle};
36use crate::error::{Error, ExternalSnafu, UnexpectedSnafu};
37use crate::repr::{ColumnType, RelationDesc, RelationType};
38impl StreamingEngine {
39    /// Get a worker handle for creating flow, using round robin to select a worker
40    pub(crate) async fn get_worker_handle_for_create_flow(&self) -> &WorkerHandle {
41        let use_idx = {
42            let mut selector = self.worker_selector.lock().await;
43            if *selector >= self.worker_handles.len() {
44                *selector = 0
45            };
46            let use_idx = *selector;
47            *selector += 1;
48            use_idx
49        };
50        // Safety: selector is always in bound
51        &self.worker_handles[use_idx]
52    }
53
54    /// Create table from given schema(will adjust to add auto column if needed), return true if table is created
55    pub(crate) async fn create_table_from_relation(
56        &self,
57        flow_name: &str,
58        table_name: &TableName,
59        relation_desc: &RelationDesc,
60    ) -> Result<bool, Error> {
61        if self.fetch_table_pk_schema(table_name).await?.is_some() {
62            return Ok(false);
63        }
64        let (pks, tys, _) = self.adjust_auto_created_table_schema(relation_desc).await?;
65
66        //create sink table using pks, column types and is_ts_auto
67
68        let proto_schema = column_schemas_to_proto(tys.clone(), &pks)?;
69
70        // create sink table
71        let create_expr = expr_helper::create_table_expr_by_column_schemas(
72            &TableReference {
73                catalog: &table_name[0],
74                schema: &table_name[1],
75                table: &table_name[2],
76            },
77            &proto_schema,
78            "mito",
79            Some(&format!("Sink table for flow {}", flow_name)),
80        )
81        .map_err(BoxedError::new)
82        .context(ExternalSnafu)?;
83
84        self.submit_create_sink_table_ddl(create_expr).await?;
85        Ok(true)
86    }
87
88    /// Try fetch table with adjusted schema(added auto column if needed)
89    pub(crate) async fn try_fetch_existing_table(
90        &self,
91        table_name: &TableName,
92    ) -> Result<Option<(bool, Vec<api::v1::ColumnSchema>)>, Error> {
93        if let Some((primary_keys, time_index, schema)) =
94            self.fetch_table_pk_schema(table_name).await?
95        {
96            // check if the last column is the auto created timestamp column, hence the table is auto created from
97            // flow's plan type
98            let is_auto_create = {
99                let correct_name = schema
100                    .last()
101                    .map(|s| s.name == AUTO_CREATED_PLACEHOLDER_TS_COL)
102                    .unwrap_or(false);
103                let correct_time_index = time_index == Some(schema.len() - 1);
104                correct_name && correct_time_index
105            };
106            let proto_schema = column_schemas_to_proto(schema, &primary_keys)?;
107            Ok(Some((is_auto_create, proto_schema)))
108        } else {
109            Ok(None)
110        }
111    }
112
113    /// submit a create table ddl
114    pub(crate) async fn submit_create_sink_table_ddl(
115        &self,
116        mut create_table: CreateTableExpr,
117    ) -> Result<(), Error> {
118        let stmt_exec = {
119            self.frontend_invoker
120                .read()
121                .await
122                .as_ref()
123                .map(|f| f.statement_executor())
124        }
125        .context(UnexpectedSnafu {
126            reason: "Failed to get statement executor",
127        })?;
128        let ctx = Arc::new(
129            QueryContextBuilder::default()
130                .current_catalog(create_table.catalog_name.clone())
131                .current_schema(create_table.schema_name.clone())
132                .build(),
133        );
134        stmt_exec
135            .create_table_inner(&mut create_table, None, ctx, TriggerReason::AutoCreate)
136            .await
137            .map_err(BoxedError::new)
138            .context(ExternalSnafu)?;
139
140        Ok(())
141    }
142}
143
144pub fn table_info_value_to_relation_desc(
145    table_info_value: TableInfoValue,
146) -> Result<TableDesc, Error> {
147    let raw_schema = table_info_value.table_info.meta.schema;
148    let (column_types, col_names): (Vec<_>, Vec<_>) = raw_schema
149        .column_schemas()
150        .to_vec()
151        .into_iter()
152        .map(|col| {
153            (
154                ColumnType {
155                    nullable: col.is_nullable(),
156                    scalar_type: col.data_type,
157                },
158                Some(col.name),
159            )
160        })
161        .unzip();
162
163    let key = table_info_value.table_info.meta.primary_key_indices;
164    let keys = vec![crate::repr::Key::from(key)];
165
166    let time_index = raw_schema.timestamp_index();
167    let relation_desc = RelationDesc {
168        typ: RelationType {
169            column_types,
170            keys,
171            time_index,
172            // by default table schema's column are all non-auto
173            auto_columns: vec![],
174        },
175        names: col_names,
176    };
177    let default_values = raw_schema
178        .column_schemas()
179        .iter()
180        .map(|c| {
181            c.default_constraint().cloned().or_else(|| {
182                if c.is_nullable() {
183                    Some(ColumnDefaultConstraint::null_value())
184                } else {
185                    None
186                }
187            })
188        })
189        .collect_vec();
190
191    Ok(TableDesc::new(relation_desc, default_values))
192}
193
194pub fn from_proto_to_data_type(
195    column_schema: &api::v1::ColumnSchema,
196) -> Result<ConcreteDataType, Error> {
197    let wrapper = ColumnDataTypeWrapper::try_new(
198        column_schema.datatype,
199        column_schema.datatype_extension.clone(),
200    )
201    .map_err(BoxedError::new)
202    .context(ExternalSnafu)?;
203    let cdt = ConcreteDataType::from(wrapper);
204
205    Ok(cdt)
206}
207
208/// convert `ColumnSchema` lists to it's corresponding proto type
209pub fn column_schemas_to_proto(
210    column_schemas: Vec<ColumnSchema>,
211    primary_keys: &[String],
212) -> Result<Vec<api::v1::ColumnSchema>, Error> {
213    let column_datatypes: Vec<(ColumnDataType, Option<ColumnDataTypeExtension>)> = column_schemas
214        .iter()
215        .map(|c| {
216            ColumnDataTypeWrapper::try_from(c.data_type.clone())
217                .map(|w| w.to_parts())
218                .map_err(BoxedError::new)
219                .context(ExternalSnafu)
220        })
221        .try_collect()?;
222
223    let ret = column_schemas
224        .iter()
225        .zip(column_datatypes)
226        .map(|(schema, datatype)| {
227            let semantic_type = if schema.is_time_index() {
228                SemanticType::Timestamp
229            } else if primary_keys.contains(&schema.name) {
230                SemanticType::Tag
231            } else {
232                SemanticType::Field
233            } as i32;
234
235            api::v1::ColumnSchema {
236                column_name: schema.name.clone(),
237                datatype: datatype.0 as i32,
238                semantic_type,
239                datatype_extension: datatype.1,
240                options: options_from_column_schema(schema),
241            }
242        })
243        .collect();
244    Ok(ret)
245}
246
247/// Convert `RelationDesc` to `ColumnSchema` list,
248/// if the column name is not present, use `col_{idx}` as the column name
249pub fn relation_desc_to_column_schemas_with_fallback(schema: &RelationDesc) -> Vec<ColumnSchema> {
250    schema
251        .typ()
252        .column_types
253        .clone()
254        .into_iter()
255        .enumerate()
256        .map(|(idx, typ)| {
257            let name = schema
258                .names
259                .get(idx)
260                .cloned()
261                .flatten()
262                .unwrap_or(format!("col_{}", idx));
263            let ret = ColumnSchema::new(name, typ.scalar_type, typ.nullable);
264            if schema.typ().time_index == Some(idx) {
265                ret.with_time_index(true)
266            } else {
267                ret
268            }
269        })
270        .collect_vec()
271}
flow/adapter/util.rs

flow/adapter/
util.rs
