Skip to content

Commit 658ebcd

Browse files
author
Dilovan Celik
committed
added functionality to handle output statement
1 parent 45420ce commit 658ebcd

File tree

6 files changed

+96
-3
lines changed

6 files changed

+96
-3
lines changed

src/ast/mod.rs

+37-2
Original file line numberDiff line numberDiff line change
@@ -3828,6 +3828,8 @@ pub enum Statement {
38283828
on: Box<Expr>,
38293829
/// Specifies the actions to perform when values match or do not match.
38303830
clauses: Vec<MergeClause>,
3831+
3832+
output: Option<Output>,
38313833
},
38323834
/// ```sql
38333835
/// CACHE [ FLAG ] TABLE <table_name> [ OPTIONS('K1' = 'V1', 'K2' = V2) ] [ AS ] [ <query> ]
@@ -5406,15 +5408,21 @@ impl fmt::Display for Statement {
54065408
table,
54075409
source,
54085410
on,
5409-
clauses,
5411+
clauses,
5412+
output
54105413
} => {
54115414
write!(
54125415
f,
54135416
"MERGE{int} {table} USING {source} ",
54145417
int = if *into { " INTO" } else { "" }
54155418
)?;
54165419
write!(f, "ON {on} ")?;
5417-
write!(f, "{}", display_separated(clauses, " "))
5420+
write!(f, "{}", display_separated(clauses, " "))?;
5421+
if output.is_some() {
5422+
let out = output.clone().unwrap();
5423+
write!(f, " {out}")?;
5424+
}
5425+
Ok(())
54185426
}
54195427
Statement::Cache {
54205428
table_name,
@@ -7900,6 +7908,33 @@ impl Display for MergeClause {
79007908
}
79017909
}
79027910

7911+
/// A Output in the end of a 'MERGE' Statement
7912+
///
7913+
/// Example:
7914+
/// OUTPUT $action, deleted.* INTO dbo.temp_products;
7915+
/// [mssql](https://learn.microsoft.com/en-us/sql/t-sql/queries/output-clause-transact-sql)
7916+
7917+
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
7918+
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
7919+
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
7920+
pub struct Output {
7921+
pub select_items: Vec<SelectItem>,
7922+
pub into_table: SelectInto,
7923+
}
7924+
7925+
impl fmt::Display for Output {
7926+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
7927+
let Output { select_items, into_table } = self;
7928+
7929+
write!(
7930+
f,
7931+
"OUTPUT {} {}",
7932+
display_comma_separated(select_items),
7933+
into_table
7934+
)
7935+
}
7936+
}
7937+
79037938
#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
79047939
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
79057940
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

src/keywords.rs

