PostgreSQL: Make DO statement a named struct

Author: devtechedge

src/ast/mod.rs

@@ -4934,12 +4934,7 @@ pub enum Statement {
     /// ```
     ///
     /// See Postgres <https://www.postgresql.org/docs/current/sql-do.html>
-    Do {
-        /// Optional procedural language name.
-        language: Option<Ident>,
-        /// The anonymous code block, a string literal.
-        code: ValueWithSpan,
-    },
+    Do(DoStatement),
     /// ```sql
     /// NOTIFY channel [ , payload ]
     /// ```
@@ -6517,13 +6512,7 @@ impl fmt::Display for Statement {
                 write!(f, "UNLISTEN {channel}")?;
                 Ok(())
             }
-            Statement::Do { language, code } => {
-                write!(f, "DO ")?;
-                if let Some(language) = language {
-                    write!(f, "LANGUAGE {language} ")?;
-                }
-                write!(f, "{code}")
-            }
+            Statement::Do(s) => write!(f, "{s}"),
             Statement::NOTIFY { channel, payload } => {
                 write!(f, "NOTIFY {channel}")?;
                 if let Some(payload) = payload {
@@ -11599,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)
@@ -12451,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,7 +498,7 @@ impl Spanned for Statement {
             Statement::ShowObjects { .. } => Span::empty(),
             Statement::ShowViews { .. } => Span::empty(),
             Statement::LISTEN { .. } => Span::empty(),
-            Statement::Do { .. } => Span::empty(),
+            Statement::Do(s) => s.span(),
             Statement::NOTIFY { .. } => Span::empty(),
             Statement::LoadData { .. } => Span::empty(),
             Statement::UNLISTEN { .. } => Span::empty(),
@@ -2633,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;
@@ -3197,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/parser/mod.rs

@@ -1520,7 +1520,7 @@ impl<'a> Parser<'a> {
         if language.is_none() {
             language = self.parse_do_language()?;
         }
-        Ok(Statement::Do { language, code })
+        Ok(Statement::Do(DoStatement { language, code }))
     }
 
     fn parse_do_language(&mut self) -> Result<Option<Ident>, ParserError> {

tests/sqlparser_common.rs

@@ -20439,7 +20439,7 @@ 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 { language, code } => {
+        Statement::Do(DoStatement { language, code }) => {
             assert_eq!(language, None);
             assert_eq!(
                 code.value,
@@ -20453,7 +20453,7 @@ fn parse_do_statement() {
     }
 
     match dialects.verified_stmt("DO LANGUAGE plpgsql $body$BEGIN NULL; END$body$") {
-        Statement::Do { language, .. } => {
+        Statement::Do(DoStatement { language, .. }) => {
             assert_eq!(language, Some(Ident::new("plpgsql")));
         }
         _ => unreachable!(),