PostgreSQL: Add support for DO statement

fuzz/fuzz_seeds/postgres_do.sql

@@ -0,0 +1 @@
+DO $$BEGIN RAISE NOTICE 'hello'; END$$

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

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$$

fuzz/fuzz_seeds/postgres_do_single_quoted.sql

@@ -0,0 +1 @@
+DO 'BEGIN RAISE NOTICE ''hi''; END'

fuzz/fuzz_seeds/postgres_do_tagged_dollar_quote.sql

@@ -0,0 +1 @@
+DO $body$BEGIN PERFORM $$nested$$; END$body$

src/ast/mod.rs

@@ -4930,6 +4930,12 @@ pub enum Statement {
         channel: Ident,
     },
     /// ```sql
+    /// DO [ LANGUAGE lang_name ] code
+    /// ```
+    ///
+    /// See Postgres <https://www.postgresql.org/docs/current/sql-do.html>
+    Do(DoStatement),
+    /// ```sql
     /// NOTIFY channel [ , payload ]
     /// ```
     /// send a notification event together with an optional "payload" string to channel
@@ -6506,6 +6512,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 {
@@ -11581,6 +11588,33 @@ impl fmt::Display for WaitForStatement {
     }
 }
 
+/// PostgreSQL `DO` statement, executing an anonymous code block.
+///
+/// ```sql
+/// DO [ LANGUAGE lang_name ] code
+/// ```
+///
+/// See <https://www.postgresql.org/docs/current/sql-do.html>
+#[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 {
+    /// Optional procedural language name.
+    pub language: Option<Ident>,
+    /// The anonymous code block, a string literal.
+    pub code: ValueWithSpan,
+}
+
+impl fmt::Display for DoStatement {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "DO ")?;
+        if let Some(language) = &self.language {
+            write!(f, "LANGUAGE {language} ")?;
+        }
+        write!(f, "{}", self.code)
+    }
+}
+
 /// Represents a `Return` statement.
 ///
 /// [MsSql triggers](https://learn.microsoft.com/en-us/sql/t-sql/statements/create-trigger-transact-sql)
@@ -12433,6 +12467,12 @@ impl From<ThrowStatement> for Statement {
     }
 }
 
+impl From<DoStatement> for Statement {
+    fn from(d: DoStatement) -> Self {
+        Self::Do(d)
+    }
+}
+
 impl From<Function> for Statement {
     fn from(f: Function) -> Self {
         Self::Call(f)

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, 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(),
@@ -2632,6 +2633,18 @@ impl Spanned for comments::CommentWithSpan {
     }
 }
 
+impl Spanned for DoStatement {
+    fn span(&self) -> Span {
+        let DoStatement { language, code } = self;
+        union_spans(
+            language
+                .iter()
+                .map(|l| l.span)
+                .chain(core::iter::once(code.span)),
+        )
+    }
+}
+
 #[cfg(test)]
 pub mod tests {
     use crate::ast::Value;
@@ -3196,4 +3209,13 @@ WHERE id = 1
             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$$");
+    }
 }

src/dialect/mod.rs

@@ -1289,6 +1289,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

src/dialect/postgresql.rs

@@ -231,6 +231,10 @@ impl Dialect for PostgreSqlDialect {
         true
     }
 
+    fn supports_do_statement(&self) -> bool {
+        true
+    }
+
     /// see <https://www.postgresql.org/docs/current/sql-createtable.html#SQL-CREATETABLE-EXCLUDE>
     fn supports_exclude_constraint(&self) -> bool {
         true

src/parser/mod.rs

@@ -717,6 +717,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 => {
@@ -1506,6 +1507,30 @@ impl<'a> Parser<'a> {
         Ok(Statement::UNLISTEN { channel })
     }
 
+    /// Parse `DO` statement.
+    pub fn parse_do(&mut self) -> Result<Statement, ParserError> {
+        let mut language = self.parse_do_language()?;
+        match &self.peek_token_ref().token {
+            Token::SingleQuotedString(_)
+            | Token::DollarQuotedString(_)
+            | Token::EscapedStringLiteral(_) => {}
+            _ => return self.expected_ref("string literal", self.peek_token_ref()),
+        }
+        let code = self.parse_value()?;
+        if language.is_none() {
+            language = self.parse_do_language()?;
+        }
+        Ok(Statement::Do(DoStatement { language, code }))
+    }
+
+    fn parse_do_language(&mut self) -> Result<Option<Ident>, ParserError> {
+        if self.parse_keyword(Keyword::LANGUAGE) {
+            Ok(Some(self.parse_identifier()?))
+        } else {
+            Ok(None)
+        }
+    }
+
     /// Parse `NOTIFY` statement.
     pub fn parse_notify(&mut self) -> Result<Statement, ParserError> {
         let channel = self.parse_identifier()?;

tests/sqlparser_common.rs

@@ -20433,3 +20433,58 @@ fn parse_bang_not_renders_apart_from_operand() {
     dialects.verified_stmt("SET eaac_cion = ! !o");
     dialects.one_statement_parses_to("SET eaac_cion = ! ! o", "SET eaac_cion = ! !o");
 }
+
+#[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 { language, code }) => {
+            assert_eq!(language, None);
+            assert_eq!(
+                code.value,
+                Value::DollarQuotedString(DollarQuotedString {
+                    value: "BEGIN RAISE NOTICE 'hi'; END".to_string(),
+                    tag: None,
+                })
+            );
+        }
+        _ => unreachable!(),
+    }
+
+    match dialects.verified_stmt("DO LANGUAGE plpgsql $body$BEGIN NULL; END$body$") {
+        Statement::Do(DoStatement { language, .. }) => {
+            assert_eq!(language, Some(Ident::new("plpgsql")));
+        }
+        _ => unreachable!(),
+    }
+
+    dialects.verified_stmt("DO 'BEGIN NULL; END'");
+    dialects.one_statement_parses_to(
+        "DO $$BEGIN NULL; END$$ LANGUAGE plpgsql",
+        "DO LANGUAGE plpgsql $$BEGIN NULL; END$$",
+    );
+    assert_eq!(
+        dialects
+            .parse_sql_statements("DO LANGUAGE plpgsql $$BEGIN NULL; END$$ LANGUAGE sql")
+            .unwrap_err(),
+        ParserError::ParserError("Expected: end of statement, found: LANGUAGE".to_string())
+    );
+
+    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())
+    );
+}