+1
Original file line numberDiff line numberDiff line change
@@ -632,6 +632,7 @@ define_keywords!(
632632
ORGANIZATION,
633633
OUT,
634634
OUTER,
635+
OUTPUT,
635636
OUTPUTFORMAT,
636637
OVER,
637638
OVERFLOW,

src/parser/mod.rs

+33-1
Original file line numberDiff line numberDiff line change
@@ -14489,7 +14489,10 @@ impl<'a> Parser<'a> {
1448914489
pub fn parse_merge_clauses(&mut self) -> Result<Vec<MergeClause>, ParserError> {
1449014490
let mut clauses = vec![];
1449114491
loop {
14492-
if self.peek_token() == Token::EOF || self.peek_token() == Token::SemiColon {
14492+
if self.peek_token() == Token::EOF
14493+
|| self.peek_token() == Token::SemiColon
14494+
|| self.peek_keyword(Keyword::OUTPUT)
14495+
{
1449314496
break;
1449414497
}
1449514498
self.expect_keyword_is(Keyword::WHEN)?;
@@ -14586,6 +14589,29 @@ impl<'a> Parser<'a> {
1458614589
Ok(clauses)
1458714590
}
1458814591

14592+
pub fn parse_output(&mut self) -> Result<Output, ParserError> {
14593+
self.expect_keyword_is(Keyword::OUTPUT)?;
14594+
let select_items = self.parse_projection()?;
14595+
self.expect_keyword_is(Keyword::INTO)?;
14596+
let temporary = self
14597+
.parse_one_of_keywords(&[Keyword::TEMP, Keyword::TEMPORARY])
14598+
.is_some();
14599+
let unlogged = self.parse_keyword(Keyword::UNLOGGED);
14600+
let table = self.parse_keyword(Keyword::TABLE);
14601+
let name = self.parse_object_name(false)?;
14602+
let into_table = SelectInto {
14603+
temporary,
14604+
unlogged,
14605+
table,
14606+
name,
14607+
};
14608+
14609+
Ok(Output {
14610+
select_items,
14611+
into_table,
14612+
})
14613+
}
14614+
1458914615
pub fn parse_merge(&mut self) -> Result<Statement, ParserError> {
1459014616
let into = self.parse_keyword(Keyword::INTO);
1459114617

@@ -14596,13 +14622,19 @@ impl<'a> Parser<'a> {
1459614622
self.expect_keyword_is(Keyword::ON)?;
1459714623
let on = self.parse_expr()?;
1459814624
let clauses = self.parse_merge_clauses()?;
14625+
let output = if self.peek_keyword(Keyword::OUTPUT) {
14626+
Some(self.parse_output()?)
14627+
} else {
14628+
None
14629+
};
1459914630

1460014631
Ok(Statement::Merge {
1460114632
into,
1460214633
table,
1460314634
source,
1460414635
on: Box::new(on),
1460514636
clauses,
14637+
output,
1460614638
})
1460714639
}
1460814640

tests/sqlparser_bigquery.rs

+2
Original file line numberDiff line numberDiff line change
@@ -1735,13 +1735,15 @@ fn parse_merge() {
17351735
},
17361736
],
17371737
};
1738+
17381739
match bigquery_and_generic().verified_stmt(sql) {
17391740
Statement::Merge {
17401741
into,
17411742
table,
17421743
source,
17431744
on,
17441745
clauses,
1746+
..
17451747
} => {
17461748
assert!(!into);
17471749
assert_eq!(

tests/sqlparser_common.rs

+17
Original file line numberDiff line numberDiff line change
@@ -9359,13 +9359,15 @@ fn parse_merge() {
93599359
source,
93609360
on,
93619361
clauses,
9362+
..
93629363
},
93639364
Statement::Merge {
93649365
into: no_into,
93659366
table: table_no_into,
93669367
source: source_no_into,
93679368
on: on_no_into,
93689369
clauses: clauses_no_into,
9370+
..
93699371
},
93709372
) => {
93719373
assert!(into);
@@ -9556,6 +9558,21 @@ fn parse_merge() {
95569558

95579559
let sql = "MERGE INTO s.bar AS dest USING newArrivals AS S ON (1 > 1) WHEN NOT MATCHED THEN INSERT VALUES (stg.A, stg.B, stg.C)";
95589560
verified_stmt(sql);
9561+
9562+
}
9563+
9564+
#[test]
9565+
fn test_merge_with_output() {
9566+
9567+
let sql = "MERGE INTO target_table USING source_table \
9568+
ON target_table.id = source_table.oooid \
9569+
WHEN MATCHED THEN \
9570+
UPDATE SET target_table.description = source_table.description \
9571+
WHEN NOT MATCHED THEN \
9572+
INSERT (ID, description) VALUES (source_table.id, source_table.description) \
9573+
OUTPUT inserted.* INTO log_target";
9574+
9575+
verified_stmt(sql);
95599576
}
95609577

95619578
#[test]

tests/sqlparser_mssql.rs

+6
Original file line numberDiff line numberDiff line change
@@ -1921,3 +1921,9 @@ fn ms() -> TestedDialects {
19211921
fn ms_and_generic() -> TestedDialects {
19221922
TestedDialects::new(vec![Box::new(MsSqlDialect {}), Box::new(GenericDialect {})])
19231923
}
1924+
1925+
#[test]
1926+
fn parse_mssql_merge_with_output() {
1927+
let stmt = "MERGE dso.products AS t USING dsi.products AS s ON s.ProductID = t.ProductID WHEN MATCHED AND NOT (t.ProductName = s.ProductName OR (ISNULL(t.ProductName, s.ProductName) IS NULL)) THEN UPDATE SET t.ProductName = s.ProductName WHEN NOT MATCHED BY TARGET THEN INSERT (ProductID, ProductName) VALUES (s.ProductID, s.ProductName) WHEN NOT MATCHED BY SOURCE THEN DELETE OUTPUT $action, deleted.ProductID INTO dsi.temp_products";
1928+
ms_and_generic().verified_stmt(stmt);
1929+
}

0 commit comments

Comments
 (0)