parser: update ast nodes to include their trailing semicolons

crates/squawk/src/reporter.rs

@@ -634,7 +634,7 @@ SELECT 1;
         );
 
         assert!(res.is_ok());
-        assert_snapshot!(String::from_utf8_lossy(&buff), @r#"[{"file":"main.sql","line":1,"column":3,"level":"Warning","message":"Missing `set lock_timeout` before potentially slow operations","help":"Configure a `lock_timeout` before this statement.","rule_name":"require-timeout-settings","column_end":62,"line_end":1},{"file":"main.sql","line":1,"column":3,"level":"Warning","message":"Missing `set statement_timeout` before potentially slow operations","help":"Configure a `statement_timeout` before this statement","rule_name":"require-timeout-settings","column_end":62,"line_end":1},{"file":"main.sql","line":1,"column":29,"level":"Warning","message":"Adding a new column that is `NOT NULL` and has no default value to an existing table effectively makes it required.","help":"Make the field nullable or add a non-VOLATILE DEFAULT","rule_name":"adding-required-field","column_end":62,"line_end":1},{"file":"main.sql","line":1,"column":29,"level":"Warning","message":"Missing `IF NOT EXISTS`, the migration can't be rerun if it fails part way through.","help":null,"rule_name":"prefer-robust-stmts","column_end":62,"line_end":1},{"file":"main.sql","line":1,"column":46,"level":"Warning","message":"Using 32-bit integer fields can result in hitting the max `int` limit.","help":"Use 64-bit integer values instead to prevent hitting this limit.","rule_name":"prefer-bigint-over-int","column_end":53,"line_end":1},{"file":"main.sql","line":2,"column":23,"level":"Warning","message":"Adding a new column that is `NOT NULL` and has no default value to an existing table effectively makes it required.","help":"Make the field nullable or add a non-VOLATILE DEFAULT","rule_name":"adding-required-field","column_end":56,"line_end":2},{"file":"main.sql","line":2,"column":23,"level":"Warning","message":"Missing `IF NOT EXISTS`, the migration can't be rerun if it fails part way through.","help":null,"rule_name":"prefer-robust-stmts","column_end":56,"line_end":2},{"file":"main.sql","line":2,"column":40,"level":"Warning","message":"Using 32-bit integer fields can result in hitting the max `int` limit.","help":"Use 64-bit integer values instead to prevent hitting this limit.","rule_name":"prefer-bigint-over-int","column_end":47,"line_end":2}]"#);
+        assert_snapshot!(String::from_utf8_lossy(&buff), @r#"[{"file":"main.sql","line":1,"column":3,"level":"Warning","message":"Missing `set lock_timeout` before potentially slow operations","help":"Configure a `lock_timeout` before this statement.","rule_name":"require-timeout-settings","column_end":63,"line_end":1},{"file":"main.sql","line":1,"column":3,"level":"Warning","message":"Missing `set statement_timeout` before potentially slow operations","help":"Configure a `statement_timeout` before this statement","rule_name":"require-timeout-settings","column_end":63,"line_end":1},{"file":"main.sql","line":1,"column":29,"level":"Warning","message":"Adding a new column that is `NOT NULL` and has no default value to an existing table effectively makes it required.","help":"Make the field nullable or add a non-VOLATILE DEFAULT","rule_name":"adding-required-field","column_end":62,"line_end":1},{"file":"main.sql","line":1,"column":29,"level":"Warning","message":"Missing `IF NOT EXISTS`, the migration can't be rerun if it fails part way through.","help":null,"rule_name":"prefer-robust-stmts","column_end":62,"line_end":1},{"file":"main.sql","line":1,"column":46,"level":"Warning","message":"Using 32-bit integer fields can result in hitting the max `int` limit.","help":"Use 64-bit integer values instead to prevent hitting this limit.","rule_name":"prefer-bigint-over-int","column_end":53,"line_end":1},{"file":"main.sql","line":2,"column":23,"level":"Warning","message":"Adding a new column that is `NOT NULL` and has no default value to an existing table effectively makes it required.","help":"Make the field nullable or add a non-VOLATILE DEFAULT","rule_name":"adding-required-field","column_end":56,"line_end":2},{"file":"main.sql","line":2,"column":23,"level":"Warning","message":"Missing `IF NOT EXISTS`, the migration can't be rerun if it fails part way through.","help":null,"rule_name":"prefer-robust-stmts","column_end":56,"line_end":2},{"file":"main.sql","line":2,"column":40,"level":"Warning","message":"Using 32-bit integer fields can result in hitting the max `int` limit.","help":"Use 64-bit integer values instead to prevent hitting this limit.","rule_name":"prefer-bigint-over-int","column_end":47,"line_end":2}]"#);
     }
 
     #[test]

crates/squawk/src/snapshots/squawk__reporter__test_check_files__check_files_invalid_syntax.snap

@@ -2,4 +2,4 @@
 source: crates/squawk/src/reporter.rs
 expression: val
 ---
-[{"column":6,"column_end":6,"file":"test.sql","help":null,"level":"Error","line":1,"line_end":1,"message":"expected SEMICOLON","rule_name":"syntax-error"},{"column":7,"column_end":7,"file":"test.sql","help":null,"level":"Error","line":1,"line_end":1,"message":"expected command, found ERROR","rule_name":"syntax-error"}]
+[{"column":7,"column_end":7,"file":"test.sql","help":null,"level":"Error","line":1,"line_end":1,"message":"expected command, found ERROR","rule_name":"syntax-error"}]

crates/squawk/src/snapshots/squawk__reporter__test_reporter__display_violations_tty.snap

@@ -6,7 +6,7 @@ warning[require-timeout-settings]: Missing `set lock_timeout` before potentially
   ╭▸ main.sql:2:4
   │
 2 │    ALTER TABLE "core_recipe" ADD COLUMN "foo" integer NOT NULL;
-  │    ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+  │    ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
   │
   ├ help: Configure a `lock_timeout` before this statement.
   ╭╴
@@ -16,7 +16,7 @@ warning[require-timeout-settings]: Missing `set statement_timeout` before potent
   ╭▸ main.sql:2:4
   │
 2 │    ALTER TABLE "core_recipe" ADD COLUMN "foo" integer NOT NULL;
-  │    ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+  │    ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
   │
   ├ help: Configure a `statement_timeout` before this statement
   ╭╴

crates/squawk/src/snapshots/squawk__reporter__test_reporter__display_violations_tty_and_github_annotations.snap

@@ -2,8 +2,8 @@
 source: crates/squawk/src/reporter.rs
 expression: "strip_ansi_codes(&String::from_utf8_lossy(&buff))"
 ---
-::warning file=main.sql,line=2,col=3,endLine=2,endColumn=62,title=require-timeout-settings::Missing `set lock_timeout` before potentially slow operations
-::warning file=main.sql,line=2,col=3,endLine=2,endColumn=62,title=require-timeout-settings::Missing `set statement_timeout` before potentially slow operations
+::warning file=main.sql,line=2,col=3,endLine=2,endColumn=63,title=require-timeout-settings::Missing `set lock_timeout` before potentially slow operations
+::warning file=main.sql,line=2,col=3,endLine=2,endColumn=63,title=require-timeout-settings::Missing `set statement_timeout` before potentially slow operations
 ::warning file=main.sql,line=2,col=29,endLine=2,endColumn=62,title=adding-required-field::Adding a new column that is `NOT NULL` and has no default value to an existing table effectively makes it required.
 ::warning file=main.sql,line=2,col=29,endLine=2,endColumn=62,title=prefer-robust-stmts::Missing `IF NOT EXISTS`, the migration can't be rerun if it fails part way through.
 ::warning file=main.sql,line=2,col=46,endLine=2,endColumn=53,title=prefer-bigint-over-int::Using 32-bit integer fields can result in hitting the max `int` limit.
@@ -14,7 +14,7 @@ warning[require-timeout-settings]: Missing `set lock_timeout` before potentially
   ╭▸ main.sql:2:4
   │
 2 │    ALTER TABLE "core_recipe" ADD COLUMN "foo" integer NOT NULL;
-  │    ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+  │    ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
   │
   ├ help: Configure a `lock_timeout` before this statement.
   ╭╴
@@ -24,7 +24,7 @@ warning[require-timeout-settings]: Missing `set statement_timeout` before potent
   ╭▸ main.sql:2:4
   │
 2 │    ALTER TABLE "core_recipe" ADD COLUMN "foo" integer NOT NULL;
-  │    ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
+  │    ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
   │
   ├ help: Configure a `statement_timeout` before this statement
   ╭╴

crates/squawk/src/snapshots/squawk__reporter__test_reporter__span_offsets.snap

@@ -10,14 +10,14 @@ CheckReport {
             file: "main.sql",
             line: 2,
             column: 3,
-            range: 5..64,
+            range: 5..65,
             level: Warning,
             message: "Missing `set lock_timeout` before potentially slow operations",
             help: Some(
                 "Configure a `lock_timeout` before this statement.",
             ),
             rule_name: "require-timeout-settings",
-            column_end: 62,
+            column_end: 63,
             line_end: 2,
             fix: Some(
                 Fix {
@@ -37,14 +37,14 @@ CheckReport {
             file: "main.sql",
             line: 2,
             column: 3,
-            range: 5..64,
+            range: 5..65,
             level: Warning,
             message: "Missing `set statement_timeout` before potentially slow operations",
             help: Some(
                 "Configure a `statement_timeout` before this statement",
             ),
             rule_name: "require-timeout-settings",
-            column_end: 62,
+            column_end: 63,
             line_end: 2,
             fix: Some(
                 Fix {

crates/squawk_fmt/src/fmt.rs

@@ -36,8 +36,6 @@ fn build_source_file(source_file: &ast::SourceFile) -> Doc<'_> {
                     } else {
                         doc = doc.append(Doc::empty_line());
                     }
-                } else if token.kind() == SyntaxKind::SEMICOLON {
-                    doc = doc.append(Doc::text(";"));
                 }
             }
         }
@@ -46,7 +44,7 @@ fn build_source_file(source_file: &ast::SourceFile) -> Doc<'_> {
 }
 
 fn build_create_table<'a>(create_table: &ast::CreateTable) -> Doc<'a> {
-    Doc::text("create")
+    let mut doc = Doc::text("create")
         .append(Doc::space())
         .append(Doc::text("table"))
         .append(Doc::space())
@@ -71,7 +69,11 @@ fn build_create_table<'a>(create_table: &ast::CreateTable) -> Doc<'a> {
                 .append(Doc::line_or_nil())
                 .group(),
         )
-        .append(Doc::text(")"))
+        .append(Doc::text(")"));
+
+    doc = doc.append(build_semicolon(create_table.semicolon_token()));
+
+    doc
 }
 
 fn build_table_arg<'a>(create_table: ast::TableArg) -> Doc<'a> {
@@ -132,9 +134,37 @@ fn build_select_doc<'a>(select: &ast::Select) -> Doc<'a> {
         );
     }
 
+    doc = doc.append(build_semicolon(select.semicolon_token()));
+
     doc.group()
 }
 
+fn build_semicolon<'a>(semi: Option<SyntaxToken>) -> Doc<'a> {
+    let Some(semi) = semi else {
+        return Doc::nil();
+    };
+    let mut doc = Doc::nil();
+    let mut comments: Vec<SyntaxToken> = vec![];
+    for next in semi.siblings_with_tokens(Direction::Prev).skip(1) {
+        match next {
+            rowan::NodeOrToken::Node(_) => break,
+            rowan::NodeOrToken::Token(token) => {
+                if token.kind() == SyntaxKind::COMMENT {
+                    comments.push(token);
+                } else if token.kind() == SyntaxKind::WHITESPACE {
+                    continue;
+                } else {
+                    break;
+                }
+            }
+        }
+    }
+    for comment in comments.iter().rev() {
+        doc = doc.append(Doc::text(comment.text().to_string()));
+    }
+    doc.append(Doc::text(";"))
+}
+
 fn build_expr<'a>(expr: ast::Expr) -> Doc<'a> {
     match expr {
         ast::Expr::ArrayExpr(array_expr) => {

crates/squawk_ide/src/code_actions/rewrite_select_as_table.rs

@@ -30,7 +30,10 @@ pub(super) fn rewrite_select_as_table(
         return None;
     };
 
-    let replacement = format!("table {}", table_name);
+    let mut replacement = format!("table {table_name}");
+    if select.semicolon_token().is_some() {
+        replacement.push(';');
+    };
 
     actions.push(CodeAction {
         title: "Rewrite as `table`".to_owned(),

crates/squawk_ide/src/code_actions/rewrite_select_as_values.rs

@@ -51,7 +51,15 @@ pub(super) fn rewrite_select_as_values(
         rows.push(format!("({})", exprs.join(", ")));
     }
 
-    let values_stmt = format!("values {}", rows.join(", "));
+    let had_semicolon = match &parent {
+        SelectContext::Compound(compound) => compound.semicolon_token().is_some(),
+        SelectContext::Single(select) => select.semicolon_token().is_some(),
+    };
+
+    let mut values_stmt = format!("values {}", rows.join(", "));
+    if had_semicolon {
+        values_stmt.push(';');
+    }
 
     let select_end = match &parent {
         SelectContext::Compound(compound) => compound.syntax().text_range().end(),

crates/squawk_ide/src/code_actions/rewrite_table_as_select.rs

@@ -17,7 +17,10 @@ pub(super) fn rewrite_table_as_select(
 
     let table_name = table.relation_name()?.syntax().text();
 
-    let replacement = format!("select * from {}", table_name);
+    let mut replacement = format!("select * from {table_name}");
+    if table.semicolon_token().is_some() {
+        replacement.push(';');
+    };
 
     actions.push(CodeAction {
         title: "Rewrite as `select`".to_owned(),

crates/squawk_ide/src/code_actions/rewrite_values_as_select.rs

@@ -47,7 +47,10 @@ pub(super) fn rewrite_values_as_select(
         select_parts.push(format!("union all\nselect {}", row_targets));
     }
 
-    let select_stmt = select_parts.join("\n");
+    let mut select_stmt = select_parts.join("\n");
+    if values.semicolon_token().is_some() {
+        select_stmt.push(';');
+    }
 
     actions.push(CodeAction {
         title: "Rewrite as `select`".to_owned(),
@@ -70,7 +73,7 @@ mod test {
     fn rewrite_values_as_select_simple() {
         assert_snapshot!(
             apply_code_action(rewrite_values_as_select, "valu$0es (1, 'one'), (2, 'two');"),
-            @r"
+            @"
         select 1 as column1, 'one' as column2
         union all
         select 2, 'two';
@@ -98,7 +101,7 @@ mod test {
     fn rewrite_values_as_select_multiple_rows() {
         assert_snapshot!(
             apply_code_action(rewrite_values_as_select, "values (1, 2), (3, 4), (5, 6$0);"),
-            @r"
+            @"
         select 1 as column1, 2 as column2
         union all
         select 3, 4
@@ -115,7 +118,7 @@ mod test {
                 rewrite_values_as_select,
                 "with cte as (select 1) val$0ues (1, 'one'), (2, 'two');"
             ),
-            @r"
+            @"
         with cte as (select 1) select 1 as column1, 'one' as column2
         union all
         select 2, 'two';
@@ -146,7 +149,7 @@ mod test {
     fn rewrite_values_as_select_on_row_content() {
         assert_snapshot!(
             apply_code_action(rewrite_values_as_select, "values (1$0, 2), (3, 4);"),
-            @r"
+            @"
         select 1 as column1, 2 as column2
         union all
         select 3, 4;