feat: include end token in ALTER TABLE statement for accurate span calculation

Author: IndexSeek

src/ast/mod.rs

@@ -3382,6 +3382,9 @@ 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)
+        /// Used for accurate span calculation
+        end_token: AttachedToken,
     },
     /// ```sql
     /// ALTER INDEX
@@ -5419,6 +5422,7 @@ impl fmt::Display for Statement {
                 location,
                 on_cluster,
                 iceberg,
+                end_token: _,
             } => {
                 if *iceberg {
                     write!(f, "ALTER ICEBERG TABLE ")?;

src/ast/spans.rs

@@ -434,10 +434,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 {
@@ -2548,4 +2550,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

@@ -9227,6 +9227,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,
@@ -9235,6 +9241,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_common.rs

@@ -4758,6 +4758,25 @@ fn parse_alter_table() {
     }
 }
 
+#[test]
+fn alter_table_span_includes_semicolon() {
+    use sqlparser::ast::Spanned;
+    use sqlparser::dialect::PostgreSqlDialect;
+    use sqlparser::parser::Parser;
+
+    let sql = r#"-- foo
+ALTER TABLE users
+  ADD COLUMN foo
+  varchar; -- hi there"#;
+    let dialect = PostgreSqlDialect {};
+    let ast = Parser::parse_sql(&dialect, sql).unwrap();
+    let stmt = &ast[0];
+    let span = stmt.span();
+
+    assert_eq!(span.end.line, 4);
+    assert_eq!(span.end.column, 11);
+}
+
 #[test]
 fn parse_rename_table() {
     match verified_stmt("RENAME TABLE test.test1 TO test_db.test2") {

tests/sqlparser_mysql.rs

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