feat: impl parser and operator for CREATE FLOW (#3806)

* feat: impl parser for CREATE TASK

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* finish parser

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* wip expr

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* finish expr

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* rename output to sink

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* fix parser

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* remove debug code

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* upload lock file

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* rename symbol

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* Apply suggestions from code review

Co-authored-by: Jeremyhi <jiachun_feng@proton.me>

* remove other task word

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* task name to flow name

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* one more comment

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

---------

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
Co-authored-by: Jeremyhi <jiachun_feng@proton.me>
This commit is contained in:
Ruihang Xia
2024-04-28 12:26:49 +00:00
committed by GitHub
co-authored by Jeremyhi
parent 1bbde15a15
commit 7ef18c0915
13 changed files with 387 additions and 17 deletions
+186 -4
View File
@@ -32,7 +32,8 @@ use crate::error::{
};
use crate::parser::ParserContext;
use crate::statements::create::{
CreateDatabase, CreateExternalTable, CreateTable, CreateTableLike, Partitions, TIME_INDEX,
CreateDatabase, CreateExternalTable, CreateFlow, CreateTable, CreateTableLike, Partitions,
TIME_INDEX,
};
use crate::statements::statement::Statement;
use crate::statements::{get_data_type_by_alias_name, OptionMap};
@@ -40,6 +41,10 @@ use crate::util::parse_option_string;
pub const ENGINE: &str = "ENGINE";
pub const MAXVALUE: &str = "MAXVALUE";
pub const FLOW: &str = "FLOW";
pub const SINK: &str = "SINK";
pub const EXPIRE: &str = "EXPIRE";
pub const WHEN: &str = "WHEN";
/// Parses create [table] statement
impl<'a> ParserContext<'a> {
@@ -52,6 +57,35 @@ impl<'a> ParserContext<'a> {
Keyword::EXTERNAL => self.parse_create_external_table(),
Keyword::OR => {
let _ = self.parser.next_token();
self.parser
.expect_keyword(Keyword::REPLACE)
.context(SyntaxSnafu)?;
match self.parser.next_token().token {
Token::Word(w) => match w.keyword {
Keyword::NoKeyword => {
let uppercase = w.value.to_uppercase();
match uppercase.as_str() {
FLOW => self.parse_create_flow(true),
_ => self.unsupported(w.to_string()),
}
}
_ => self.unsupported(w.to_string()),
},
_ => self.unsupported(w.to_string()),
}
}
Keyword::NoKeyword => {
let _ = self.parser.next_token();
let uppercase = w.value.to_uppercase();
match uppercase.as_str() {
FLOW => self.parse_create_flow(false),
_ => self.unsupported(w.to_string()),
}
}
_ => self.unsupported(w.to_string()),
},
unexpected => self.unsupported(unexpected.to_string()),
@@ -137,6 +171,64 @@ impl<'a> ParserContext<'a> {
Ok(Statement::CreateTable(create_table))
}
/// "CREATE FLOW" clause
fn parse_create_flow(&mut self, or_replace: bool) -> Result<Statement> {
let if_not_exists = self.parse_if_not_exist()?;
let task_name = self.intern_parse_table_name()?;
self.parser
.expect_token(&Token::make_keyword(SINK))
.context(SyntaxSnafu)?;
self.parser
.expect_keyword(Keyword::TO)
.context(SyntaxSnafu)?;
let output_table_name = self.intern_parse_table_name()?;
let expire_when = if self
.parser
.consume_tokens(&[Token::make_keyword(EXPIRE), Token::make_keyword(WHEN)])
{
Some(self.parser.parse_expr().context(error::SyntaxSnafu)?)
} else {
None
};
let comment = if self.parser.parse_keyword(Keyword::COMMENT) {
match self.parser.next_token() {
TokenWithLocation {
token: Token::SingleQuotedString(value, ..),
..
} => Some(value),
unexpected => {
return self
.parser
.expected("string", unexpected)
.context(SyntaxSnafu)
}
}
} else {
None
};
self.parser
.expect_keyword(Keyword::AS)
.context(SyntaxSnafu)?;
let query = Box::new(self.parser.parse_query().context(error::SyntaxSnafu)?);
Ok(Statement::CreateFlow(CreateFlow {
flow_name: task_name,
sink_table_name: output_table_name,
or_replace,
if_not_exists,
expire_when,
comment,
query,
}))
}
fn parse_if_not_exist(&mut self) -> Result<bool> {
match self.parser.peek_token().token {
Token::Word(w) if Keyword::IF != w.keyword => return Ok(false),
@@ -185,8 +277,7 @@ impl<'a> ParserContext<'a> {
Ok(options.into())
}
/// "PARTITION BY ..." syntax:
// TODO(ruihang): docs
/// "PARTITION BY ..." clause
fn parse_partitions(&mut self) -> Result<Option<Partitions>> {
if !self.parser.parse_keyword(Keyword::PARTITION) {
return Ok(None);
@@ -737,7 +828,7 @@ mod tests {
use common_catalog::consts::FILE_ENGINE;
use common_error::ext::ErrorExt;
use sqlparser::ast::ColumnOption::NotNull;
use sqlparser::ast::{BinaryOperator, ObjectName, Value};
use sqlparser::ast::{BinaryOperator, Expr, Function, Interval, ObjectName, Value};
use super::*;
use crate::dialect::GreptimeDbDialect;
@@ -945,6 +1036,97 @@ mod tests {
);
}
#[test]
fn test_parse_create_flow() {
let sql = r"
CREATE OR REPLACE FLOW IF NOT EXISTS task_1
SINK TO schema_1.table_1
EXPIRE WHEN timestamp < now() - INTERVAL '5m'
COMMENT 'test comment'
AS
SELECT max(c1), min(c2) FROM schema_2.table_2;";
let stmts =
ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
.unwrap();
assert_eq!(1, stmts.len());
let create_task = match &stmts[0] {
Statement::CreateFlow(c) => c,
_ => unreachable!(),
};
let expected = CreateFlow {
flow_name: ObjectName(vec![Ident {
value: "task_1".to_string(),
quote_style: None,
}]),
sink_table_name: ObjectName(vec![
Ident {
value: "schema_1".to_string(),
quote_style: None,
},
Ident {
value: "table_1".to_string(),
quote_style: None,
},
]),
or_replace: true,
if_not_exists: true,
expire_when: Some(Expr::BinaryOp {
left: Box::new(Expr::Identifier(Ident {
value: "timestamp".to_string(),
quote_style: None,
})),
op: BinaryOperator::Lt,
right: Box::new(Expr::BinaryOp {
left: Box::new(Expr::Function(Function {
name: ObjectName(vec![Ident {
value: "now".to_string(),
quote_style: None,
}]),
args: vec![],
filter: None,
null_treatment: None,
over: None,
distinct: false,
special: false,
order_by: vec![],
})),
op: BinaryOperator::Minus,
right: Box::new(Expr::Interval(Interval {
value: Box::new(Expr::Value(Value::SingleQuotedString("5m".to_string()))),
leading_field: None,
leading_precision: None,
last_field: None,
fractional_seconds_precision: None,
})),
}),
}),
comment: Some("test comment".to_string()),
// ignore query parse result
query: create_task.query.clone(),
};
assert_eq!(create_task, &expected);
// create flow without `OR REPLACE`, `IF NOT EXISTS`, `EXPIRE WHEN` and `COMMENT`
let sql = r"
CREATE FLOW task_2
SINK TO schema_1.table_1
AS
SELECT max(c1), min(c2) FROM schema_2.table_2;";
let stmts =
ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
.unwrap();
assert_eq!(1, stmts.len());
let create_task = match &stmts[0] {
Statement::CreateFlow(c) => c,
_ => unreachable!(),
};
assert!(!create_task.or_replace);
assert!(!create_task.if_not_exists);
assert!(create_task.expire_when.is_none());
assert!(create_task.comment.is_none());
}
#[test]
fn test_validate_create() {
let sql = r"
+41 -1
View File
@@ -16,7 +16,7 @@ use std::fmt::{Display, Formatter};
use common_catalog::consts::FILE_ENGINE;
use itertools::Itertools;
use sqlparser::ast::Expr;
use sqlparser::ast::{Expr, Query};
use sqlparser_derive::{Visit, VisitMut};
use crate::ast::{ColumnDef, Ident, ObjectName, TableConstraint, Value as SqlValue};
@@ -237,6 +237,46 @@ impl Display for CreateTableLike {
}
}
#[derive(Debug, PartialEq, Eq, Clone, Visit, VisitMut)]
pub struct CreateFlow {
/// Flow name
pub flow_name: ObjectName,
/// Output (sink) table name
pub sink_table_name: ObjectName,
/// Whether to replace existing task
pub or_replace: bool,
/// Create if not exist
pub if_not_exists: bool,
/// `EXPIRE_WHEN`
pub expire_when: Option<Expr>,
/// Comment string
pub comment: Option<String>,
/// SQL statement
pub query: Box<Query>,
}
impl Display for CreateFlow {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "CREATE ")?;
if self.or_replace {
write!(f, "OR REPLACE ")?;
}
write!(f, "TASK ")?;
if self.if_not_exists {
write!(f, "IF NOT EXISTS ")?;
}
write!(f, "{} ", &self.flow_name)?;
write!(f, "OUTPUT AS {} ", &self.sink_table_name)?;
if let Some(expire_when) = &self.expire_when {
write!(f, "EXPIRE WHEN {} ", expire_when)?;
}
if let Some(comment) = &self.comment {
write!(f, "COMMENT '{}' ", comment)?;
}
write!(f, "AS {}", &self.query)
}
}
#[cfg(test)]
mod tests {
use std::assert_matches::assert_matches;
+4
View File
@@ -18,6 +18,7 @@ use datafusion_sql::parser::Statement as DfStatement;
use sqlparser::ast::Statement as SpStatement;
use sqlparser_derive::{Visit, VisitMut};
use super::create::CreateFlow;
use super::drop::DropDatabase;
use super::show::ShowVariables;
use crate::error::{ConvertToDfStatementSnafu, Error};
@@ -54,6 +55,8 @@ pub enum Statement {
CreateExternalTable(CreateExternalTable),
// CREATE TABLE ... LIKE
CreateTableLike(CreateTableLike),
// CREATE FLOW
CreateFlow(CreateFlow),
// DROP TABLE
DropTable(DropTable),
// DROP DATABASE
@@ -100,6 +103,7 @@ impl Display for Statement {
Statement::CreateTable(s) => s.fmt(f),
Statement::CreateExternalTable(s) => s.fmt(f),
Statement::CreateTableLike(s) => s.fmt(f),
Statement::CreateFlow(s) => s.fmt(f),
Statement::DropTable(s) => s.fmt(f),
Statement::DropDatabase(s) => s.fmt(f),
Statement::CreateDatabase(s) => s.fmt(f),
+44 -1
View File
@@ -12,9 +12,10 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use std::collections::HashSet;
use std::fmt::{Display, Formatter};
use sqlparser::ast::{Expr, ObjectName, SqlOption, Value};
use sqlparser::ast::{Expr, ObjectName, Query, SetExpr, SqlOption, TableFactor, Value};
use crate::error::{InvalidTableOptionValueSnafu, Result};
@@ -50,3 +51,45 @@ pub fn parse_option_string(option: SqlOption) -> Result<(String, String)> {
let k = key.value.to_lowercase();
Ok((k, v))
}
/// Walk through a [Query] and extract all the tables referenced in it.
pub fn extract_tables_from_query(query: &Query) -> impl Iterator<Item = ObjectName> {
let mut names = HashSet::new();
extract_tables_from_set_expr(&query.body, &mut names);
names.into_iter()
}
/// Helper function for [extract_tables_from_query].
///
/// Handle [SetExpr].
fn extract_tables_from_set_expr(set_expr: &SetExpr, names: &mut HashSet<ObjectName>) {
match set_expr {
SetExpr::Select(select) => {
for from in &select.from {
table_factor_to_object_name(&from.relation, names);
for join in &from.joins {
table_factor_to_object_name(&join.relation, names);
}
}
}
SetExpr::Query(query) => {
extract_tables_from_set_expr(&query.body, names);
}
SetExpr::SetOperation { left, right, .. } => {
extract_tables_from_set_expr(left, names);
extract_tables_from_set_expr(right, names);
}
SetExpr::Values(_) | SetExpr::Insert(_) | SetExpr::Update(_) | SetExpr::Table(_) => {}
};
}
/// Helper function for [extract_tables_from_query].
///
/// Handle [TableFactor].
fn table_factor_to_object_name(table_factor: &TableFactor, names: &mut HashSet<ObjectName>) {
if let TableFactor::Table { name, .. } = table_factor {
names.insert(name.to_owned());
}
}