feat: Include end token in ALTER TABLE statement (#1999)

Co-authored-by: Ifeanyi Ubah <[email protected]>

Author: IndexSeek

src/ast/mod.rs

@@ -3384,6 +3384,8 @@ pub enum Statement {
         /// Snowflake "ICEBERG" clause for Iceberg tables
         /// <https://docs.snowflake.com/en/sql-reference/sql/alter-iceberg-table>
         iceberg: bool,
+        /// Token that represents the end of the statement (semicolon or EOF)
+        end_token: AttachedToken,
     },
     /// ```sql
     /// ALTER INDEX
@@ -5442,6 +5444,7 @@ impl fmt::Display for Statement {
                 location,
                 on_cluster,
                 iceberg,
+                end_token: _,
             } => {
                 if *iceberg {
                     write!(f, "ALTER ICEBERG TABLE ")?;

src/ast/spans.rs

@@ -435,10 +435,12 @@ impl Spanned for Statement {
                 location: _,
                 on_cluster,
                 iceberg: _,
+                end_token,
             } => union_spans(
                 core::iter::once(name.span())
                     .chain(operations.iter().map(|i| i.span()))
-                    .chain(on_cluster.iter().map(|i| i.span)),
+                    .chain(on_cluster.iter().map(|i| i.span))
+                    .chain(core::iter::once(end_token.0.span)),
             ),
             Statement::AlterIndex { name, operation } => name.span().union(&operation.span()),
             Statement::AlterView {
@@ -2553,4 +2555,20 @@ pub mod tests {
             stmt => panic!("expected query; got {stmt:?}"),
         }
     }
+
+    #[test]
+    fn test_alter_table_multiline_span() {
+        let sql = r#"-- foo
+ALTER TABLE users
+  ADD COLUMN foo
+  varchar; -- hi there"#;
+
+        let r = Parser::parse_sql(&crate::dialect::PostgreSqlDialect {}, sql).unwrap();
+        assert_eq!(1, r.len());
+
+        let stmt_span = r[0].span();
+
+        assert_eq!(stmt_span.start, (2, 13).into());
+        assert_eq!(stmt_span.end, (4, 11).into());
+    }
 }

src/parser/mod.rs

@@ -9263,6 +9263,12 @@ impl<'a> Parser<'a> {
             });
         }
 
+        let end_token = if self.peek_token_ref().token == Token::SemiColon {
+            self.peek_token_ref().clone()
+        } else {
+            self.get_current_token().clone()
+        };
+
         Ok(Statement::AlterTable {
             name: table_name,
             if_exists,
@@ -9271,6 +9277,7 @@ impl<'a> Parser<'a> {
             location,
             on_cluster,
             iceberg,
+            end_token: AttachedToken(end_token),
         })
     }
 

src/test_utils.rs

@@ -351,6 +351,7 @@ pub fn alter_table_op_with_name(stmt: Statement, expected_name: &str) -> AlterTa
             on_cluster: _,
             location: _,
             iceberg,
+            end_token: _,
         } => {
             assert_eq!(name.to_string(), expected_name);
             assert!(!if_exists);

tests/sqlparser_mysql.rs

@@ -2643,6 +2643,7 @@ fn parse_alter_table_add_column() {
             iceberg,
             location: _,
             on_cluster: _,
+            end_token: _,
         } => {
             assert_eq!(name.to_string(), "tab");
             assert!(!if_exists);