Merge branch 'update-grammar-rules-1' of github.com:supabase-community/postgres-language-server into test/update-insert-copy-grammar

Author: juleswritescode

crates/pgls_completions/src/providers/keywords.rs

@@ -63,6 +63,8 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("boolean"),
     SqlKeyword::new("brin"),
     SqlKeyword::new("btree"),
+    SqlKeyword::new("buffer_usage_limit"),
+    SqlKeyword::new("buffers"),
     SqlKeyword::new("by"),
     SqlKeyword::new("bytea"),
     SqlKeyword::new("cache"),
@@ -77,7 +79,6 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("check"),
     SqlKeyword::new("collate"),
     SqlKeyword::new("column"),
-    SqlKeyword::new("columns"),
     SqlKeyword::new("comment").starts_statement(),
     SqlKeyword::new("commit").starts_statement(),
     SqlKeyword::new("committed"),
@@ -89,6 +90,7 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("constraints"),
     SqlKeyword::new("copy").starts_statement(),
     SqlKeyword::new("cost"),
+    SqlKeyword::new("costs"),
     SqlKeyword::new("create").starts_statement(),
     SqlKeyword::new("cross"),
     SqlKeyword::new("csv"),
@@ -109,7 +111,6 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("delete").starts_statement(),
     SqlKeyword::new("delimiter"),
     SqlKeyword::new("desc"),
-    SqlKeyword::new("disable_page_skipping"),
     SqlKeyword::new("distinct"),
     SqlKeyword::new("do").starts_statement(),
     SqlKeyword::new("double"),
@@ -149,6 +150,7 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("function"),
     SqlKeyword::new("functions"),
     SqlKeyword::new("generated"),
+    SqlKeyword::new("generic_plan"),
     SqlKeyword::new("gin"),
     SqlKeyword::new("gist"),
     SqlKeyword::new("grant").starts_statement(),
@@ -159,12 +161,12 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("having"),
     SqlKeyword::new("header"),
     SqlKeyword::new("if"),
+    SqlKeyword::new("ignore"),
     SqlKeyword::new("immediate"),
     SqlKeyword::new("immutable"),
     SqlKeyword::new("in").require_prefix(),
     SqlKeyword::new("increment"),
     SqlKeyword::new("index"),
-    SqlKeyword::new("index_cleanup"),
     SqlKeyword::new("inet"),
     SqlKeyword::new("inherit"),
     SqlKeyword::new("initially"),
@@ -194,14 +196,15 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("limit"),
     SqlKeyword::new("list"),
     SqlKeyword::new("local"),
-    SqlKeyword::new("location"),
+    SqlKeyword::new("log_verbosity"),
     SqlKeyword::new("logged"),
     SqlKeyword::new("main"),
     SqlKeyword::new("maintain"),
     SqlKeyword::new("match"),
     SqlKeyword::new("matched"),
     SqlKeyword::new("materialized"),
     SqlKeyword::new("maxvalue"),
+    SqlKeyword::new("memory"),
     SqlKeyword::new("merge").starts_statement(),
     SqlKeyword::new("minvalue"),
     SqlKeyword::new("money"),
@@ -213,7 +216,6 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("none"),
     SqlKeyword::new("not").require_prefix(),
     SqlKeyword::new("nothing"),
-    SqlKeyword::new("nowait"),
     SqlKeyword::new("null").require_prefix(),
     SqlKeyword::new("nulls"),
     SqlKeyword::new("numeric"),
@@ -224,6 +226,7 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("oids"),
     SqlKeyword::new("old"),
     SqlKeyword::new("on"),
+    SqlKeyword::new("on_error"),
     SqlKeyword::new("only").require_prefix(),
     SqlKeyword::new("option"),
     SqlKeyword::new("or").require_prefix(),
@@ -238,19 +241,16 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("owner"),
     SqlKeyword::new("parallel"),
     SqlKeyword::new("partition"),
-    SqlKeyword::new("partitioned"),
     SqlKeyword::new("password"),
     SqlKeyword::new("permissive"),
     SqlKeyword::new("plain"),
     SqlKeyword::new("policy"),
-    SqlKeyword::new("precedes"),
     SqlKeyword::new("preceding"),
     SqlKeyword::new("precision"),
     SqlKeyword::new("primary"),
     SqlKeyword::new("privileges"),
     SqlKeyword::new("procedure"),
     SqlKeyword::new("procedures"),
-    SqlKeyword::new("process_toast"),
     SqlKeyword::new("program"),
     SqlKeyword::new("public"),
     SqlKeyword::new("quote"),
@@ -264,10 +264,10 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("regnamespace"),
     SqlKeyword::new("regproc"),
     SqlKeyword::new("regtype"),
+    SqlKeyword::new("reject_limit"),
     SqlKeyword::new("rename"),
     SqlKeyword::new("repeatable"),
     SqlKeyword::new("replace"),
-    SqlKeyword::new("replication"),
     SqlKeyword::new("reset").starts_statement(),
     SqlKeyword::new("restart"),
     SqlKeyword::new("restrict"),
