diff --git a/fuzz/fuzz_seeds/postgres_do.sql b/fuzz/fuzz_seeds/postgres_do.sql new file mode 100644 index 0000000000..a27d1d18f5 --- /dev/null +++ b/fuzz/fuzz_seeds/postgres_do.sql @@ -0,0 +1 @@ +DO $$BEGIN RAISE NOTICE 'hello'; END$$ \ No newline at end of file diff --git a/fuzz/fuzz_seeds/postgres_do_language_after.sql b/fuzz/fuzz_seeds/postgres_do_language_after.sql new file mode 100644 index 0000000000..30f6e3aeac --- /dev/null +++ b/fuzz/fuzz_seeds/postgres_do_language_after.sql @@ -0,0 +1 @@ +DO $$BEGIN UPDATE t SET a = a + 1 WHERE b > 0; END$$ LANGUAGE plpgsql \ No newline at end of file diff --git a/fuzz/fuzz_seeds/postgres_do_language_before.sql b/fuzz/fuzz_seeds/postgres_do_language_before.sql new file mode 100644 index 0000000000..259368c967 --- /dev/null +++ b/fuzz/fuzz_seeds/postgres_do_language_before.sql @@ -0,0 +1 @@ +DO LANGUAGE plpgsql $$DECLARE n int; BEGIN SELECT count(*) INTO n FROM t; END$$ \ No newline at end of file diff --git a/fuzz/fuzz_seeds/postgres_do_single_quoted.sql b/fuzz/fuzz_seeds/postgres_do_single_quoted.sql new file mode 100644 index 0000000000..8165b2478b --- /dev/null +++ b/fuzz/fuzz_seeds/postgres_do_single_quoted.sql @@ -0,0 +1 @@ +DO 'BEGIN RAISE NOTICE ''hi''; END' \ No newline at end of file diff --git a/fuzz/fuzz_seeds/postgres_do_tagged_dollar_quote.sql b/fuzz/fuzz_seeds/postgres_do_tagged_dollar_quote.sql new file mode 100644 index 0000000000..32728aa87e --- /dev/null +++ b/fuzz/fuzz_seeds/postgres_do_tagged_dollar_quote.sql @@ -0,0 +1 @@ +DO $body$BEGIN PERFORM $$nested$$; END$body$ \ No newline at end of file diff --git a/src/ast/mod.rs b/src/ast/mod.rs index f585121add..8b14269ec6 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -4937,6 +4937,12 @@ pub enum Statement { channel: Ident, }, /// ```sql + /// DO { code | LANGUAGE lang_name } [ ... ] + /// ``` + /// + /// See Postgres + Do(DoStatement), + /// ```sql /// NOTIFY channel [ , payload ] /// ``` /// send a notification event together with an optional "payload" string to channel @@ -6514,6 +6520,7 @@ impl fmt::Display for Statement { write!(f, "UNLISTEN {channel}")?; Ok(()) } + Statement::Do(s) => write!(f, "{s}"), Statement::NOTIFY { channel, payload } => { write!(f, "NOTIFY {channel}")?; if let Some(payload) = payload { @@ -11587,6 +11594,50 @@ impl fmt::Display for WaitForStatement { } } +/// PostgreSQL `DO` statement, executing an anonymous code block. +/// +/// ```sql +/// DO { code | LANGUAGE lang_name } [ ... ] +/// ``` +/// +/// See +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +pub struct DoStatement { + /// The code blocks and `LANGUAGE` options, in source order. + pub options: Vec, +} + +impl fmt::Display for DoStatement { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "DO {}", display_separated(&self.options, " ")) + } +} + +/// An item of a [`DoStatement`]. +#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)] +#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))] +#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] +pub enum DoOption { + /// `LANGUAGE lang_name` + Language(Ident), + /// `LANGUAGE` followed by the name as a string literal. + LanguageLiteral(ValueWithSpan), + /// The anonymous code block, a string literal. + Code(ValueWithSpan), +} + +impl fmt::Display for DoOption { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + DoOption::Language(language) => write!(f, "LANGUAGE {language}"), + DoOption::LanguageLiteral(language) => write!(f, "LANGUAGE {language}"), + DoOption::Code(code) => write!(f, "{code}"), + } + } +} + /// Represents a `Return` statement. /// /// [MsSql triggers](https://learn.microsoft.com/en-us/sql/t-sql/statements/create-trigger-transact-sql) @@ -12439,6 +12490,12 @@ impl From for Statement { } } +impl From for Statement { + fn from(d: DoStatement) -> Self { + Self::Do(d) + } +} + impl From for Statement { fn from(f: Function) -> Self { Self::Call(f) diff --git a/src/ast/spans.rs b/src/ast/spans.rs index a17eba20cc..6cdbc3e375 100644 --- a/src/ast/spans.rs +++ b/src/ast/spans.rs @@ -33,13 +33,13 @@ use super::{ AttachedToken, BeginEndStatements, CaseStatement, CloseCursor, ClusteredIndex, ColumnDef, ColumnOption, ColumnOptionDef, ConditionalStatementBlock, ConditionalStatements, ConflictTarget, ConnectByKind, ConstraintCharacteristics, CopySource, CreateIndex, CreateTable, - CreateTableOptions, Cte, Delete, DoUpdate, ExceptSelectItem, ExcludeConstraintElement, - ExcludeSelectItem, Expr, ExprWithAlias, Fetch, ForValues, FromTable, Function, FunctionArg, - FunctionArgExpr, FunctionArgumentClause, FunctionArgumentList, FunctionArguments, GroupByExpr, - HavingBound, IfStatement, IlikeSelectItem, IndexColumn, Insert, Interpolate, InterpolateExpr, - Join, JoinConstraint, JoinOperator, JsonPath, JsonPathElem, LateralView, LimitClause, - MatchRecognizePattern, Measure, Merge, MergeAction, MergeClause, MergeInsertExpr, - MergeInsertKind, MergeUpdateExpr, MergeUpdateKind, NamedParenthesizedList, + CreateTableOptions, Cte, Delete, DoOption, DoStatement, DoUpdate, ExceptSelectItem, + ExcludeConstraintElement, ExcludeSelectItem, Expr, ExprWithAlias, Fetch, ForValues, FromTable, + Function, FunctionArg, FunctionArgExpr, FunctionArgumentClause, FunctionArgumentList, + FunctionArguments, GroupByExpr, HavingBound, IfStatement, IlikeSelectItem, IndexColumn, Insert, + Interpolate, InterpolateExpr, Join, JoinConstraint, JoinOperator, JsonPath, JsonPathElem, + LateralView, LimitClause, MatchRecognizePattern, Measure, Merge, MergeAction, MergeClause, + MergeInsertExpr, MergeInsertKind, MergeUpdateExpr, MergeUpdateKind, NamedParenthesizedList, NamedWindowDefinition, ObjectName, ObjectNamePart, Offset, OnConflict, OnConflictAction, OnInsert, OpenStatement, OrderBy, OrderByExpr, OrderByKind, OutputClause, Parens, Partition, PartitionBoundValue, PivotValueSource, ProjectionSelect, Query, RaiseStatement, @@ -498,6 +498,7 @@ impl Spanned for Statement { Statement::ShowObjects { .. } => Span::empty(), Statement::ShowViews { .. } => Span::empty(), Statement::LISTEN { .. } => Span::empty(), + Statement::Do(s) => s.span(), Statement::NOTIFY { .. } => Span::empty(), Statement::LoadData { .. } => Span::empty(), Statement::UNLISTEN { .. } => Span::empty(), @@ -1105,9 +1106,7 @@ impl Spanned for CreateTableOptions { match self { CreateTableOptions::None => Span::empty(), CreateTableOptions::With(vec) => union_spans(vec.iter().map(|i| i.span())), - CreateTableOptions::Options(vec) => { - union_spans(vec.as_slice().iter().map(|i| i.span())) - } + CreateTableOptions::Options(vec) => union_spans(vec.iter().map(|i| i.span())), CreateTableOptions::Plain(vec) => union_spans(vec.iter().map(|i| i.span())), CreateTableOptions::TableProperties(vec) => union_spans(vec.iter().map(|i| i.span())), } @@ -2656,6 +2655,21 @@ impl Spanned for comments::CommentWithSpan { } } +impl Spanned for DoStatement { + fn span(&self) -> Span { + union_spans(self.options.iter().map(|option| option.span())) + } +} + +impl Spanned for DoOption { + fn span(&self) -> Span { + match self { + DoOption::Language(language) => language.span, + DoOption::LanguageLiteral(value) | DoOption::Code(value) => value.span, + } + } +} + #[cfg(test)] pub mod tests { use crate::ast::Value; @@ -3251,4 +3265,13 @@ SETTINGS update_parallel_mode = 'auto'"#; stmt => panic!("expected ALTER TABLE; got {stmt:?}"), } } + + #[test] + fn test_do_statement_span() { + let dialect = &crate::dialect::PostgreSqlDialect {}; + let sql = "DO LANGUAGE plpgsql $$BEGIN NULL; END$$"; + let test = SpanTest::new(dialect, sql); + let r = Parser::parse_sql(dialect, sql).unwrap(); + assert_eq!(test.get_source(r[0].span()), "plpgsql $$BEGIN NULL; END$$"); + } } diff --git a/src/dialect/mod.rs b/src/dialect/mod.rs index d8f93e49ef..627b2d07cb 100644 --- a/src/dialect/mod.rs +++ b/src/dialect/mod.rs @@ -1308,6 +1308,11 @@ pub trait Dialect: Debug + Any { false } + /// Returns true if the dialect supports the `DO` statement + fn supports_do_statement(&self) -> bool { + false + } + /// Returns true if the dialect supports the `LISTEN`, `UNLISTEN` and `NOTIFY` statements fn supports_listen_notify(&self) -> bool { false diff --git a/src/dialect/postgresql.rs b/src/dialect/postgresql.rs index 3bec6ceba3..bfbfe913cd 100644 --- a/src/dialect/postgresql.rs +++ b/src/dialect/postgresql.rs @@ -231,6 +231,10 @@ impl Dialect for PostgreSqlDialect { true } + fn supports_do_statement(&self) -> bool { + true + } + /// see fn supports_exclude_constraint(&self) -> bool { true diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 7d2006b2d1..d039b02578 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -723,6 +723,7 @@ impl<'a> Parser<'a> { Keyword::LISTEN if self.dialect.supports_listen_notify() => self.parse_listen(), Keyword::UNLISTEN if self.dialect.supports_listen_notify() => self.parse_unlisten(), Keyword::NOTIFY if self.dialect.supports_listen_notify() => self.parse_notify(), + Keyword::DO if self.dialect.supports_do_statement() => self.parse_do(), // `PRAGMA` is sqlite specific https://www.sqlite.org/pragma.html Keyword::PRAGMA => self.parse_pragma(), Keyword::UNLOAD => { @@ -1512,6 +1513,46 @@ impl<'a> Parser<'a> { Ok(Statement::UNLISTEN { channel }) } + /// Parse `DO` statement. + pub fn parse_do(&mut self) -> Result { + let mut options = Vec::new(); + loop { + let language = self.parse_keyword(Keyword::LANGUAGE); + let string_literal = matches!( + self.peek_token_ref().token, + Token::SingleQuotedString(_) + | Token::DollarQuotedString(_) + | Token::EscapedStringLiteral(_) + | Token::UnicodeStringLiteral(_) + ); + let option = match (language, string_literal) { + (true, true) => DoOption::LanguageLiteral(self.parse_value()?), + (true, false) + if matches!(self.peek_token_ref().token, Token::NationalStringLiteral(_)) => + { + // PostgreSQL lexes `N'...'` as the `NCHAR` keyword followed + // by a string literal, so `LANGUAGE N'...'` is the language + // name `nchar` followed by the code block. + let next_token = self.next_token(); + let span = next_token.span; + let Token::NationalStringLiteral(code) = next_token.token else { + return self.expected("national string literal", next_token); + }; + options.push(DoOption::Language(Ident::with_span(span, "nchar"))); + DoOption::Code(Value::SingleQuotedString(code).with_span(span)) + } + (true, false) => DoOption::Language(self.parse_identifier()?), + (false, true) => DoOption::Code(self.parse_value()?), + (false, false) if options.is_empty() => { + return self.expected_ref("string literal", self.peek_token_ref()); + } + (false, false) => break, + }; + options.push(option); + } + Ok(Statement::Do(DoStatement { options })) + } + /// Parse `NOTIFY` statement. pub fn parse_notify(&mut self) -> Result { let channel = self.parse_identifier()?; diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs index 53fbe1b49f..f9147b5d5b 100644 --- a/tests/sqlparser_common.rs +++ b/tests/sqlparser_common.rs @@ -6095,9 +6095,9 @@ fn parse_window_functions() { expr_from_projection(&select.projection[0]) ); - for i in 0..EXPECTED_PROJ_QTY { + for proj in select.projection.iter().take(EXPECTED_PROJ_QTY) { assert!(matches!( - expr_from_projection(&select.projection[i]), + expr_from_projection(proj), Expr::Function(Function { over: Some(WindowType::WindowSpec(WindowSpec { window_name: None, @@ -20642,6 +20642,59 @@ fn parse_position_in_connect_by_list_after_state_change() { } } +#[test] +fn parse_do_statement() { + let dialects = all_dialects_where(|d| d.supports_do_statement()); + + match dialects.verified_stmt("DO $$BEGIN RAISE NOTICE 'hi'; END$$") { + Statement::Do(DoStatement { options }) => assert_eq!( + options, + vec![DoOption::Code( + Value::DollarQuotedString(DollarQuotedString { + value: "BEGIN RAISE NOTICE 'hi'; END".to_string(), + tag: None, + }) + .with_empty_span() + )] + ), + _ => unreachable!(), + } + + match dialects.verified_stmt("DO LANGUAGE plpgsql $body$BEGIN NULL; END$body$") { + Statement::Do(DoStatement { options }) => { + assert_eq!(options[0], DoOption::Language(Ident::new("plpgsql"))); + } + _ => unreachable!(), + } + + dialects.verified_stmt("DO 'BEGIN NULL; END'"); + dialects.verified_stmt("DO U&'BEGIN NULL; END'"); + dialects.verified_stmt("DO $$BEGIN NULL; END$$ LANGUAGE plpgsql"); + dialects.verified_stmt("DO LANGUAGE 'plpgsql' $$BEGIN NULL; END$$"); + dialects.verified_stmt("DO LANGUAGE $$plpgsql$$ $$BEGIN NULL; END$$"); + dialects.verified_stmt("DO $$BEGIN NULL; END$$ LANGUAGE E'plpgsql'"); + dialects.verified_stmt("DO LANGUAGE plpgsql"); + dialects.verified_stmt("DO $$BEGIN NULL; END$$ $$BEGIN NULL; END$$"); + dialects.verified_stmt("DO LANGUAGE plpgsql LANGUAGE sql $$BEGIN NULL; END$$"); + + assert_eq!( + dialects.parse_sql_statements("DO 1").unwrap_err(), + ParserError::ParserError("Expected: string literal, found: 1".to_string()) + ); + assert_eq!( + dialects.parse_sql_statements("DO LANGUAGE").unwrap_err(), + ParserError::ParserError("Expected: identifier, found: EOF".to_string()) + ); + + let dialects = all_dialects_where(|d| !d.supports_do_statement()); + assert_eq!( + dialects + .parse_sql_statements("DO $$BEGIN NULL; END$$") + .unwrap_err(), + ParserError::ParserError("Expected: an SQL statement, found: DO".to_string()) + ); +} + #[test] fn parse_array_suffix_after_nested_type() { // When `>>` closes two types at once, a following array suffix belongs to the outer one. diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index 11a6b38713..813b1ff681 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -4176,6 +4176,16 @@ fn parse_xml_typed_string() { } } +#[test] +fn parse_do_statement() { + let dialects = pg(); + dialects.verified_stmt("DO LANGUAGE plpgsql LANGUAGE sql $$BEGIN NULL; END$$"); + dialects.one_statement_parses_to( + "DO LANGUAGE N'BEGIN NULL; END'", + "DO LANGUAGE nchar 'BEGIN NULL; END'", + ); +} + fn pg() -> TestedDialects { TestedDialects::new(vec![Box::new(PostgreSqlDialect {})]) }