From 356308b348a57f8e243428866c2163cbb270d86c Mon Sep 17 00:00:00 2001
From: Tyler White <50381805+IndexSeek@users.noreply.github.com>
Date: Sat, 9 Aug 2025 02:31:58 -0400
Subject: [PATCH] feat: Include end token in `ALTER TABLE` statement (#1999)

Co-authored-by: Ifeanyi Ubah <ify1992@yahoo.com>
---
 src/ast/mod.rs           |  3 +++
 src/ast/spans.rs         | 20 +++++++++++++++++++-
 src/parser/mod.rs        |  7 +++++++
 src/test_utils.rs        |  1 +
 tests/sqlparser_mysql.rs |  1 +
 5 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 93f44e1b..5b50d020 100644
--- a/src/ast/mod.rs
+++ b/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 ")?;
diff --git a/src/ast/spans.rs b/src/ast/spans.rs
index 9705a007..e1709026 100644
--- a/src/ast/spans.rs
+++ b/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());
+    }
 }
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index d7a07d38..d6ace1eb 100644
--- a/src/parser/mod.rs
+++ b/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),
         })
     }
 
diff --git a/src/test_utils.rs b/src/test_utils.rs
index 654f2723..ab2cf89b 100644
--- a/src/test_utils.rs
+++ b/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);
diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs
index 0dcf7c0d..c47750b5 100644
--- a/tests/sqlparser_mysql.rs
+++ b/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);