@@ -277,7 +277,6 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("returning"),
     SqlKeyword::new("returns"),
     SqlKeyword::new("revoke").starts_statement(),
-    SqlKeyword::new("rewrite"),
     SqlKeyword::new("right"),
     SqlKeyword::new("role"),
     SqlKeyword::new("rollback").starts_statement(),
@@ -296,7 +295,10 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("session_user"),
     SqlKeyword::new("set").starts_statement(),
     SqlKeyword::new("setof"),
+    SqlKeyword::new("settings"),
+    SqlKeyword::new("share"),
     SqlKeyword::new("show").starts_statement(),
+    SqlKeyword::new("silent"),
     SqlKeyword::new("similar").require_prefix(),
     SqlKeyword::new("skip_locked"),
     SqlKeyword::new("smallint"),
@@ -309,9 +311,12 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("statement"),
     SqlKeyword::new("statistics"),
     SqlKeyword::new("stdin"),
+    SqlKeyword::new("stdout"),
+    SqlKeyword::new("stop"),
     SqlKeyword::new("storage"),
     SqlKeyword::new("stored"),
     SqlKeyword::new("strict"),
+    SqlKeyword::new("summary"),
     SqlKeyword::new("support"),
     SqlKeyword::new("system"),
     SqlKeyword::new("table"),
@@ -325,6 +330,7 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("time"),
     SqlKeyword::new("timestamp"),
     SqlKeyword::new("timestamptz"),
+    SqlKeyword::new("timing"),
     SqlKeyword::new("to"),
     SqlKeyword::new("transaction"),
     SqlKeyword::new("trigger"),
@@ -353,13 +359,15 @@ pub static ALL_KEYWORDS: &[SqlKeyword] = &[
     SqlKeyword::new("version"),
     SqlKeyword::new("view"),
     SqlKeyword::new("volatile"),
+    SqlKeyword::new("wal"),
     SqlKeyword::new("when"),
     SqlKeyword::new("where"),
     SqlKeyword::new("window"),
     SqlKeyword::new("with").starts_statement(),
     SqlKeyword::new("without"),
     SqlKeyword::new("write"),
     SqlKeyword::new("xml"),
+    SqlKeyword::new("yaml"),
     SqlKeyword::new("zone"),
 ];
 
@@ -409,17 +417,92 @@ pub fn complete_keywords<'a>(
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashSet;
+
     use pgls_test_utils::QueryWithCursorPosition;
+    use regex::Regex;
     use sqlx::PgPool;
+    use tree_sitter::Language;
 
     use crate::{
         CompletionItemKind,
+        providers::keywords::ALL_KEYWORDS,
         test_helper::{
             CompletionAssertion, TestCompletionsCase, TestCompletionsSuite,
             assert_complete_results, assert_no_complete_results,
         },
     };
 
