Skip to main content

operator/
expr_helper.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#[cfg(feature = "enterprise")]
16pub mod trigger;
17
18use std::collections::{HashMap, HashSet};
19
20use api::helper::ColumnDataTypeWrapper;
21use api::v1::alter_database_expr::Kind as AlterDatabaseKind;
22use api::v1::alter_table_expr::Kind as AlterTableKind;
23use api::v1::column_def::{options_from_column_schema, try_as_column_schema};
24use api::v1::{
25    AddColumn, AddColumns, AlterDatabaseExpr, AlterTableExpr, Analyzer, ColumnDataType,
26    ColumnDataTypeExtension, CreateFlowExpr, CreateTableExpr, CreateViewExpr, DropColumn,
27    DropColumns, DropDefaults, ExpireAfter, FulltextBackend as PbFulltextBackend, ModifyColumnType,
28    ModifyColumnTypes, RenameTable, SemanticType, SetDatabaseOptions, SetDefaults, SetFulltext,
29    SetIndex, SetIndexes, SetInverted, SetSkipping, SetTableOptions,
30    SkippingIndexType as PbSkippingIndexType, TableName, UnsetDatabaseOptions, UnsetFulltext,
31    UnsetIndex, UnsetIndexes, UnsetInverted, UnsetSkipping, UnsetTableOptions, set_index,
32    unset_index,
33};
34use common_error::ext::BoxedError;
35use common_grpc_expr::util::ColumnExpr;
36use common_time::Timezone;
37use datafusion::sql::planner::object_name_to_table_reference;
38use datatypes::schema::{
39    COLUMN_FULLTEXT_OPT_KEY_ANALYZER, COLUMN_FULLTEXT_OPT_KEY_BACKEND,
40    COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE, COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE,
41    COLUMN_FULLTEXT_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE,
42    COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY, COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE, COMMENT_KEY,
43    ColumnDefaultConstraint, ColumnSchema, FulltextAnalyzer, FulltextBackend, Schema,
44    SkippingIndexType,
45};
46use file_engine::FileOptions;
47use query::sql::{
48    check_file_to_table_schema_compatibility, file_column_schemas_to_table,
49    infer_file_table_schema, prepare_file_table_files,
50};
51use session::context::QueryContextRef;
52use session::table_name::table_idents_to_full_name;
53use snafu::{OptionExt, ResultExt, ensure};
54use sql::ast::{
55    ColumnDef, ColumnOption, ColumnOptionDef, Expr, Ident, ObjectName, ObjectNamePartExt,
56};
57use sql::dialect::GreptimeDbDialect;
58use sql::parser::ParserContext;
59use sql::statements::alter::{
60    AlterDatabase, AlterDatabaseOperation, AlterTable, AlterTableOperation,
61};
62use sql::statements::create::{
63    Column as SqlColumn, ColumnExtensions, CreateExternalTable, CreateFlow, CreateTable,
64    CreateView, TableConstraint,
65};
66use sql::statements::{
67    OptionMap, column_to_schema, concrete_data_type_to_sql_data_type,
68    sql_column_def_to_grpc_column_def, sql_data_type_to_concrete_data_type, value_to_sql_value,
69};
70use sql::util::extract_tables_from_query;
71use store_api::mito_engine_options::{COMPACTION_OVERRIDE, COMPACTION_TYPE};
72use table::requests::{FILE_TABLE_META_KEY, TableOptions};
73use table::table_reference::TableReference;
74#[cfg(feature = "enterprise")]
75pub use trigger::to_create_trigger_task_expr;
76
77use crate::error::{
78    BuildCreateExprOnInsertionSnafu, ColumnDataTypeSnafu, ConvertColumnDefaultConstraintSnafu,
79    ConvertIdentifierSnafu, EncodeJsonSnafu, ExternalSnafu, FindNewColumnsOnInsertionSnafu,
80    IllegalPrimaryKeysDefSnafu, InferFileTableSchemaSnafu, InvalidColumnDefSnafu,
81    InvalidFlowNameSnafu, InvalidSqlSnafu, NotSupportedSnafu, ParseSqlSnafu, ParseSqlValueSnafu,
82    PrepareFileTableSnafu, Result, SchemaIncompatibleSnafu, UnrecognizedTableOptionSnafu,
83};
84
85pub fn create_table_expr_by_column_schemas(
86    table_name: &TableReference<'_>,
87    column_schemas: &[api::v1::ColumnSchema],
88    engine: &str,
89    desc: Option<&str>,
90) -> Result<CreateTableExpr> {
91    let column_exprs = ColumnExpr::from_column_schemas(column_schemas);
92    let expr = common_grpc_expr::util::build_create_table_expr(
93        None,
94        table_name,
95        column_exprs,
96        engine,
97        desc.unwrap_or("Created on insertion"),
98    )
99    .context(BuildCreateExprOnInsertionSnafu)?;
100
101    validate_create_expr(&expr)?;
102    Ok(expr)
103}
104
105pub fn extract_add_columns_expr(
106    schema: &Schema,
107    column_exprs: Vec<ColumnExpr>,
108) -> Result<Option<AddColumns>> {
109    let add_columns = common_grpc_expr::util::extract_new_columns(schema, column_exprs)
110        .context(FindNewColumnsOnInsertionSnafu)?;
111    if let Some(add_columns) = &add_columns {
112        validate_add_columns_expr(add_columns)?;
113    }
114    Ok(add_columns)
115}
116
117//   cpu float64,
118//   memory float64,
119//   TIME INDEX (ts),
120//   PRIMARY KEY(host)
121// ) WITH (location='/var/data/city.csv', format='csv');
122// ```
123// The user needs to specify the TIME INDEX column. If there is no suitable
124// column in the file to use as TIME INDEX, an additional placeholder column
125// needs to be created as the TIME INDEX, and a `DEFAULT <value>` constraint
126// should be added.
127//
128//
129// When the `CREATE EXTERNAL TABLE` statement is in inferred form, like
130// ```sql
131// CREATE EXTERNAL TABLE IF NOT EXISTS city WITH (location='/var/data/city.csv',format='csv');
132// ```
133// 1. If the TIME INDEX column can be inferred from metadata, use that column
134//    as the TIME INDEX. Otherwise,
135// 2. If a column named `greptime_timestamp` exists (with the requirement that
136//    the column is with type TIMESTAMP, otherwise an error is thrown), use
137//    that column as the TIME INDEX. Otherwise,
138// 3. Automatically create the `greptime_timestamp` column and add a `DEFAULT 0`
139//    constraint.
140pub(crate) async fn create_external_expr(
141    create: CreateExternalTable,
142    query_ctx: &QueryContextRef,
143) -> Result<CreateTableExpr> {
144    let (catalog_name, schema_name, table_name) =
145        table_idents_to_full_name(&create.name, query_ctx)
146            .map_err(BoxedError::new)
147            .context(ExternalSnafu)?;
148
149    let mut table_options = create.options.into_map();
150
151    let (object_store, files) = prepare_file_table_files(&table_options)
152        .await
153        .context(PrepareFileTableSnafu)?;
154
155    let file_column_schemas = infer_file_table_schema(&object_store, &files, &table_options)
156        .await
157        .context(InferFileTableSchemaSnafu)?
158        .column_schemas()
159        .to_vec();
160
161    let (time_index, primary_keys, table_column_schemas) = if !create.columns.is_empty() {
162        // expanded form
163        let time_index = find_time_index(&create.constraints)?;
164        let primary_keys = find_primary_keys(&create.columns, &create.constraints)?;
165        let column_schemas =
166            columns_to_column_schemas(&create.columns, &time_index, Some(&query_ctx.timezone()))?;
167        (time_index, primary_keys, column_schemas)
168    } else {
169        // inferred form
170        let (column_schemas, time_index) = file_column_schemas_to_table(&file_column_schemas);
171        let primary_keys = vec![];
172        (time_index, primary_keys, column_schemas)
173    };
174
175    check_file_to_table_schema_compatibility(&file_column_schemas, &table_column_schemas)
176        .context(SchemaIncompatibleSnafu)?;
177
178    let meta = FileOptions {
179        files,
180        file_column_schemas,
181    };
182    table_options.insert(
183        FILE_TABLE_META_KEY.to_string(),
184        serde_json::to_string(&meta).context(EncodeJsonSnafu)?,
185    );
186
187    let column_defs = column_schemas_to_defs(table_column_schemas, &primary_keys)?;
188    let expr = CreateTableExpr {
189        catalog_name,
190        schema_name,
191        table_name,
192        desc: String::default(),
193        column_defs,
194        time_index,
195        primary_keys,
196        create_if_not_exists: create.if_not_exists,
197        table_options,
198        table_id: None,
199        engine: create.engine.clone(),
200    };
201
202    Ok(expr)
203}
204
205/// Convert `CreateTable` statement to [`CreateTableExpr`] gRPC request.
206pub fn create_to_expr(
207    create: &CreateTable,
208    query_ctx: &QueryContextRef,
209) -> Result<CreateTableExpr> {
210    let (catalog_name, schema_name, table_name) =
211        table_idents_to_full_name(&create.name, query_ctx)
212            .map_err(BoxedError::new)
213            .context(ExternalSnafu)?;
214
215    let time_index = find_time_index(&create.constraints)?;
216    let table_options = HashMap::from(
217        &TableOptions::try_from_iter(create.options.to_str_map())
218            .context(UnrecognizedTableOptionSnafu)?,
219    );
220
221    let mut table_options = table_options;
222    if table_options.contains_key(COMPACTION_TYPE) {
223        table_options.insert(COMPACTION_OVERRIDE.to_string(), "true".to_string());
224    }
225
226    let primary_keys = find_primary_keys(&create.columns, &create.constraints)?;
227
228    let expr = CreateTableExpr {
229        catalog_name,
230        schema_name,
231        table_name,
232        desc: String::default(),
233        column_defs: columns_to_expr(
234            &create.columns,
235            &time_index,
236            &primary_keys,
237            Some(&query_ctx.timezone()),
238        )?,
239        time_index,
240        primary_keys,
241        create_if_not_exists: create.if_not_exists,
242        table_options,
243        table_id: None,
244        engine: create.engine.clone(),
245    };
246
247    validate_create_expr(&expr)?;
248    Ok(expr)
249}
250
251/// Convert gRPC's [`CreateTableExpr`] back to `CreateTable` statement.
252/// You can use `create_table_expr_by_column_schemas` to create a `CreateTableExpr` from column schemas.
253///
254/// # Parameters
255///
256/// * `expr` - The `CreateTableExpr` to convert
257/// * `quote_style` - Optional quote style for identifiers (defaults to MySQL style ` backtick)
258pub fn expr_to_create(expr: &CreateTableExpr, quote_style: Option<char>) -> Result<CreateTable> {
259    let quote_style = quote_style.unwrap_or('`');
260
261    // Convert table name
262    let table_name = ObjectName(vec![sql::ast::ObjectNamePart::Identifier(
263        sql::ast::Ident::with_quote(quote_style, &expr.table_name),
264    )]);
265
266    // Convert columns
267    let mut columns = Vec::with_capacity(expr.column_defs.len());
268    for column_def in &expr.column_defs {
269        let column_schema = try_as_column_schema(column_def).context(InvalidColumnDefSnafu {
270            column: &column_def.name,
271        })?;
272
273        let mut options = Vec::new();
274
275        // Add NULL/NOT NULL constraint
276        if column_def.is_nullable {
277            options.push(ColumnOptionDef {
278                name: None,
279                option: ColumnOption::Null,
280            });
281        } else {
282            options.push(ColumnOptionDef {
283                name: None,
284                option: ColumnOption::NotNull,
285            });
286        }
287
288        // Add DEFAULT constraint if present
289        if let Some(default_constraint) = column_schema.default_constraint() {
290            let expr = match default_constraint {
291                ColumnDefaultConstraint::Value(v) => {
292                    Expr::Value(value_to_sql_value(v).context(ParseSqlValueSnafu)?.into())
293                }
294                ColumnDefaultConstraint::Function(func_expr) => {
295                    ParserContext::parse_function(func_expr, &GreptimeDbDialect {})
296                        .context(ParseSqlSnafu)?
297                }
298            };
299            options.push(ColumnOptionDef {
300                name: None,
301                option: ColumnOption::Default(expr),
302            });
303        }
304
305        // Add COMMENT if present
306        if !column_def.comment.is_empty() {
307            options.push(ColumnOptionDef {
308                name: None,
309                option: ColumnOption::Comment(column_def.comment.clone()),
310            });
311        }
312
313        // Note: We don't add inline PRIMARY KEY options here,
314        // we'll handle all primary keys as constraints instead for consistency
315
316        // Handle column extensions (fulltext, inverted index, skipping index)
317        let mut extensions = ColumnExtensions::default();
318
319        // Add fulltext index options if present
320        if let Ok(Some(opt)) = column_schema.fulltext_options()
321            && opt.enable
322        {
323            let mut map = HashMap::from([
324                (
325                    COLUMN_FULLTEXT_OPT_KEY_ANALYZER.to_string(),
326                    opt.analyzer.to_string(),
327                ),
328                (
329                    COLUMN_FULLTEXT_OPT_KEY_CASE_SENSITIVE.to_string(),
330                    opt.case_sensitive.to_string(),
331                ),
332                (
333                    COLUMN_FULLTEXT_OPT_KEY_BACKEND.to_string(),
334                    opt.backend.to_string(),
335                ),
336            ]);
337            if opt.backend == FulltextBackend::Bloom {
338                map.insert(
339                    COLUMN_FULLTEXT_OPT_KEY_GRANULARITY.to_string(),
340                    opt.granularity.to_string(),
341                );
342                map.insert(
343                    COLUMN_FULLTEXT_OPT_KEY_FALSE_POSITIVE_RATE.to_string(),
344                    opt.false_positive_rate().to_string(),
345                );
346            }
347            extensions.fulltext_index_options = Some(map.into());
348        }
349
350        // Add skipping index options if present
351        if let Ok(Some(opt)) = column_schema.skipping_index_options() {
352            let map = HashMap::from([
353                (
354                    COLUMN_SKIPPING_INDEX_OPT_KEY_GRANULARITY.to_string(),
355                    opt.granularity.to_string(),
356                ),
357                (
358                    COLUMN_SKIPPING_INDEX_OPT_KEY_FALSE_POSITIVE_RATE.to_string(),
359                    opt.false_positive_rate().to_string(),
360                ),
361                (
362                    COLUMN_SKIPPING_INDEX_OPT_KEY_TYPE.to_string(),
363                    opt.index_type.to_string(),
364                ),
365            ]);
366            extensions.skipping_index_options = Some(map.into());
367        }
368
369        // Add inverted index options if present
370        if column_schema.is_inverted_indexed() {
371            extensions.inverted_index_options = Some(HashMap::new().into());
372        }
373
374        let sql_column = SqlColumn {
375            column_def: ColumnDef {
376                name: Ident::with_quote(quote_style, &column_def.name),
377                data_type: concrete_data_type_to_sql_data_type(&column_schema.data_type)
378                    .context(ParseSqlSnafu)?,
379                options,
380            },
381            extensions,
382        };
383
384        columns.push(sql_column);
385    }
386
387    // Convert constraints
388    let mut constraints = Vec::new();
389
390    // Add TIME INDEX constraint
391    constraints.push(TableConstraint::TimeIndex {
392        column: Ident::with_quote(quote_style, &expr.time_index),
393    });
394
395    // Add PRIMARY KEY constraint (always add as constraint for consistency)
396    if !expr.primary_keys.is_empty() {
397        let primary_key_columns: Vec<Ident> = expr
398            .primary_keys
399            .iter()
400            .map(|pk| Ident::with_quote(quote_style, pk))
401            .collect();
402
403        constraints.push(TableConstraint::PrimaryKey {
404            columns: primary_key_columns,
405        });
406    }
407
408    // Convert table options
409    let mut options = OptionMap::default();
410    for (key, value) in &expr.table_options {
411        options.insert(key.clone(), value.clone());
412    }
413
414    Ok(CreateTable {
415        if_not_exists: expr.create_if_not_exists,
416        table_id: expr.table_id.as_ref().map(|tid| tid.id).unwrap_or(0),
417        name: table_name,
418        columns,
419        engine: expr.engine.clone(),
420        constraints,
421        options,
422        partitions: None,
423    })
424}
425
426/// Validate the [`CreateTableExpr`] request.
427pub fn validate_create_expr(create: &CreateTableExpr) -> Result<()> {
428    // construct column list
429    let mut column_to_indices = HashMap::with_capacity(create.column_defs.len());
430    for (idx, column) in create.column_defs.iter().enumerate() {
431        if let Some(indices) = column_to_indices.get(&column.name) {
432            return InvalidSqlSnafu {
433                err_msg: format!(
434                    "column name `{}` is duplicated at index {} and {}",
435                    column.name, indices, idx
436                ),
437            }
438            .fail();
439        }
440        column_to_indices.insert(&column.name, idx);
441    }
442
443    // verify time_index exists
444    let _ = column_to_indices
445        .get(&create.time_index)
446        .with_context(|| InvalidSqlSnafu {
447            err_msg: format!(
448                "column name `{}` is not found in column list",
449                create.time_index
450            ),
451        })?;
452
453    // verify primary_key exists
454    for pk in &create.primary_keys {
455        let _ = column_to_indices
456            .get(&pk)
457            .with_context(|| InvalidSqlSnafu {
458                err_msg: format!("column name `{}` is not found in column list", pk),
459            })?;
460    }
461
462    // construct primary_key set
463    let mut pk_set = HashSet::new();
464    for pk in &create.primary_keys {
465        if !pk_set.insert(pk) {
466            return InvalidSqlSnafu {
467                err_msg: format!("column name `{}` is duplicated in primary keys", pk),
468            }
469            .fail();
470        }
471    }
472
473    // verify time index is not primary key
474    if pk_set.contains(&create.time_index) {
475        return InvalidSqlSnafu {
476            err_msg: format!(
477                "column name `{}` is both primary key and time index",
478                create.time_index
479            ),
480        }
481        .fail();
482    }
483
484    for column in &create.column_defs {
485        // verify do not contain interval type column issue #3235
486        if is_interval_type(&column.data_type()) {
487            return InvalidSqlSnafu {
488                err_msg: format!(
489                    "column name `{}` is interval type, which is not supported",
490                    column.name
491                ),
492            }
493            .fail();
494        }
495        // verify do not contain datetime type column issue #5489
496        if is_date_time_type(&column.data_type()) {
497            return InvalidSqlSnafu {
498                err_msg: format!(
499                    "column name `{}` is datetime type, which is not supported, please use `timestamp` type instead",
500                    column.name
501                ),
502            }
503            .fail();
504        }
505    }
506    Ok(())
507}
508
509fn validate_add_columns_expr(add_columns: &AddColumns) -> Result<()> {
510    for add_column in &add_columns.add_columns {
511        let Some(column_def) = &add_column.column_def else {
512            continue;
513        };
514        if is_date_time_type(&column_def.data_type()) {
515            return InvalidSqlSnafu {
516                    err_msg: format!("column name `{}` is datetime type, which is not supported, please use `timestamp` type instead", column_def.name),
517                }
518                .fail();
519        }
520        if is_interval_type(&column_def.data_type()) {
521            return InvalidSqlSnafu {
522                err_msg: format!(
523                    "column name `{}` is interval type, which is not supported",
524                    column_def.name
525                ),
526            }
527            .fail();
528        }
529    }
530    Ok(())
531}
532
533fn is_date_time_type(data_type: &ColumnDataType) -> bool {
534    matches!(data_type, ColumnDataType::Datetime)
535}
536
537fn is_interval_type(data_type: &ColumnDataType) -> bool {
538    matches!(
539        data_type,
540        ColumnDataType::IntervalYearMonth
541            | ColumnDataType::IntervalDayTime
542            | ColumnDataType::IntervalMonthDayNano
543    )
544}
545
546fn find_primary_keys(
547    columns: &[SqlColumn],
548    constraints: &[TableConstraint],
549) -> Result<Vec<String>> {
550    let columns_pk = columns
551        .iter()
552        .filter_map(|x| {
553            if x.options()
554                .iter()
555                .any(|o| matches!(o.option, ColumnOption::PrimaryKey(_)))
556            {
557                Some(x.name().value.clone())
558            } else {
559                None
560            }
561        })
562        .collect::<Vec<String>>();
563
564    ensure!(
565        columns_pk.len() <= 1,
566        IllegalPrimaryKeysDefSnafu {
567            msg: "not allowed to inline multiple primary keys in columns options"
568        }
569    );
570
571    let constraints_pk = constraints
572        .iter()
573        .filter_map(|constraint| match constraint {
574            TableConstraint::PrimaryKey { columns, .. } => {
575                Some(columns.iter().map(|ident| ident.value.clone()))
576            }
577            _ => None,
578        })
579        .flatten()
580        .collect::<Vec<String>>();
581
582    ensure!(
583        columns_pk.is_empty() || constraints_pk.is_empty(),
584        IllegalPrimaryKeysDefSnafu {
585            msg: "found definitions of primary keys in multiple places"
586        }
587    );
588
589    let mut primary_keys = Vec::with_capacity(columns_pk.len() + constraints_pk.len());
590    primary_keys.extend(columns_pk);
591    primary_keys.extend(constraints_pk);
592    Ok(primary_keys)
593}
594
595pub fn find_time_index(constraints: &[TableConstraint]) -> Result<String> {
596    let time_index = constraints
597        .iter()
598        .filter_map(|constraint| match constraint {
599            TableConstraint::TimeIndex { column, .. } => Some(&column.value),
600            _ => None,
601        })
602        .collect::<Vec<&String>>();
603    ensure!(
604        time_index.len() == 1,
605        InvalidSqlSnafu {
606            err_msg: "must have one and only one TimeIndex columns",
607        }
608    );
609    Ok(time_index[0].clone())
610}
611
612fn columns_to_expr(
613    column_defs: &[SqlColumn],
614    time_index: &str,
615    primary_keys: &[String],
616    timezone: Option<&Timezone>,
617) -> Result<Vec<api::v1::ColumnDef>> {
618    let column_schemas = columns_to_column_schemas(column_defs, time_index, timezone)?;
619    column_schemas_to_defs(column_schemas, primary_keys)
620}
621
622fn columns_to_column_schemas(
623    columns: &[SqlColumn],
624    time_index: &str,
625    timezone: Option<&Timezone>,
626) -> Result<Vec<ColumnSchema>> {
627    columns
628        .iter()
629        .map(|c| column_to_schema(c, time_index, timezone).context(ParseSqlSnafu))
630        .collect::<Result<Vec<ColumnSchema>>>()
631}
632
633// TODO(weny): refactor this function to use `try_as_column_def`
634pub fn column_schemas_to_defs(
635    column_schemas: Vec<ColumnSchema>,
636    primary_keys: &[String],
637) -> Result<Vec<api::v1::ColumnDef>> {
638    let column_datatypes: Vec<(ColumnDataType, Option<ColumnDataTypeExtension>)> = column_schemas
639        .iter()
640        .map(|c| {
641            ColumnDataTypeWrapper::try_from(c.data_type.clone())
642                .map(|w| w.to_parts())
643                .context(ColumnDataTypeSnafu)
644        })
645        .collect::<Result<Vec<_>>>()?;
646
647    column_schemas
648        .iter()
649        .zip(column_datatypes)
650        .map(|(schema, datatype)| {
651            let semantic_type = if schema.is_time_index() {
652                SemanticType::Timestamp
653            } else if primary_keys.contains(&schema.name) {
654                SemanticType::Tag
655            } else {
656                SemanticType::Field
657            } as i32;
658            let comment = schema
659                .metadata()
660                .get(COMMENT_KEY)
661                .cloned()
662                .unwrap_or_default();
663
664            Ok(api::v1::ColumnDef {
665                name: schema.name.clone(),
666                data_type: datatype.0 as i32,
667                is_nullable: schema.is_nullable(),
668                default_constraint: match schema.default_constraint() {
669                    None => vec![],
670                    Some(v) => {
671                        v.clone()
672                            .try_into()
673                            .context(ConvertColumnDefaultConstraintSnafu {
674                                column_name: &schema.name,
675                            })?
676                    }
677                },
678                semantic_type,
679                comment,
680                datatype_extension: datatype.1,
681                options: options_from_column_schema(schema),
682            })
683        })
684        .collect()
685}
686
687#[derive(Debug, Clone, PartialEq, Eq)]
688pub struct RepartitionRequest {
689    pub catalog_name: String,
690    pub schema_name: String,
691    pub table_name: String,
692    pub source: RepartitionSource,
693    pub into_exprs: Vec<Expr>,
694    pub options: OptionMap,
695}
696
697#[derive(Debug, Clone, PartialEq, Eq)]
698pub enum RepartitionSource {
699    Partitions {
700        from_exprs: Vec<Expr>,
701        target_partition_columns: Option<Vec<String>>,
702    },
703    Unpartitioned {
704        partition_columns: Vec<String>,
705    },
706}
707
708pub(crate) fn to_repartition_request(
709    alter_table: AlterTable,
710    query_ctx: &QueryContextRef,
711) -> Result<RepartitionRequest> {
712    let AlterTable {
713        table_name,
714        alter_operation,
715        options,
716    } = alter_table;
717
718    let (catalog_name, schema_name, table_name) = table_idents_to_full_name(&table_name, query_ctx)
719        .map_err(BoxedError::new)
720        .context(ExternalSnafu)?;
721
722    let (source, into_exprs) = match alter_operation {
723        AlterTableOperation::Repartition { operation } => (
724            RepartitionSource::Partitions {
725                from_exprs: operation.from_exprs,
726                target_partition_columns: operation.partition_columns.map(|columns| {
727                    columns
728                        .into_iter()
729                        .map(|ident| ident.value)
730                        .collect::<Vec<_>>()
731                }),
732            },
733            operation.into_exprs,
734        ),
735        AlterTableOperation::Partition { partitions } => (
736            RepartitionSource::Unpartitioned {
737                partition_columns: partitions
738                    .column_list
739                    .into_iter()
740                    .map(|ident| ident.value)
741                    .collect(),
742            },
743            partitions.exprs,
744        ),
745        _ => {
746            return InvalidSqlSnafu {
747                err_msg: "expected REPARTITION or PARTITION operation",
748            }
749            .fail();
750        }
751    };
752
753    Ok(RepartitionRequest {
754        catalog_name,
755        schema_name,
756        table_name,
757        source,
758        into_exprs,
759        options,
760    })
761}
762
763/// Converts a SQL alter table statement into a gRPC alter table expression.
764pub(crate) fn to_alter_table_expr(
765    alter_table: AlterTable,
766    query_ctx: &QueryContextRef,
767) -> Result<AlterTableExpr> {
768    let (catalog_name, schema_name, table_name) =
769        table_idents_to_full_name(alter_table.table_name(), query_ctx)
770            .map_err(BoxedError::new)
771            .context(ExternalSnafu)?;
772
773    let kind = match alter_table.alter_operation {
774        AlterTableOperation::AddConstraint(_) => {
775            return NotSupportedSnafu {
776                feat: "ADD CONSTRAINT",
777            }
778            .fail();
779        }
780        AlterTableOperation::AddColumns { add_columns } => AlterTableKind::AddColumns(AddColumns {
781            add_columns: add_columns
782                .into_iter()
783                .map(|add_column| {
784                    let column_def = sql_column_def_to_grpc_column_def(
785                        &add_column.column_def,
786                        Some(&query_ctx.timezone()),
787                    )
788                    .map_err(BoxedError::new)
789                    .context(ExternalSnafu)?;
790                    if is_interval_type(&column_def.data_type()) {
791                        return NotSupportedSnafu {
792                            feat: "Add column with interval type",
793                        }
794                        .fail();
795                    }
796                    Ok(AddColumn {
797                        column_def: Some(column_def),
798                        location: add_column.location.as_ref().map(From::from),
799                        add_if_not_exists: add_column.add_if_not_exists,
800                    })
801                })
802                .collect::<Result<Vec<AddColumn>>>()?,
803        }),
804        AlterTableOperation::ModifyColumnType {
805            column_name,
806            target_type,
807        } => {
808            let target_type =
809                sql_data_type_to_concrete_data_type(&target_type).context(ParseSqlSnafu)?;
810            let (target_type, target_type_extension) = ColumnDataTypeWrapper::try_from(target_type)
811                .map(|w| w.to_parts())
812                .context(ColumnDataTypeSnafu)?;
813            if is_interval_type(&target_type) {
814                return NotSupportedSnafu {
815                    feat: "Modify column type to interval type",
816                }
817                .fail();
818            }
819            AlterTableKind::ModifyColumnTypes(ModifyColumnTypes {
820                modify_column_types: vec![ModifyColumnType {
821                    column_name: column_name.value,
822                    target_type: target_type as i32,
823                    target_type_extension,
824                }],
825            })
826        }
827        AlterTableOperation::DropColumn { name } => AlterTableKind::DropColumns(DropColumns {
828            drop_columns: vec![DropColumn {
829                name: name.value.clone(),
830            }],
831        }),
832        AlterTableOperation::RenameTable { new_table_name } => {
833            AlterTableKind::RenameTable(RenameTable {
834                new_table_name: new_table_name.clone(),
835            })
836        }
837        AlterTableOperation::SetTableOptions { options } => {
838            AlterTableKind::SetTableOptions(SetTableOptions {
839                table_options: options.into_iter().map(Into::into).collect(),
840            })
841        }
842        AlterTableOperation::UnsetTableOptions { keys } => {
843            AlterTableKind::UnsetTableOptions(UnsetTableOptions { keys })
844        }
845        AlterTableOperation::Repartition { .. } => {
846            return NotSupportedSnafu {
847                feat: "ALTER TABLE ... REPARTITION",
848            }
849            .fail();
850        }
851        AlterTableOperation::Partition { .. } => {
852            return NotSupportedSnafu {
853                feat: "ALTER TABLE ... PARTITION ON COLUMNS",
854            }
855            .fail();
856        }
857        AlterTableOperation::SetIndex { options } => {
858            let option = match options {
859                sql::statements::alter::SetIndexOperation::Fulltext {
860                    column_name,
861                    options,
862                } => SetIndex {
863                    options: Some(set_index::Options::Fulltext(SetFulltext {
864                        column_name: column_name.value,
865                        enable: options.enable,
866                        analyzer: match options.analyzer {
867                            FulltextAnalyzer::English => Analyzer::English.into(),
868                            FulltextAnalyzer::Chinese => Analyzer::Chinese.into(),
869                        },
870                        case_sensitive: options.case_sensitive,
871                        backend: match options.backend {
872                            FulltextBackend::Bloom => PbFulltextBackend::Bloom.into(),
873                            FulltextBackend::Tantivy => PbFulltextBackend::Tantivy.into(),
874                        },
875                        granularity: options.granularity as u64,
876                        false_positive_rate: options.false_positive_rate(),
877                    })),
878                },
879                sql::statements::alter::SetIndexOperation::Inverted { column_name } => SetIndex {
880                    options: Some(set_index::Options::Inverted(SetInverted {
881                        column_name: column_name.value,
882                    })),
883                },
884                sql::statements::alter::SetIndexOperation::Skipping {
885                    column_name,
886                    options,
887                } => SetIndex {
888                    options: Some(set_index::Options::Skipping(SetSkipping {
889                        column_name: column_name.value,
890                        enable: true,
891                        granularity: options.granularity as u64,
892                        false_positive_rate: options.false_positive_rate(),
893                        skipping_index_type: match options.index_type {
894                            SkippingIndexType::BloomFilter => {
895                                PbSkippingIndexType::BloomFilter.into()
896                            }
897                        },
898                    })),
899                },
900            };
901            AlterTableKind::SetIndexes(SetIndexes {
902                set_indexes: vec![option],
903            })
904        }
905        AlterTableOperation::UnsetIndex { options } => {
906            let option = match options {
907                sql::statements::alter::UnsetIndexOperation::Fulltext { column_name } => {
908                    UnsetIndex {
909                        options: Some(unset_index::Options::Fulltext(UnsetFulltext {
910                            column_name: column_name.value,
911                        })),
912                    }
913                }
914                sql::statements::alter::UnsetIndexOperation::Inverted { column_name } => {
915                    UnsetIndex {
916                        options: Some(unset_index::Options::Inverted(UnsetInverted {
917                            column_name: column_name.value,
918                        })),
919                    }
920                }
921                sql::statements::alter::UnsetIndexOperation::Skipping { column_name } => {
922                    UnsetIndex {
923                        options: Some(unset_index::Options::Skipping(UnsetSkipping {
924                            column_name: column_name.value,
925                        })),
926                    }
927                }
928            };
929
930            AlterTableKind::UnsetIndexes(UnsetIndexes {
931                unset_indexes: vec![option],
932            })
933        }
934        AlterTableOperation::DropDefaults { columns } => {
935            AlterTableKind::DropDefaults(DropDefaults {
936                drop_defaults: columns
937                    .into_iter()
938                    .map(|col| {
939                        let column_name = col.0.to_string();
940                        Ok(api::v1::DropDefault { column_name })
941                    })
942                    .collect::<Result<Vec<_>>>()?,
943            })
944        }
945        AlterTableOperation::SetDefaults { defaults } => AlterTableKind::SetDefaults(SetDefaults {
946            set_defaults: defaults
947                .into_iter()
948                .map(|col| {
949                    let column_name = col.column_name.to_string();
950                    let default_constraint = serde_json::to_string(&col.default_constraint)
951                        .context(EncodeJsonSnafu)?
952                        .into_bytes();
953                    Ok(api::v1::SetDefault {
954                        column_name,
955                        default_constraint,
956                    })
957                })
958                .collect::<Result<Vec<_>>>()?,
959        }),
960    };
961
962    Ok(AlterTableExpr {
963        catalog_name,
964        schema_name,
965        table_name,
966        kind: Some(kind),
967    })
968}
969
970/// Try to cast the `[AlterDatabase]` statement into gRPC `[AlterDatabaseExpr]`.
971pub fn to_alter_database_expr(
972    alter_database: AlterDatabase,
973    query_ctx: &QueryContextRef,
974) -> Result<AlterDatabaseExpr> {
975    let catalog = query_ctx.current_catalog();
976    let schema = alter_database.database_name;
977
978    let kind = match alter_database.alter_operation {
979        AlterDatabaseOperation::SetDatabaseOption { options } => {
980            let options = options.into_iter().map(Into::into).collect();
981            AlterDatabaseKind::SetDatabaseOptions(SetDatabaseOptions {
982                set_database_options: options,
983            })
984        }
985        AlterDatabaseOperation::UnsetDatabaseOption { keys } => {
986            AlterDatabaseKind::UnsetDatabaseOptions(UnsetDatabaseOptions { keys })
987        }
988    };
989
990    Ok(AlterDatabaseExpr {
991        catalog_name: catalog.to_string(),
992        schema_name: schema.to_string(),
993        kind: Some(kind),
994    })
995}
996
997/// Try to cast the `[CreateViewExpr]` statement into gRPC `[CreateViewExpr]`.
998pub fn to_create_view_expr(
999    stmt: CreateView,
1000    logical_plan: Vec<u8>,
1001    table_names: Vec<TableName>,
1002    columns: Vec<String>,
1003    plan_columns: Vec<String>,
1004    definition: String,
1005    query_ctx: QueryContextRef,
1006) -> Result<CreateViewExpr> {
1007    let (catalog_name, schema_name, view_name) = table_idents_to_full_name(&stmt.name, &query_ctx)
1008        .map_err(BoxedError::new)
1009        .context(ExternalSnafu)?;
1010
1011    let expr = CreateViewExpr {
1012        catalog_name,
1013        schema_name,
1014        view_name,
1015        logical_plan,
1016        create_if_not_exists: stmt.if_not_exists,
1017        or_replace: stmt.or_replace,
1018        table_names,
1019        columns,
1020        plan_columns,
1021        definition,
1022    };
1023
1024    Ok(expr)
1025}
1026
1027pub fn to_create_flow_task_expr(
1028    create_flow: CreateFlow,
1029    query_ctx: &QueryContextRef,
1030) -> Result<CreateFlowExpr> {
1031    // retrieve sink table name
1032    let sink_table_ref = object_name_to_table_reference(create_flow.sink_table_name.clone(), true)
1033        .with_context(|_| ConvertIdentifierSnafu {
1034            ident: create_flow.sink_table_name.to_string(),
1035        })?;
1036    let catalog = sink_table_ref
1037        .catalog()
1038        .unwrap_or(query_ctx.current_catalog())
1039        .to_string();
1040    let schema = sink_table_ref
1041        .schema()
1042        .map(|s| s.to_owned())
1043        .unwrap_or(query_ctx.current_schema());
1044
1045    let sink_table_name = TableName {
1046        catalog_name: catalog,
1047        schema_name: schema,
1048        table_name: sink_table_ref.table().to_string(),
1049    };
1050
1051    let source_table_names = extract_tables_from_query(&create_flow.query)
1052        .map(|name| {
1053            let reference =
1054                object_name_to_table_reference(name.clone(), true).with_context(|_| {
1055                    ConvertIdentifierSnafu {
1056                        ident: name.to_string(),
1057                    }
1058                })?;
1059            let catalog = reference
1060                .catalog()
1061                .unwrap_or(query_ctx.current_catalog())
1062                .to_string();
1063            let schema = reference
1064                .schema()
1065                .map(|s| s.to_string())
1066                .unwrap_or(query_ctx.current_schema());
1067
1068            let table_name = TableName {
1069                catalog_name: catalog,
1070                schema_name: schema,
1071                table_name: reference.table().to_string(),
1072            };
1073            Ok(table_name)
1074        })
1075        .collect::<Result<Vec<_>>>()?;
1076
1077    let eval_interval = create_flow.eval_interval;
1078
1079    Ok(CreateFlowExpr {
1080        catalog_name: query_ctx.current_catalog().to_string(),
1081        flow_name: sanitize_flow_name(create_flow.flow_name)?,
1082        source_table_names,
1083        sink_table_name: Some(sink_table_name),
1084        or_replace: create_flow.or_replace,
1085        create_if_not_exists: create_flow.if_not_exists,
1086        expire_after: create_flow.expire_after.map(|value| ExpireAfter { value }),
1087        eval_interval: eval_interval.map(|seconds| api::v1::EvalInterval { seconds }),
1088        comment: create_flow.comment.unwrap_or_default(),
1089        sql: create_flow.query.to_string(),
1090        flow_options: stringify_flow_options(create_flow.flow_options)?,
1091    })
1092}
1093
1094fn stringify_flow_options(flow_options: OptionMap) -> Result<HashMap<String, String>> {
1095    let options_len = flow_options.len();
1096    let flow_options = flow_options.into_map();
1097    ensure!(
1098        flow_options.len() == options_len,
1099        InvalidSqlSnafu {
1100            err_msg: "flow options only support scalar string-compatible values".to_string(),
1101        }
1102    );
1103    Ok(flow_options)
1104}
1105
1106/// sanitize the flow name, remove possible quotes
1107fn sanitize_flow_name(mut flow_name: ObjectName) -> Result<String> {
1108    ensure!(
1109        flow_name.0.len() == 1,
1110        InvalidFlowNameSnafu {
1111            name: flow_name.to_string(),
1112        }
1113    );
1114    // safety: we've checked flow_name.0 has exactly one element.
1115    Ok(flow_name.0.swap_remove(0).to_string_unquoted())
1116}
1117
1118#[cfg(test)]
1119mod tests {
1120    use std::collections::HashMap;
1121
1122    use api::v1::{SetDatabaseOptions, UnsetDatabaseOptions};
1123    use datatypes::value::Value;
1124    use session::context::{QueryContext, QueryContextBuilder};
1125    use sql::dialect::GreptimeDbDialect;
1126    use sql::parser::{ParseOptions, ParserContext};
1127    use sql::statements::statement::Statement;
1128    use store_api::storage::ColumnDefaultConstraint;
1129
1130    use super::*;
1131
1132    #[test]
1133    fn test_create_flow_tql_expr() {
1134        let sql = r#"
1135CREATE FLOW calc_reqs SINK TO cnt_reqs AS
1136TQL EVAL (0, 15, '5s') count_values("status_code", http_requests);"#;
1137        let stmt =
1138            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default());
1139
1140        assert!(
1141            stmt.is_err(),
1142            "Expected error for invalid TQL EVAL parameters: {:#?}",
1143            stmt
1144        );
1145
1146        let sql = r#"
1147CREATE FLOW calc_reqs SINK TO cnt_reqs AS
1148TQL EVAL (now() - '15s'::interval, now(), '5s') count_values("status_code", http_requests);"#;
1149        let stmt =
1150            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1151                .unwrap()
1152                .pop()
1153                .unwrap();
1154
1155        let Statement::CreateFlow(create_flow) = stmt else {
1156            unreachable!()
1157        };
1158        let expr = to_create_flow_task_expr(create_flow, &QueryContext::arc()).unwrap();
1159
1160        let to_dot_sep =
1161            |c: TableName| format!("{}.{}.{}", c.catalog_name, c.schema_name, c.table_name);
1162        assert_eq!("calc_reqs", expr.flow_name);
1163        assert_eq!("greptime", expr.catalog_name);
1164        assert_eq!(
1165            "greptime.public.cnt_reqs",
1166            expr.sink_table_name.map(to_dot_sep).unwrap()
1167        );
1168        assert_eq!(1, expr.source_table_names.len());
1169        assert_eq!(
1170            "greptime.public.http_requests",
1171            to_dot_sep(expr.source_table_names[0].clone())
1172        );
1173        assert_eq!(
1174            r#"TQL EVAL (now() - '15s'::interval, now(), '5s') count_values("status_code", http_requests)"#,
1175            expr.sql
1176        );
1177
1178        let sql = r#"
1179CREATE FLOW calc_reqs SINK TO cnt_reqs AS
1180TQL EVAL (now() - '15s'::interval, now(), '5s') count_values("status_code", http_requests{__schema__="greptime_private"});"#;
1181        let stmt =
1182            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1183                .unwrap()
1184                .pop()
1185                .unwrap();
1186        let Statement::CreateFlow(create_flow) = stmt else {
1187            unreachable!()
1188        };
1189        let expr = to_create_flow_task_expr(create_flow, &QueryContext::arc()).unwrap();
1190        assert_eq!(1, expr.source_table_names.len());
1191        assert_eq!(
1192            "greptime.greptime_private.http_requests",
1193            to_dot_sep(expr.source_table_names[0].clone())
1194        );
1195
1196        let sql = r#"
1197CREATE FLOW calc_reqs SINK TO cnt_reqs AS
1198TQL EVAL (now() - '15s'::interval, now(), '5s') count_values("status_code", http_requests{__database__="greptime_private"});"#;
1199        let stmt =
1200            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1201                .unwrap()
1202                .pop()
1203                .unwrap();
1204        let Statement::CreateFlow(create_flow) = stmt else {
1205            unreachable!()
1206        };
1207        let expr = to_create_flow_task_expr(create_flow, &QueryContext::arc()).unwrap();
1208        assert_eq!(1, expr.source_table_names.len());
1209        assert_eq!(
1210            "greptime.greptime_private.http_requests",
1211            to_dot_sep(expr.source_table_names[0].clone())
1212        );
1213    }
1214
1215    #[test]
1216    fn test_create_flow_tql_cte_source_tables() {
1217        let sql = r#"
1218CREATE FLOW calc_cte
1219SINK TO metric_cte_sink
1220EVAL INTERVAL '1m'
1221AS
1222WITH tql(ts, the_value) AS (
1223  TQL EVAL (now() - '1m'::interval, now(), '5s') metric_cte
1224)
1225SELECT * FROM tql;
1226"#;
1227
1228        let stmt =
1229            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1230                .unwrap()
1231                .pop()
1232                .unwrap();
1233
1234        let Statement::CreateFlow(create_flow) = stmt else {
1235            unreachable!()
1236        };
1237        let expr = to_create_flow_task_expr(create_flow, &QueryContext::arc()).unwrap();
1238
1239        let to_dot_sep =
1240            |c: TableName| format!("{}.{}.{}", c.catalog_name, c.schema_name, c.table_name);
1241        assert_eq!(1, expr.source_table_names.len());
1242        assert_eq!(
1243            "greptime.public.metric_cte",
1244            to_dot_sep(expr.source_table_names[0].clone())
1245        );
1246    }
1247
1248    #[test]
1249    fn test_create_flow_tql_cte_source_tables_quoted_cte_name() {
1250        let sql = r#"
1251CREATE FLOW calc_cte
1252SINK TO metric_cte_sink
1253EVAL INTERVAL '1m'
1254AS
1255WITH "TQL"(ts, the_value) AS (
1256  TQL EVAL (now() - '1m'::interval, now(), '5s') metric_cte
1257)
1258SELECT * FROM "TQL";
1259"#;
1260
1261        let stmt =
1262            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1263                .unwrap()
1264                .pop()
1265                .unwrap();
1266
1267        let Statement::CreateFlow(create_flow) = stmt else {
1268            unreachable!()
1269        };
1270        let expr = to_create_flow_task_expr(create_flow, &QueryContext::arc()).unwrap();
1271
1272        let to_dot_sep =
1273            |c: TableName| format!("{}.{}.{}", c.catalog_name, c.schema_name, c.table_name);
1274        assert_eq!(1, expr.source_table_names.len());
1275        assert_eq!(
1276            "greptime.public.metric_cte",
1277            to_dot_sep(expr.source_table_names[0].clone())
1278        );
1279    }
1280
1281    #[test]
1282    fn test_create_flow_tql_cte_source_tables_same_name() {
1283        let sql = r#"
1284CREATE FLOW calc_cte
1285SINK TO metric_cte_sink
1286EVAL INTERVAL '1m'
1287AS
1288WITH tql(ts, the_value) AS (
1289  TQL EVAL (now() - '1m'::interval, now(), '5s') tql
1290)
1291SELECT * FROM tql;
1292"#;
1293
1294        let stmt =
1295            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1296                .unwrap()
1297                .pop()
1298                .unwrap();
1299
1300        let Statement::CreateFlow(create_flow) = stmt else {
1301            unreachable!()
1302        };
1303        let expr = to_create_flow_task_expr(create_flow, &QueryContext::arc()).unwrap();
1304
1305        let to_dot_sep =
1306            |c: TableName| format!("{}.{}.{}", c.catalog_name, c.schema_name, c.table_name);
1307        assert_eq!(1, expr.source_table_names.len());
1308        assert_eq!(
1309            "greptime.public.tql",
1310            to_dot_sep(expr.source_table_names[0].clone())
1311        );
1312    }
1313
1314    #[test]
1315    fn test_create_flow_expr() {
1316        let sql = r"
1317CREATE FLOW test_distinct_basic SINK TO out_distinct_basic AS
1318SELECT
1319    DISTINCT number as dis
1320FROM
1321    distinct_basic;";
1322        let stmt =
1323            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1324                .unwrap()
1325                .pop()
1326                .unwrap();
1327
1328        let Statement::CreateFlow(create_flow) = stmt else {
1329            unreachable!()
1330        };
1331        let expr = to_create_flow_task_expr(create_flow, &QueryContext::arc()).unwrap();
1332
1333        let to_dot_sep =
1334            |c: TableName| format!("{}.{}.{}", c.catalog_name, c.schema_name, c.table_name);
1335        assert_eq!("test_distinct_basic", expr.flow_name);
1336        assert_eq!("greptime", expr.catalog_name);
1337        assert_eq!(
1338            "greptime.public.out_distinct_basic",
1339            expr.sink_table_name.map(to_dot_sep).unwrap()
1340        );
1341        assert_eq!(1, expr.source_table_names.len());
1342        assert_eq!(
1343            "greptime.public.distinct_basic",
1344            to_dot_sep(expr.source_table_names[0].clone())
1345        );
1346        assert_eq!(
1347            r"SELECT
1348    DISTINCT number as dis
1349FROM
1350    distinct_basic",
1351            expr.sql
1352        );
1353
1354        let sql = r"
1355CREATE FLOW `task_2`
1356SINK TO schema_1.table_1
1357AS
1358SELECT max(c1), min(c2) FROM schema_2.table_2;";
1359        let stmt =
1360            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1361                .unwrap()
1362                .pop()
1363                .unwrap();
1364
1365        let Statement::CreateFlow(create_flow) = stmt else {
1366            unreachable!()
1367        };
1368        let expr = to_create_flow_task_expr(create_flow, &QueryContext::arc()).unwrap();
1369
1370        let to_dot_sep =
1371            |c: TableName| format!("{}.{}.{}", c.catalog_name, c.schema_name, c.table_name);
1372        assert_eq!("task_2", expr.flow_name);
1373        assert_eq!("greptime", expr.catalog_name);
1374        assert_eq!(
1375            "greptime.schema_1.table_1",
1376            expr.sink_table_name.map(to_dot_sep).unwrap()
1377        );
1378        assert_eq!(1, expr.source_table_names.len());
1379        assert_eq!(
1380            "greptime.schema_2.table_2",
1381            to_dot_sep(expr.source_table_names[0].clone())
1382        );
1383        assert_eq!("SELECT max(c1), min(c2) FROM schema_2.table_2", expr.sql);
1384        assert!(expr.flow_options.is_empty());
1385
1386        let sql = r"
1387CREATE FLOW task_3
1388SINK TO schema_1.table_1
1389WITH (defer_on_missing_source = 'true', foo = 'bar')
1390AS
1391SELECT max(c1), min(c2) FROM schema_2.table_2;";
1392        let stmt =
1393            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1394                .unwrap()
1395                .pop()
1396                .unwrap();
1397
1398        let Statement::CreateFlow(create_flow) = stmt else {
1399            unreachable!()
1400        };
1401        let expr = to_create_flow_task_expr(create_flow, &QueryContext::arc()).unwrap();
1402        assert_eq!(
1403            expr.flow_options,
1404            HashMap::from([
1405                ("defer_on_missing_source".to_string(), "true".to_string()),
1406                ("foo".to_string(), "bar".to_string()),
1407            ])
1408        );
1409
1410        let sql = r"
1411CREATE FLOW task_4
1412SINK TO schema_1.table_1
1413WITH (defer_on_missing_source = true)
1414AS
1415SELECT max(c1), min(c2) FROM schema_2.table_2;";
1416        let stmt =
1417            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1418                .unwrap()
1419                .pop()
1420                .unwrap();
1421
1422        let Statement::CreateFlow(create_flow) = stmt else {
1423            unreachable!()
1424        };
1425        let expr = to_create_flow_task_expr(create_flow, &QueryContext::arc()).unwrap();
1426        assert_eq!(
1427            expr.flow_options,
1428            HashMap::from([("defer_on_missing_source".to_string(), "true".to_string(),)])
1429        );
1430
1431        let sql = r"
1432CREATE FLOW task_5
1433SINK TO schema_1.table_1
1434WITH (defer_on_missing_source = [true])
1435AS
1436SELECT max(c1), min(c2) FROM schema_2.table_2;";
1437        let stmt =
1438            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1439                .unwrap()
1440                .pop()
1441                .unwrap();
1442
1443        let Statement::CreateFlow(create_flow) = stmt else {
1444            unreachable!()
1445        };
1446        let res = to_create_flow_task_expr(create_flow, &QueryContext::arc());
1447        assert!(res.is_err());
1448        assert!(
1449            res.unwrap_err()
1450                .to_string()
1451                .contains("flow options only support scalar string-compatible values")
1452        );
1453
1454        let sql = r"
1455CREATE FLOW abc.`task_2`
1456SINK TO schema_1.table_1
1457AS
1458SELECT max(c1), min(c2) FROM schema_2.table_2;";
1459        let stmt =
1460            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1461                .unwrap()
1462                .pop()
1463                .unwrap();
1464
1465        let Statement::CreateFlow(create_flow) = stmt else {
1466            unreachable!()
1467        };
1468        let res = to_create_flow_task_expr(create_flow, &QueryContext::arc());
1469
1470        assert!(res.is_err());
1471        assert!(
1472            res.unwrap_err()
1473                .to_string()
1474                .contains("Invalid flow name: abc.`task_2`")
1475        );
1476    }
1477
1478    #[test]
1479    fn test_create_to_expr() {
1480        let sql = "CREATE TABLE monitor (host STRING,ts TIMESTAMP,TIME INDEX (ts),PRIMARY KEY(host)) ENGINE=mito WITH(ttl='3days', write_buffer_size='1024KB');";
1481        let stmt =
1482            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1483                .unwrap()
1484                .pop()
1485                .unwrap();
1486
1487        let Statement::CreateTable(create_table) = stmt else {
1488            unreachable!()
1489        };
1490        let expr = create_to_expr(&create_table, &QueryContext::arc()).unwrap();
1491        assert_eq!("3days", expr.table_options.get("ttl").unwrap());
1492        assert_eq!(
1493            "1.0MiB",
1494            expr.table_options.get("write_buffer_size").unwrap()
1495        );
1496    }
1497
1498    #[test]
1499    fn test_invalid_create_to_expr() {
1500        let cases = [
1501            // duplicate column declaration
1502            "CREATE TABLE monitor (host STRING primary key, ts TIMESTAMP TIME INDEX, some_column text, some_column string);",
1503            // duplicate primary key
1504            "CREATE TABLE monitor (host STRING, ts TIMESTAMP TIME INDEX, some_column STRING, PRIMARY KEY (some_column, host, some_column));",
1505            // time index is primary key
1506            "CREATE TABLE monitor (host STRING, ts TIMESTAMP TIME INDEX, PRIMARY KEY (host, ts));",
1507        ];
1508
1509        for sql in cases {
1510            let stmt = ParserContext::create_with_dialect(
1511                sql,
1512                &GreptimeDbDialect {},
1513                ParseOptions::default(),
1514            )
1515            .unwrap()
1516            .pop()
1517            .unwrap();
1518            let Statement::CreateTable(create_table) = stmt else {
1519                unreachable!()
1520            };
1521            create_to_expr(&create_table, &QueryContext::arc()).unwrap_err();
1522        }
1523    }
1524
1525    #[test]
1526    fn test_create_to_expr_with_default_timestamp_value() {
1527        let sql = "CREATE TABLE monitor (v double,ts TIMESTAMP default '2024-01-30T00:01:01',TIME INDEX (ts)) engine=mito;";
1528        let stmt =
1529            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1530                .unwrap()
1531                .pop()
1532                .unwrap();
1533
1534        let Statement::CreateTable(create_table) = stmt else {
1535            unreachable!()
1536        };
1537
1538        // query context with system timezone UTC.
1539        let expr = create_to_expr(&create_table, &QueryContext::arc()).unwrap();
1540        let ts_column = &expr.column_defs[1];
1541        let constraint = assert_ts_column(ts_column);
1542        assert!(
1543            matches!(constraint, ColumnDefaultConstraint::Value(Value::Timestamp(ts))
1544                         if ts.to_iso8601_string() == "2024-01-30 00:01:01+0000")
1545        );
1546
1547        // query context with timezone `+08:00`
1548        let ctx = QueryContextBuilder::default()
1549            .timezone(Timezone::from_tz_string("+08:00").unwrap())
1550            .build()
1551            .into();
1552        let expr = create_to_expr(&create_table, &ctx).unwrap();
1553        let ts_column = &expr.column_defs[1];
1554        let constraint = assert_ts_column(ts_column);
1555        assert!(
1556            matches!(constraint, ColumnDefaultConstraint::Value(Value::Timestamp(ts))
1557                         if ts.to_iso8601_string() == "2024-01-29 16:01:01+0000")
1558        );
1559    }
1560
1561    fn assert_ts_column(ts_column: &api::v1::ColumnDef) -> ColumnDefaultConstraint {
1562        assert_eq!("ts", ts_column.name);
1563        assert_eq!(
1564            ColumnDataType::TimestampMillisecond as i32,
1565            ts_column.data_type
1566        );
1567        assert!(!ts_column.default_constraint.is_empty());
1568
1569        ColumnDefaultConstraint::try_from(&ts_column.default_constraint[..]).unwrap()
1570    }
1571
1572    #[test]
1573    fn test_to_alter_expr() {
1574        let sql = "ALTER DATABASE greptime SET key1='value1', key2='value2';";
1575        let stmt =
1576            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1577                .unwrap()
1578                .pop()
1579                .unwrap();
1580
1581        let Statement::AlterDatabase(alter_database) = stmt else {
1582            unreachable!()
1583        };
1584
1585        let expr = to_alter_database_expr(alter_database, &QueryContext::arc()).unwrap();
1586        let kind = expr.kind.unwrap();
1587
1588        let AlterDatabaseKind::SetDatabaseOptions(SetDatabaseOptions {
1589            set_database_options,
1590        }) = kind
1591        else {
1592            unreachable!()
1593        };
1594
1595        assert_eq!(2, set_database_options.len());
1596        assert_eq!("key1", set_database_options[0].key);
1597        assert_eq!("value1", set_database_options[0].value);
1598        assert_eq!("key2", set_database_options[1].key);
1599        assert_eq!("value2", set_database_options[1].value);
1600
1601        let sql = "ALTER DATABASE greptime UNSET key1, key2;";
1602        let stmt =
1603            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1604                .unwrap()
1605                .pop()
1606                .unwrap();
1607
1608        let Statement::AlterDatabase(alter_database) = stmt else {
1609            unreachable!()
1610        };
1611
1612        let expr = to_alter_database_expr(alter_database, &QueryContext::arc()).unwrap();
1613        let kind = expr.kind.unwrap();
1614
1615        let AlterDatabaseKind::UnsetDatabaseOptions(UnsetDatabaseOptions { keys }) = kind else {
1616            unreachable!()
1617        };
1618
1619        assert_eq!(2, keys.len());
1620        assert!(keys.contains(&"key1".to_string()));
1621        assert!(keys.contains(&"key2".to_string()));
1622
1623        let sql = "ALTER TABLE monitor add column ts TIMESTAMP default '2024-01-30T00:01:01';";
1624        let stmt =
1625            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1626                .unwrap()
1627                .pop()
1628                .unwrap();
1629
1630        let Statement::AlterTable(alter_table) = stmt else {
1631            unreachable!()
1632        };
1633
1634        // query context with system timezone UTC.
1635        let expr = to_alter_table_expr(alter_table.clone(), &QueryContext::arc()).unwrap();
1636        let kind = expr.kind.unwrap();
1637
1638        let AlterTableKind::AddColumns(AddColumns { add_columns, .. }) = kind else {
1639            unreachable!()
1640        };
1641
1642        assert_eq!(1, add_columns.len());
1643        let ts_column = add_columns[0].column_def.clone().unwrap();
1644        let constraint = assert_ts_column(&ts_column);
1645        assert!(
1646            matches!(constraint, ColumnDefaultConstraint::Value(Value::Timestamp(ts))
1647                         if ts.to_iso8601_string() == "2024-01-30 00:01:01+0000")
1648        );
1649
1650        //
1651        // query context with timezone `+08:00`
1652        let ctx = QueryContextBuilder::default()
1653            .timezone(Timezone::from_tz_string("+08:00").unwrap())
1654            .build()
1655            .into();
1656        let expr = to_alter_table_expr(alter_table, &ctx).unwrap();
1657        let kind = expr.kind.unwrap();
1658
1659        let AlterTableKind::AddColumns(AddColumns { add_columns, .. }) = kind else {
1660            unreachable!()
1661        };
1662
1663        assert_eq!(1, add_columns.len());
1664        let ts_column = add_columns[0].column_def.clone().unwrap();
1665        let constraint = assert_ts_column(&ts_column);
1666        assert!(
1667            matches!(constraint, ColumnDefaultConstraint::Value(Value::Timestamp(ts))
1668                         if ts.to_iso8601_string() == "2024-01-29 16:01:01+0000")
1669        );
1670    }
1671
1672    #[test]
1673    fn test_to_alter_modify_column_type_expr() {
1674        let sql = "ALTER TABLE monitor MODIFY COLUMN mem_usage STRING;";
1675        let stmt =
1676            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1677                .unwrap()
1678                .pop()
1679                .unwrap();
1680
1681        let Statement::AlterTable(alter_table) = stmt else {
1682            unreachable!()
1683        };
1684
1685        // query context with system timezone UTC.
1686        let expr = to_alter_table_expr(alter_table.clone(), &QueryContext::arc()).unwrap();
1687        let kind = expr.kind.unwrap();
1688
1689        let AlterTableKind::ModifyColumnTypes(ModifyColumnTypes {
1690            modify_column_types,
1691        }) = kind
1692        else {
1693            unreachable!()
1694        };
1695
1696        assert_eq!(1, modify_column_types.len());
1697        let modify_column_type = &modify_column_types[0];
1698
1699        assert_eq!("mem_usage", modify_column_type.column_name);
1700        assert_eq!(
1701            ColumnDataType::String as i32,
1702            modify_column_type.target_type
1703        );
1704        assert!(modify_column_type.target_type_extension.is_none());
1705    }
1706
1707    #[test]
1708    fn test_to_repartition_request() {
1709        let sql = r#"
1710ALTER TABLE metrics REPARTITION (
1711  device_id < 100
1712) INTO (
1713  device_id < 100 AND area < 'South',
1714  device_id < 100 AND area >= 'South'
1715);"#;
1716        let stmt =
1717            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1718                .unwrap()
1719                .pop()
1720                .unwrap();
1721
1722        let Statement::AlterTable(alter_table) = stmt else {
1723            unreachable!()
1724        };
1725
1726        let request = to_repartition_request(alter_table, &QueryContext::arc()).unwrap();
1727        assert_eq!("greptime", request.catalog_name);
1728        assert_eq!("public", request.schema_name);
1729        assert_eq!("metrics", request.table_name);
1730        let RepartitionSource::Partitions {
1731            from_exprs,
1732            target_partition_columns,
1733        } = request.source
1734        else {
1735            unreachable!()
1736        };
1737        assert!(target_partition_columns.is_none());
1738        assert_eq!(
1739            from_exprs
1740                .into_iter()
1741                .map(|x| x.to_string())
1742                .collect::<Vec<_>>(),
1743            vec!["device_id < 100".to_string()]
1744        );
1745        assert_eq!(
1746            request
1747                .into_exprs
1748                .into_iter()
1749                .map(|x| x.to_string())
1750                .collect::<Vec<_>>(),
1751            vec![
1752                "device_id < 100 AND area < 'South'".to_string(),
1753                "device_id < 100 AND area >= 'South'".to_string()
1754            ]
1755        );
1756    }
1757
1758    #[test]
1759    fn test_to_repartition_request_with_target_partition_columns() {
1760        let sql = r#"
1761ALTER TABLE metrics REPARTITION (
1762  device_id < 100
1763) ON COLUMNS (device_id, area) INTO (
1764  device_id < 100 AND area < 'South',
1765  device_id < 100 AND area >= 'South'
1766);"#;
1767        let stmt =
1768            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1769                .unwrap()
1770                .pop()
1771                .unwrap();
1772
1773        let Statement::AlterTable(alter_table) = stmt else {
1774            unreachable!()
1775        };
1776
1777        let request = to_repartition_request(alter_table, &QueryContext::arc()).unwrap();
1778        let RepartitionSource::Partitions {
1779            target_partition_columns,
1780            ..
1781        } = request.source
1782        else {
1783            unreachable!()
1784        };
1785
1786        assert_eq!(
1787            target_partition_columns,
1788            Some(vec!["device_id".to_string(), "area".to_string()])
1789        );
1790    }
1791
1792    #[test]
1793    fn test_to_repartition_request_with_unpartitioned_source() {
1794        let sql = r#"
1795ALTER TABLE metrics PARTITION ON COLUMNS (device_id, area) (
1796  device_id < 100 AND area < 'South',
1797  device_id < 100 AND area >= 'South'
1798);"#;
1799        let stmt =
1800            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1801                .unwrap()
1802                .pop()
1803                .unwrap();
1804
1805        let Statement::AlterTable(alter_table) = stmt else {
1806            unreachable!()
1807        };
1808
1809        let request = to_repartition_request(alter_table, &QueryContext::arc()).unwrap();
1810        assert_eq!("greptime", request.catalog_name);
1811        assert_eq!("public", request.schema_name);
1812        assert_eq!("metrics", request.table_name);
1813        let RepartitionSource::Unpartitioned { partition_columns } = request.source else {
1814            unreachable!()
1815        };
1816        assert_eq!(partition_columns, vec!["device_id", "area"]);
1817        assert_eq!(
1818            request
1819                .into_exprs
1820                .into_iter()
1821                .map(|x| x.to_string())
1822                .collect::<Vec<_>>(),
1823            vec![
1824                "device_id < 100 AND area < 'South'".to_string(),
1825                "device_id < 100 AND area >= 'South'".to_string()
1826            ]
1827        );
1828    }
1829
1830    fn new_test_table_names() -> Vec<TableName> {
1831        vec![
1832            TableName {
1833                catalog_name: "greptime".to_string(),
1834                schema_name: "public".to_string(),
1835                table_name: "a_table".to_string(),
1836            },
1837            TableName {
1838                catalog_name: "greptime".to_string(),
1839                schema_name: "public".to_string(),
1840                table_name: "b_table".to_string(),
1841            },
1842        ]
1843    }
1844
1845    #[test]
1846    fn test_to_create_view_expr() {
1847        let sql = "CREATE VIEW test AS SELECT * FROM NUMBERS";
1848        let stmt =
1849            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1850                .unwrap()
1851                .pop()
1852                .unwrap();
1853
1854        let Statement::CreateView(stmt) = stmt else {
1855            unreachable!()
1856        };
1857
1858        let logical_plan = vec![1, 2, 3];
1859        let table_names = new_test_table_names();
1860        let columns = vec!["a".to_string()];
1861        let plan_columns = vec!["number".to_string()];
1862
1863        let expr = to_create_view_expr(
1864            stmt,
1865            logical_plan.clone(),
1866            table_names.clone(),
1867            columns.clone(),
1868            plan_columns.clone(),
1869            sql.to_string(),
1870            QueryContext::arc(),
1871        )
1872        .unwrap();
1873
1874        assert_eq!("greptime", expr.catalog_name);
1875        assert_eq!("public", expr.schema_name);
1876        assert_eq!("test", expr.view_name);
1877        assert!(!expr.create_if_not_exists);
1878        assert!(!expr.or_replace);
1879        assert_eq!(logical_plan, expr.logical_plan);
1880        assert_eq!(table_names, expr.table_names);
1881        assert_eq!(sql, expr.definition);
1882        assert_eq!(columns, expr.columns);
1883        assert_eq!(plan_columns, expr.plan_columns);
1884    }
1885
1886    #[test]
1887    fn test_to_create_view_expr_complex() {
1888        let sql = "CREATE OR REPLACE VIEW IF NOT EXISTS test.test_view AS SELECT * FROM NUMBERS";
1889        let stmt =
1890            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1891                .unwrap()
1892                .pop()
1893                .unwrap();
1894
1895        let Statement::CreateView(stmt) = stmt else {
1896            unreachable!()
1897        };
1898
1899        let logical_plan = vec![1, 2, 3];
1900        let table_names = new_test_table_names();
1901        let columns = vec!["a".to_string()];
1902        let plan_columns = vec!["number".to_string()];
1903
1904        let expr = to_create_view_expr(
1905            stmt,
1906            logical_plan.clone(),
1907            table_names.clone(),
1908            columns.clone(),
1909            plan_columns.clone(),
1910            sql.to_string(),
1911            QueryContext::arc(),
1912        )
1913        .unwrap();
1914
1915        assert_eq!("greptime", expr.catalog_name);
1916        assert_eq!("test", expr.schema_name);
1917        assert_eq!("test_view", expr.view_name);
1918        assert!(expr.create_if_not_exists);
1919        assert!(expr.or_replace);
1920        assert_eq!(logical_plan, expr.logical_plan);
1921        assert_eq!(table_names, expr.table_names);
1922        assert_eq!(sql, expr.definition);
1923        assert_eq!(columns, expr.columns);
1924        assert_eq!(plan_columns, expr.plan_columns);
1925    }
1926
1927    #[test]
1928    fn test_expr_to_create() {
1929        let sql = r#"CREATE TABLE IF NOT EXISTS `tt` (
1930  `timestamp` TIMESTAMP(9) NOT NULL,
1931  `ip_address` STRING NULL SKIPPING INDEX WITH(false_positive_rate = '0.01', granularity = '10240', type = 'BLOOM'),
1932  `username` STRING NULL,
1933  `http_method` STRING NULL INVERTED INDEX,
1934  `request_line` STRING NULL FULLTEXT INDEX WITH(analyzer = 'English', backend = 'bloom', case_sensitive = 'false', false_positive_rate = '0.01', granularity = '10240'),
1935  `protocol` STRING NULL,
1936  `status_code` INT NULL INVERTED INDEX,
1937  `response_size` BIGINT NULL,
1938  `message` STRING NULL,
1939  TIME INDEX (`timestamp`),
1940  PRIMARY KEY (`username`, `status_code`)
1941)
1942ENGINE=mito
1943WITH(
1944  append_mode = 'true'
1945)"#;
1946        let stmt =
1947            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
1948                .unwrap()
1949                .pop()
1950                .unwrap();
1951
1952        let Statement::CreateTable(original_create) = stmt else {
1953            unreachable!()
1954        };
1955
1956        // Convert CreateTable -> CreateTableExpr -> CreateTable
1957        let expr = create_to_expr(&original_create, &QueryContext::arc()).unwrap();
1958
1959        let create_table = expr_to_create(&expr, Some('`')).unwrap();
1960        let new_sql = format!("{:#}", create_table);
1961        assert_eq!(sql, new_sql);
1962    }
1963}