use std::collections::{BTreeMap, BTreeSet};
use sqlparser::ast::{Query, SetExpr, TableFactor};
use sqlparser::dialect::GenericDialect;
use sqlparser::parser::Parser;
use sqlparser::tokenizer::Token;
use crate::error::{Result, ThundError};
use crate::ir::{
Dataset, DatasetSchema, Expectation, Flow, OnViolation, OutputType, Pipeline, SourceSpec,
};
pub fn from_sql(pipeline_name: impl Into<String>, text: &str) -> Result<Pipeline> {
let name = pipeline_name.into();
let result: Result<Pipeline> = (|| {
let mut p = Pipeline::new(name.clone());
for stmt in ThundParser::new(text)?.parse_all()? {
match stmt {
ThundStatement::Dataset { dataset, flow } => {
p = p.with_dataset(dataset);
if let Some(f) = flow {
p = p.with_flow(f);
}
}
ThundStatement::Flow(f) => p = p.with_flow(f),
}
}
p.validate()?;
Ok(p)
})();
crate::functional_status(
"knut-thund/sql",
"from_sql",
result.is_ok(),
&match &result {
Ok(p) => format!("{} ({} datasets)", p.name, p.datasets.len()),
Err(e) => e.to_string(),
},
);
result
}
pub fn to_sql(pipeline: &Pipeline) -> Result<String> {
let result: Result<String> = (|| {
let order = pipeline.topo_order().ok_or(ThundError::Cyclic)?;
let mut out = String::new();
for name in &order {
let Some(d) = pipeline.dataset(name) else {
continue;
};
let flow = pipeline.flows.iter().find(|f| &f.target == name);
out.push_str(&render_dataset(d, flow));
out.push('\n');
}
for f in &pipeline.flows {
if pipeline.dataset(&f.target).is_some()
&& pipeline.flows.iter().position(|g| g.target == f.target)
!= pipeline.flows.iter().position(|g| g.name == f.name)
{
out.push_str(&render_standalone_flow(f));
out.push('\n');
}
}
Ok(out)
})();
crate::functional_status("knut-thund/sql", "to_sql", result.is_ok(), &pipeline.name);
result
}
enum ThundStatement {
Dataset {
dataset: Dataset,
flow: Option<Flow>,
},
Flow(Flow),
}
struct ThundParser<'a> {
parser: Parser<'a>,
}
#[derive(Default)]
struct TableClauses {
partition_cols: Vec<String>,
comment: Option<String>,
properties: BTreeMap<String, String>,
}
const DIALECT: GenericDialect = GenericDialect {};
impl<'a> ThundParser<'a> {
fn new(sql: &str) -> Result<ThundParser<'_>> {
let parser = Parser::new(&DIALECT)
.try_with_sql(sql)
.map_err(|e| ThundError::Sql(e.to_string()))?;
Ok(ThundParser { parser })
}
fn parse_all(&mut self) -> Result<Vec<ThundStatement>> {
let mut out = Vec::new();
loop {
while self.parser.consume_token(&Token::SemiColon) {}
if self.parser.peek_token().token == Token::EOF {
break;
}
out.push(self.parse_statement()?);
}
Ok(out)
}
fn word(&mut self, w: &str) -> bool {
let tok = self.parser.peek_token();
if matches!(&tok.token, Token::Word(word) if word.value.eq_ignore_ascii_case(w)) {
self.parser.next_token();
return true;
}
false
}
fn words(&mut self, ws: &[&str]) -> bool {
for (i, w) in ws.iter().enumerate() {
match &self.parser.peek_nth_token(i).token {
Token::Word(word) if word.value.eq_ignore_ascii_case(w) => {}
_ => return false,
}
}
for _ in ws {
self.parser.next_token();
}
true
}
fn expect_word(&mut self, w: &str) -> Result<()> {
if self.word(w) {
Ok(())
} else {
Err(ThundError::Sql(format!(
"expected `{w}`, found `{}`",
self.parser.peek_token()
)))
}
}
fn parse_statement(&mut self) -> Result<ThundStatement> {
if !self.word("CREATE") {
return Err(ThundError::Sql(format!(
"expected a CREATE definition, found `{}` — a Thund SQL script \
contains only dataset/flow definitions",
self.parser.peek_token()
)));
}
if self.words(&["MATERIALIZED", "VIEW"]) {
self.parse_dataset(OutputType::MaterializedView, false)
} else if self.words(&["STREAMING", "TABLE"]) {
self.parse_dataset(OutputType::Table, true)
} else if self.words(&["TEMPORARY", "VIEW"])
|| self.words(&["TEMP", "VIEW"])
|| self.word("VIEW")
{
self.parse_dataset(OutputType::TemporaryView, false)
} else if self.word("FLOW") {
self.parse_create_flow()
} else {
Err(ThundError::Sql(format!(
"unsupported CREATE form at `{}` — expected MATERIALIZED VIEW, \
STREAMING TABLE, [TEMPORARY] VIEW or FLOW",
self.parser.peek_token()
)))
}
}
fn parse_dataset(
&mut self,
output_type: OutputType,
streaming: bool,
) -> Result<ThundStatement> {
let name = self.object_name()?;
let mut schema = DatasetSchema::new();
let mut expectations = Vec::new();
if self.parser.consume_token(&Token::LParen) {
self.parse_column_list(&mut schema, &mut expectations)?;
}
let clauses = self.parse_table_clauses()?;
let mut dataset = Dataset::new(&name, output_type);
dataset.comment = clauses.comment;
dataset.partition_cols = clauses.partition_cols;
dataset.properties = clauses.properties;
if !schema.is_empty() {
dataset = dataset.with_schema(schema);
}
if !self.word("AS") {
return Ok(ThundStatement::Dataset {
dataset,
flow: None,
});
}
let flow = self.parse_flow_body(format!("f_{name}"), &name, streaming, expectations)?;
Ok(ThundStatement::Dataset {
dataset,
flow: Some(flow),
})
}
fn parse_create_flow(&mut self) -> Result<ThundStatement> {
let flow_name = self.object_name()?;
self.expect_word("AS")?;
self.expect_word("INSERT")?;
self.expect_word("INTO")?;
let target = self.object_name()?;
self.words(&["BY", "NAME"]);
let flow = self.parse_flow_body(flow_name, &target, true, Vec::new())?;
Ok(ThundStatement::Flow(flow))
}
fn parse_flow_body(
&mut self,
flow_name: impl Into<String>,
target: &str,
prefer_streaming: bool,
expectations: Vec<Expectation>,
) -> Result<Flow> {
let mut query = self
.parser
.parse_query()
.map_err(|e| ThundError::Sql(e.to_string()))?;
let streamed = rewrite_stream(&mut query);
let reads = collect_reads(&query);
let sql = query.to_string();
let mut flow = if let Some(source) = streamed.first() {
let mut f = Flow::streaming(
flow_name,
target,
SourceSpec::Cdc {
source: source.clone(),
},
);
f.reads = reads;
f
} else if prefer_streaming && !reads.is_empty() {
let mut f = Flow::streaming(
flow_name,
target,
SourceSpec::Cdc {
source: reads[0].clone(),
},
);
f.reads = reads;
f
} else {
Flow::batch(flow_name, target, reads)
};
flow.query = Some(sql);
flow.expectations = expectations;
Ok(flow)
}
fn parse_column_list(
&mut self,
schema: &mut DatasetSchema,
expectations: &mut Vec<Expectation>,
) -> Result<()> {
loop {
if self.word("CONSTRAINT") {
expectations.push(self.parse_expectation()?);
} else {
let col = self
.parser
.parse_identifier()
.map_err(|e| ThundError::Sql(e.to_string()))?;
let ty = self
.parser
.parse_data_type()
.map_err(|e| ThundError::Sql(e.to_string()))?;
let nullable = !self.words(&["NOT", "NULL"]);
*schema = std::mem::take(schema).field(col.value, sql_type_to_arrow(&ty), nullable);
}
if self.parser.consume_token(&Token::Comma) {
continue;
}
self.parser
.expect_token(&Token::RParen)
.map_err(|e| ThundError::Sql(e.to_string()))?;
return Ok(());
}
}
fn parse_expectation(&mut self) -> Result<Expectation> {
let name = self
.parser
.parse_identifier()
.map_err(|e| ThundError::Sql(e.to_string()))?;
self.expect_word("EXPECT")?;
self.parser
.expect_token(&Token::LParen)
.map_err(|e| ThundError::Sql(e.to_string()))?;
let expr = self
.parser
.parse_expr()
.map_err(|e| ThundError::Sql(e.to_string()))?;
self.parser
.expect_token(&Token::RParen)
.map_err(|e| ThundError::Sql(e.to_string()))?;
let policy = if self.words(&["ON", "VIOLATION"]) {
if self.words(&["DROP", "ROW"]) {
OnViolation::Drop
} else if self.words(&["FAIL", "UPDATE"]) {
OnViolation::Fail
} else {
return Err(ThundError::Sql(format!(
"expected DROP ROW or FAIL UPDATE after ON VIOLATION, found `{}`",
self.parser.peek_token()
)));
}
} else {
OnViolation::Warn
};
Ok(Expectation::new(name.value, expr.to_string()).on(policy))
}
fn parse_table_clauses(&mut self) -> Result<TableClauses> {
let mut c = TableClauses::default();
loop {
if self.words(&["PARTITIONED", "BY"]) {
c.partition_cols = self.parse_partition_list()?;
} else if self.word("COMMENT") {
c.comment = Some(
self.parser
.parse_literal_string()
.map_err(|e| ThundError::Sql(e.to_string()))?,
);
} else if self.word("TBLPROPERTIES") {
c.properties = self.parse_property_list()?;
} else {
return Ok(c);
}
}
}
fn parse_property_list(&mut self) -> Result<BTreeMap<String, String>> {
self.parser
.expect_token(&Token::LParen)
.map_err(|e| ThundError::Sql(e.to_string()))?;
let mut props = BTreeMap::new();
if self.parser.consume_token(&Token::RParen) {
return Ok(props);
}
loop {
let key = self.property_token()?;
let _ = self.parser.consume_token(&Token::Eq);
let value = self.property_token()?;
props.insert(key, value);
if self.parser.consume_token(&Token::Comma) {
continue;
}
self.parser
.expect_token(&Token::RParen)
.map_err(|e| ThundError::Sql(e.to_string()))?;
return Ok(props);
}
}
fn property_token(&mut self) -> Result<String> {
let mut s = match self.parser.next_token().token {
Token::SingleQuotedString(v) | Token::DoubleQuotedString(v) => return Ok(v),
Token::Word(w) => w.value,
Token::Number(n, _) => n,
other => {
return Err(ThundError::Sql(format!(
"expected a TBLPROPERTIES key or value, found `{other}`"
)));
}
};
while self.parser.consume_token(&Token::Period) {
match self.parser.next_token().token {
Token::Word(w) => {
s.push('.');
s.push_str(&w.value);
}
other => {
return Err(ThundError::Sql(format!(
"expected an identifier after `.` in a TBLPROPERTIES key, found `{other}`"
)));
}
}
}
Ok(s)
}
fn parse_partition_list(&mut self) -> Result<Vec<String>> {
self.parser
.expect_token(&Token::LParen)
.map_err(|e| ThundError::Sql(e.to_string()))?;
let mut cols = Vec::new();
loop {
let expr = self
.parser
.parse_expr()
.map_err(|e| ThundError::Sql(e.to_string()))?;
cols.push(expr.to_string());
if self.parser.consume_token(&Token::Comma) {
continue;
}
self.parser
.expect_token(&Token::RParen)
.map_err(|e| ThundError::Sql(e.to_string()))?;
return Ok(cols);
}
}
fn object_name(&mut self) -> Result<String> {
self.parser
.parse_object_name(false)
.map(|n| n.to_string())
.map_err(|e| ThundError::Sql(e.to_string()))
}
}
fn rewrite_stream(query: &mut Query) -> Vec<String> {
let mut found = Vec::new();
walk_factors_mut(query, &mut |tf: &mut TableFactor| {
if let TableFactor::Table {
name, args, alias, ..
} = tf
{
let is_stream = name
.0
.last()
.map(|p| {
p.to_string()
.trim_matches('"')
.eq_ignore_ascii_case("STREAM")
})
.unwrap_or(false);
if !is_stream {
return;
}
if let Some(a) = args.take() {
if let Some(first) = a.args.first() {
let src = first.to_string();
if let Ok(parsed) = Parser::new(&DIALECT)
.try_with_sql(&src)
.and_then(|mut p| p.parse_object_name(false))
{
*name = parsed;
found.push(src);
return;
}
}
return;
}
if let Some(al) = alias.take() {
let src = al.name.value.clone();
if let Ok(parsed) = Parser::new(&DIALECT)
.try_with_sql(&src)
.and_then(|mut p| p.parse_object_name(false))
{
*name = parsed;
found.push(src);
}
}
}
});
found
}
fn collect_reads(query: &Query) -> Vec<String> {
let mut ctes = BTreeSet::new();
if let Some(with) = &query.with {
for cte in &with.cte_tables {
ctes.insert(cte.alias.name.value.to_ascii_lowercase());
}
}
let mut names = BTreeSet::new();
walk_factors(query, &mut |tf: &TableFactor| {
if let TableFactor::Table { name, .. } = tf {
let n = name.to_string();
if !ctes.contains(&n.to_ascii_lowercase()) {
names.insert(n);
}
}
});
names.into_iter().collect()
}
fn walk_factors(query: &Query, f: &mut impl FnMut(&TableFactor)) {
if let Some(with) = &query.with {
for cte in &with.cte_tables {
walk_factors(&cte.query, f);
}
}
walk_set_expr(&query.body, f);
}
fn walk_set_expr(body: &SetExpr, f: &mut impl FnMut(&TableFactor)) {
match body {
SetExpr::Select(select) => {
for twj in &select.from {
walk_factor(&twj.relation, f);
for join in &twj.joins {
walk_factor(&join.relation, f);
}
}
}
SetExpr::Query(q) => walk_factors(q, f),
SetExpr::SetOperation { left, right, .. } => {
walk_set_expr(left, f);
walk_set_expr(right, f);
}
_ => {}
}
}
fn walk_factor(tf: &TableFactor, f: &mut impl FnMut(&TableFactor)) {
f(tf);
match tf {
TableFactor::Derived { subquery, .. } => walk_factors(subquery, f),
TableFactor::NestedJoin {
table_with_joins, ..
} => {
walk_factor(&table_with_joins.relation, f);
for join in &table_with_joins.joins {
walk_factor(&join.relation, f);
}
}
_ => {}
}
}
fn walk_factors_mut(query: &mut Query, f: &mut impl FnMut(&mut TableFactor)) {
if let Some(with) = &mut query.with {
for cte in &mut with.cte_tables {
walk_factors_mut(&mut cte.query, f);
}
}
walk_set_expr_mut(&mut query.body, f);
}
fn walk_set_expr_mut(body: &mut SetExpr, f: &mut impl FnMut(&mut TableFactor)) {
match body {
SetExpr::Select(select) => {
for twj in &mut select.from {
walk_factor_mut(&mut twj.relation, f);
for join in &mut twj.joins {
walk_factor_mut(&mut join.relation, f);
}
}
}
SetExpr::Query(q) => walk_factors_mut(q, f),
SetExpr::SetOperation { left, right, .. } => {
walk_set_expr_mut(left, f);
walk_set_expr_mut(right, f);
}
_ => {}
}
}
fn walk_factor_mut(tf: &mut TableFactor, f: &mut impl FnMut(&mut TableFactor)) {
f(tf);
match tf {
TableFactor::Derived { subquery, .. } => walk_factors_mut(subquery, f),
TableFactor::NestedJoin {
table_with_joins, ..
} => {
walk_factor_mut(&mut table_with_joins.relation, f);
for join in &mut table_with_joins.joins {
walk_factor_mut(&mut join.relation, f);
}
}
_ => {}
}
}
fn sql_type_to_arrow(ty: &sqlparser::ast::DataType) -> String {
use sqlparser::ast::DataType as D;
match ty {
D::Boolean | D::Bool => "Boolean".into(),
D::TinyInt(_) => "Int8".into(),
D::SmallInt(_) => "Int16".into(),
D::Int(_) | D::Integer(_) => "Int32".into(),
D::BigInt(_) => "Int64".into(),
D::TinyIntUnsigned(_) | D::UTinyInt => "UInt8".into(),
D::SmallIntUnsigned(_) => "UInt16".into(),
D::IntUnsigned(_) | D::IntegerUnsigned(_) => "UInt32".into(),
D::BigIntUnsigned(_) | D::UBigInt => "UInt64".into(),
D::Real | D::Float(_) | D::Float4 => "Float32".into(),
D::Double(_) | D::DoublePrecision | D::Float8 => "Float64".into(),
D::Date => "Date32".into(),
D::Timestamp(_, _) | D::Datetime(_) => "Timestamp(Microsecond, None)".into(),
D::Time(_, _) => "Time64(Microsecond)".into(),
D::Decimal(info) | D::Numeric(info) | D::Dec(info) => match info {
sqlparser::ast::ExactNumberInfo::PrecisionAndScale(p, s) => {
format!("Decimal128({p}, {s})")
}
sqlparser::ast::ExactNumberInfo::Precision(p) => format!("Decimal128({p}, 0)"),
sqlparser::ast::ExactNumberInfo::None => "Decimal128(38, 10)".into(),
},
D::Bytea | D::Blob(_) | D::Binary(_) | D::Varbinary(_) => "Binary".into(),
_ => "Utf8".into(),
}
}
fn render_dataset(d: &Dataset, flow: Option<&Flow>) -> String {
let kind = match d.output_type {
OutputType::MaterializedView => "CREATE MATERIALIZED VIEW",
OutputType::Table => "CREATE STREAMING TABLE",
OutputType::TemporaryView => "CREATE TEMPORARY VIEW",
OutputType::Sink => "CREATE STREAMING TABLE",
};
let mut s = format!("{kind} {}", d.name);
let expectations = flow.map(|f| f.expectations.as_slice()).unwrap_or(&[]);
if !d.schema.is_empty() || !expectations.is_empty() {
let mut parts: Vec<String> = d
.schema
.fields
.iter()
.map(|f| {
format!(
"{} {}{}",
f.name,
arrow_to_sql_type(&f.arrow_type),
if f.nullable { "" } else { " NOT NULL" }
)
})
.collect();
for e in expectations {
let policy = match e.on_violation {
OnViolation::Warn => "",
OnViolation::Drop => " ON VIOLATION DROP ROW",
OnViolation::Fail => " ON VIOLATION FAIL UPDATE",
};
parts.push(format!(
"CONSTRAINT {} EXPECT ({}){}",
e.name, e.constraint, policy
));
}
s.push_str(&format!(" (\n {}\n)", parts.join(",\n ")));
}
if !d.partition_cols.is_empty() {
s.push_str(&format!(
"\nPARTITIONED BY ({})",
d.partition_cols.join(", ")
));
}
if let Some(c) = &d.comment {
s.push_str(&format!("\nCOMMENT '{}'", c.replace('\'', "''")));
}
if !d.properties.is_empty() {
let props = d
.properties
.iter()
.map(|(k, v)| format!("'{}' = '{}'", k.replace('\'', "''"), v.replace('\'', "''")))
.collect::<Vec<_>>()
.join(", ");
s.push_str(&format!("\nTBLPROPERTIES ({props})"));
}
match flow.and_then(|f| f.query.clone()) {
Some(q) => s.push_str(&format!("\nAS {q};\n")),
None => s.push_str(";\n"),
}
s
}
fn render_standalone_flow(f: &Flow) -> String {
let body = f.query.clone().unwrap_or_else(|| "SELECT 1".into());
format!(
"CREATE FLOW {} AS INSERT INTO {} BY NAME {body};\n",
f.name, f.target
)
}
fn arrow_to_sql_type(arrow: &str) -> &str {
match arrow {
"Boolean" => "BOOLEAN",
"Int8" => "TINYINT",
"Int16" => "SMALLINT",
"Int32" => "INT",
"Int64" => "BIGINT",
"Float32" => "REAL",
"Float64" => "DOUBLE",
"Date32" => "DATE",
"Timestamp(Microsecond, None)" => "TIMESTAMP",
"Binary" => "BINARY",
"Utf8" => "STRING",
other => other,
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn stream_function_and_bare_forms_both_resolve() {
for sql in [
"CREATE MATERIALIZED VIEW src AS SELECT 1 AS x;
CREATE STREAMING TABLE st AS SELECT * FROM STREAM(src);",
"CREATE MATERIALIZED VIEW src AS SELECT 1 AS x;
CREATE STREAMING TABLE st AS SELECT * FROM STREAM src;",
] {
let p = from_sql("p", sql).unwrap();
let f = p.flows.iter().find(|f| f.target == "st").unwrap();
assert_eq!(f.reads, vec!["src".to_string()]);
assert!(
f.kind.is_unbounded(),
"STREAM(...) must produce a streaming flow"
);
assert!(!f.query.as_ref().unwrap().to_uppercase().contains("STREAM"));
}
}
#[test]
fn partition_transform_survives_verbatim() {
let p = from_sql(
"p",
"CREATE MATERIALIZED VIEW e PARTITIONED BY (months(event_ts), bucket(16, id))
AS SELECT 1 AS id, CAST(1 AS TIMESTAMP) AS event_ts;",
)
.unwrap();
assert_eq!(
p.dataset("e").unwrap().partition_cols,
vec!["months(event_ts)".to_string(), "bucket(16, id)".to_string()]
);
}
#[test]
fn cte_names_are_not_graph_edges() {
let p = from_sql(
"p",
"CREATE MATERIALIZED VIEW base AS SELECT 1 AS x;
CREATE MATERIALIZED VIEW roll AS
WITH tmp AS (SELECT x FROM base) SELECT x FROM tmp;",
)
.unwrap();
let f = p.flows.iter().find(|f| f.target == "roll").unwrap();
assert_eq!(
f.reads,
vec!["base".to_string()],
"`tmp` is a CTE, not a dataset"
);
}
}