+    #[test]
+    fn has_all_keywords_from_treesitter_grammar() {
+        let expected_keywords = keywords_from_treesitter_grammar();
+        let actual_keywords = all_keywords_set();
+
+        let mut missing_keywords = expected_keywords
+            .difference(&actual_keywords)
+            .cloned()
+            .collect::<Vec<_>>();
+        missing_keywords.sort_unstable();
+
+        assert!(
+            missing_keywords.is_empty(),
+            "Found {} keyword(s) from tree-sitter grammar that are missing from ALL_KEYWORDS.\n\
+             Add missing entries to ALL_KEYWORDS in crates/pgls_completions/src/providers/keywords.rs.\n\
+             Missing keywords:\n{}",
+            missing_keywords.len(),
+            missing_keywords.join("\n")
+        );
+    }
+
+    #[test]
+    fn has_no_extra_keywords_outside_treesitter_grammar() {
+        let expected_keywords = keywords_from_treesitter_grammar();
+        let actual_keywords = all_keywords_set();
+
+        let mut extra_keywords = actual_keywords
+            .difference(&expected_keywords)
+            .cloned()
+            .collect::<Vec<_>>();
+        extra_keywords.sort_unstable();
+
+        assert!(
+            extra_keywords.is_empty(),
+            "Found {} keyword(s) in ALL_KEYWORDS that are not tree-sitter grammar keywords.\n\
+             Remove these from ALL_KEYWORDS or add matching keyword_* symbols to grammar.js.\n\
+             Extra keywords:\n{}",
+            extra_keywords.len(),
+            extra_keywords.join("\n")
+        );
+    }
+
+    fn all_keywords_set() -> HashSet<String> {
+        ALL_KEYWORDS
+            .iter()
+            .map(|keyword| keyword.name.to_string())
+            .collect::<HashSet<_>>()
+    }
+
+    fn keywords_from_treesitter_grammar() -> HashSet<String> {
+        let language: Language = pgls_treesitter_grammar::LANGUAGE.into();
+        // Tree-sitter generates auxiliary symbol names like `keyword_int_token2` for keyword
+        // rules that use `choice(...)`. Strip the suffix to compare canonical keyword names.
+        let generated_token_suffix =
+            Regex::new(r"_token\d+$").expect("valid regex for generated token suffix");
+        let mut keywords = HashSet::new();
+
+        for id in 0..language.node_kind_count() {
+            let Some(kind) = language.node_kind_for_id(id as u16) else {
+                continue;
+            };
+            if let Some(keyword) = kind.strip_prefix("keyword_") {
+                let keyword = generated_token_suffix.replace(keyword, "").to_string();
+                keywords.insert(keyword);
+            }
+        }
+
+        keywords
+    }
+
     #[sqlx::test]
     async fn completes_stmt_start_keywords(pool: PgPool) {
         let setup = r#"

crates/pgls_treesitter/src/context/ancestors.rs

@@ -3,7 +3,15 @@ pub struct Scope {
     pub ancestors: AncestorTracker,
 }
 
-static SCOPE_BOUNDARIES: &[&str] = &["statement", "ERROR", "program"];
+static SCOPE_BOUNDARIES: &[&str] = &[
+    "statement",
+    "ERROR",
+    "program",
+    "block",
+    "transaction",
+    "psql_meta_command",
+    "copy_data_stream",
+];
 
 #[derive(Debug)]
 pub struct ScopeTracker {
@@ -22,7 +30,7 @@ impl ScopeTracker {
 
         self.scopes
             .last_mut()
-            .unwrap_or_else(|| panic!("Unhandled node kind: {}", node.kind()))
+            .unwrap_or_else(|| panic!("No top-level grammar-rule found. Please create an issue with the entire Postgres file, noting cursor/hover position."))
             .ancestors
             .register(node, position);
     }
@@ -105,3 +113,104 @@ impl AncestorTracker {
         true
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::context::{TreeSitterContextParams, TreesitterContext};
+
+    fn get_tree(input: &str) -> tree_sitter::Tree {
+        let mut parser = tree_sitter::Parser::new();
+        parser
+            .set_language(&pgls_treesitter_grammar::LANGUAGE.into())
+            .expect("Couldn't set language");
+
+        parser.parse(input, None).expect("Unable to parse tree")
+    }
+
+    fn assert_no_panic_for_all_positions(sql: &str) {
+        let tree = get_tree(sql);
+        for pos in 0..sql.len() {
+            let params = TreeSitterContextParams {
+                position: (pos as u32).into(),
+                text: sql,
+                tree: &tree,
+            };
+            let _ = TreesitterContext::new(params);
+        }
+    }
+
+    #[test]
+    fn scope_boundary_block() {
+        assert_no_panic_for_all_positions("BEGIN; SELECT 1; END;");
+    }
+
+    #[test]
+    fn scope_boundary_transaction() {
+        assert_no_panic_for_all_positions("BEGIN TRANSACTION; SELECT 1; COMMIT;");
+        assert_no_panic_for_all_positions("BEGIN; INSERT INTO t VALUES (1); ROLLBACK;");
+    }
+
+    #[test]
+    fn scope_boundary_psql_meta_command() {
+        assert_no_panic_for_all_positions("\\dt\n\\d users");
+    }
+
+    #[test]
+    fn scope_boundary_copy_data_stream() {
+        assert_no_panic_for_all_positions("COPY t FROM STDIN;\n1\tAlice\n\\.\n");
+    }
+
+    #[test]
+    fn scope_boundary_comment() {
+        assert_no_panic_for_all_positions("-- a comment\nSELECT 1;");
+    }
+
+    #[test]
+    fn issue_704_regression() {
+        let statements = vec![
+            r#"
+            CREATE OR REPLACE FUNCTION my_schema.my_function1(
+                pi_1 character varying, 
+                pi_2 character varying, 
+                pi_3 jsonb, 
+                OUT po_1 integer, 
+                OUT po_2 integer, 
+                OUT result integer
+            )
+            RETURNS record
+            LANGUAGE plpgsql
+            AS $function$
+            "#
+            .trim(),
+
+            r#"
+            CREATE OR REPLACE FUNCTION my_schema.my_function2(
+                pi_1 character varying, 
+                pi_2 character varying, 
+                pi_3 jsonb, 
+                OUT po_1 integer, 
+                OUT po_2 integer, 
+                OUT result integer
+            )
+            RETURNS record
+            LANGUAGE plpgsql
+            AS $function$
+            DECLARE
+            BEGIN
+                -- Function logic goes here
+                -- For example, you can perform some operations using the input parameters and set the output parameters accordingly
+
+                -- Example logic (replace with actual implementation):
+                po_1 := length(pi_1); -- Set po_1 to the length of pi_1
+                po_2 := length(pi_2); -- Set po_2 to the length of pi_2
+                result := po_1 + po_2; -- Set result to the sum of po_1 and po_2
+            END;
+            $function$;
+            "#.trim(),
+        ];
+
+        for stmt in statements {
+            assert_no_panic_for_all_positions(stmt);
+        }
+    }
+}