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

flow/adapter/
util.rs
