feats(completions): complete insert, drop/alter table, ignore many situations, improve WHERE

crates/pgt_completions/src/providers/columns.rs

@@ -39,8 +39,8 @@ mod tests {
     use crate::{
         CompletionItem, CompletionItemKind, complete,
         test_helper::{
-            CURSOR_POS, CompletionAssertion, InputQuery, assert_complete_results, get_test_deps,
-            get_test_params,
+            CURSOR_POS, CompletionAssertion, InputQuery, assert_complete_results,
+            assert_no_complete_results, get_test_deps, get_test_params,
         },
     };
 
@@ -573,4 +573,74 @@ mod tests {
         )
         .await;
     }
+
+    #[tokio::test]
+    async fn suggests_columns_in_insert_clause() {
+        let setup = r#"
+            create table instruments (
+                id bigint primary key generated always as identity,
+                name text not null,
+                z text
+            );
+
+            create table others (
+                id serial primary key,
+                a text,
+                b text
+            );
+        "#;
+
+        // We should prefer the instrument columns, even though they
+        // are lower in the alphabet
+
+        assert_complete_results(
+            format!("insert into instruments ({})", CURSOR_POS).as_str(),
+            vec![
+                CompletionAssertion::Label("id".to_string()),
+                CompletionAssertion::Label("name".to_string()),
+                CompletionAssertion::Label("z".to_string()),
+            ],
+            setup,
+        )
+        .await;
+
+        assert_complete_results(
+            format!("insert into instruments (id, {})", CURSOR_POS).as_str(),
+            vec![
+                CompletionAssertion::Label("name".to_string()),
+                CompletionAssertion::Label("z".to_string()),
+            ],
+            setup,
+        )
+        .await;
+
+        assert_complete_results(
+            format!("insert into instruments (id, {}, name)", CURSOR_POS).as_str(),
+            vec![CompletionAssertion::Label("z".to_string())],
+            setup,
+        )
+        .await;
+
+        // works with completed statement
+        assert_complete_results(
+            format!(
+                "insert into instruments (name, {}) values ('my_bass');",
+                CURSOR_POS
+            )
+            .as_str(),
+            vec![
+                CompletionAssertion::Label("id".to_string()),
+                CompletionAssertion::Label("z".to_string()),
+            ],
+            setup,
+        )
+        .await;
+
+        // no completions in the values list!
+        assert_no_complete_results(
+            format!("insert into instruments (id, name) values ({})", CURSOR_POS).as_str(),
+            setup,
+        )
+        .await;
+    }
 }

crates/pgt_completions/src/providers/tables.rs

@@ -310,4 +310,123 @@ mod tests {
         )
         .await;
     }
+
+    #[tokio::test]
+    async fn suggests_tables_in_alter_and_drop_statements() {
+        let setup = r#"
+            create schema auth;
+
+            create table auth.users (
+                uid serial primary key,
+                name text not null,
+                email text unique not null
+            );
+
+            create table auth.posts (
+                pid serial primary key,
+                user_id int not null references auth.users(uid),
+                title text not null,
+                content text,
+                created_at timestamp default now()
+            );
+        "#;
+
+        assert_complete_results(
+            format!("alter table {}", CURSOR_POS).as_str(),
+            vec![
+                CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
+                CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
+            ],
+            setup,
+        )
+        .await;
+
+        assert_complete_results(
+            format!("alter table if exists {}", CURSOR_POS).as_str(),
+            vec![
+                CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
+                CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
+            ],
+            setup,
+        )
+        .await;
+
+        assert_complete_results(
+            format!("drop table {}", CURSOR_POS).as_str(),
+            vec![
+                CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table),
+                CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
+            ],
+            setup,
+        )
+        .await;
+
+        assert_complete_results(
+            format!("drop table if exists {}", CURSOR_POS).as_str(),
+            vec![
+                CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("posts".into(), CompletionItemKind::Table), // self-join
+                CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
+            ],
+            setup,
+        )
+        .await;
+    }
+
+    #[tokio::test]
+    async fn suggests_tables_in_insert_into() {
+        let setup = r#"
+            create schema auth;
+
+            create table auth.users (
+                uid serial primary key,
+                name text not null,
+                email text unique not null
+            );
+        "#;
+
+        assert_complete_results(
+            format!("insert into {}", CURSOR_POS).as_str(),
+            vec![
+                CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
+            ],
+            setup,
+        )
+        .await;
+
+        assert_complete_results(
+            format!("insert into auth.{}", CURSOR_POS).as_str(),
+            vec![CompletionAssertion::LabelAndKind(
+                "users".into(),
+                CompletionItemKind::Table,
+            )],
+            setup,
+        )
+        .await;
+
+        // works with complete statement.
+        assert_complete_results(
+            format!(
+                "insert into {} (name, email) values ('jules', 'a@b.com');",
+                CURSOR_POS
+            )
+            .as_str(),
+            vec![
+                CompletionAssertion::LabelAndKind("public".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("auth".into(), CompletionItemKind::Schema),
+                CompletionAssertion::LabelAndKind("users".into(), CompletionItemKind::Table),
+            ],
+            setup,
+        )
+        .await;
+    }
 }

crates/pgt_completions/src/relevance/filtering.rs

@@ -1,4 +1,4 @@
-use crate::context::{CompletionContext, NodeUnderCursor, WrappingClause};
+use crate::context::{CompletionContext, NodeUnderCursor, WrappingClause, WrappingNode};
 
 use super::CompletionRelevanceData;
 
@@ -24,6 +24,10 @@ impl CompletionFilter<'_> {
     }
 
     fn completable_context(&self, ctx: &CompletionContext) -> Option<()> {
+        if ctx.wrapping_node_kind.is_none() && ctx.wrapping_clause_type.is_none() {
+            return None;
+        }
+
         let current_node_kind = ctx
             .node_under_cursor
             .as_ref()
@@ -65,55 +69,99 @@ impl CompletionFilter<'_> {
     }
 
     fn check_clause(&self, ctx: &CompletionContext) -> Option<()> {
-        let clause = ctx.wrapping_clause_type.as_ref();
-
-        let in_clause = |compare: WrappingClause| clause.is_some_and(|c| c == &compare);
-
-        match self.data {
-            CompletionRelevanceData::Table(_) => {
-                if in_clause(WrappingClause::Select)
-                    || in_clause(WrappingClause::Where)
-                    || in_clause(WrappingClause::PolicyName)
-                {
-                    return None;
-                };
-            }
-            CompletionRelevanceData::Column(_) => {
-                if in_clause(WrappingClause::From) || in_clause(WrappingClause::PolicyName) {
-                    return None;
-                }
-
-                // We can complete columns in JOIN cluases, but only if we are after the
-                // ON node in the "ON u.id = posts.user_id" part.
-                let in_join_clause_before_on_node = clause.is_some_and(|c| match c {
-                    // we are in a JOIN, but definitely not after an ON
-                    WrappingClause::Join { on_node: None } => true,
-
-                    WrappingClause::Join { on_node: Some(on) } => ctx
-                        .node_under_cursor
-                        .as_ref()
-                        .is_some_and(|n| n.end_byte() < on.start_byte()),
-
-                    _ => false,
-                });
-
-                if in_join_clause_before_on_node {
-                    return None;
-                }
-            }
-            CompletionRelevanceData::Policy(_) => {
-                if clause.is_none_or(|c| c != &WrappingClause::PolicyName) {
-                    return None;
-                }
-            }
-            _ => {
-                if in_clause(WrappingClause::PolicyName) {
-                    return None;
+        ctx.wrapping_clause_type
+            .as_ref()
+            .map(|clause| {
+                match self.data {
+                    CompletionRelevanceData::Table(_) => match clause {
+                        WrappingClause::Select
+                        | WrappingClause::Where
+                        | WrappingClause::ColumnDefinitions => false,
+
+                        WrappingClause::Insert => {
+                            ctx.wrapping_node_kind
+                                .as_ref()
+                                .is_none_or(|n| n != &WrappingNode::List)
+                                && (ctx.before_cursor_matches_kind(&["keyword_into"])
+                                    || (ctx.before_cursor_matches_kind(&["."])
+                                        && ctx.parent_matches_one_of_kind(&["object_reference"])))
+                        }
+
+                        WrappingClause::DropTable | WrappingClause::AlterTable => ctx
+                            .before_cursor_matches_kind(&[
+                                "keyword_exists",
+                                "keyword_only",
+                                "keyword_table",
+                            ]),
+
+                        _ => true,
+                    },
+
+                    CompletionRelevanceData::Column(_) => {
+                        match clause {
+                            WrappingClause::From
+                            | WrappingClause::ColumnDefinitions
+                            | WrappingClause::AlterTable
+                            | WrappingClause::DropTable => false,
+
+                            // We can complete columns in JOIN cluases, but only if we are after the
+                            // ON node in the "ON u.id = posts.user_id" part.
+                            WrappingClause::Join { on_node: Some(on) } => ctx
+                                .node_under_cursor
+                                .as_ref()
+                                .is_some_and(|cn| cn.start_byte() >= on.end_byte()),
+
+                            // we are in a JOIN, but definitely not after an ON
+                            WrappingClause::Join { on_node: None } => false,
+
+                            WrappingClause::Insert => ctx
+                                .wrapping_node_kind
+                                .as_ref()
+                                .is_some_and(|n| n == &WrappingNode::List),
+
+                            _ => true,
+                        }
+                    }
+
+                    CompletionRelevanceData::Function(_) => matches!(
+                        clause,
+                        WrappingClause::From
+                            | WrappingClause::Select
+                            | WrappingClause::Where
+                            | WrappingClause::Join { .. }
+                    ),
+
+                    CompletionRelevanceData::Schema(_) => match clause {
+                        WrappingClause::Select
+                        | WrappingClause::Where
+                        | WrappingClause::From
+                        | WrappingClause::Join { .. }
+                        | WrappingClause::Update
+                        | WrappingClause::Delete => true,
+
+                        WrappingClause::DropTable | WrappingClause::AlterTable => ctx
+                            .before_cursor_matches_kind(&[
+                                "keyword_exists",
+                                "keyword_only",
+                                "keyword_table",
+                            ]),
+
+                        WrappingClause::Insert => {
+                            ctx.wrapping_node_kind
+                                .as_ref()
+                                .is_none_or(|n| n != &WrappingNode::List)
+                                && ctx.before_cursor_matches_kind(&["keyword_into"])
+                        }
+
+                        _ => false,
+                    },
+
+                    CompletionRelevanceData::Policy(_) => {
+                        matches!(clause, WrappingClause::PolicyName)
+                    }
                 }
-            }
-        }
-
-        Some(())
+            })
+            .and_then(|is_ok| if is_ok { Some(()) } else { None })
     }
 
     fn check_invocation(&self, ctx: &CompletionContext) -> Option<()> {
@@ -188,4 +236,15 @@ mod tests {
         )
         .await;
     }
+
+    #[tokio::test]
+    async fn completion_after_create_table() {
+        assert_no_complete_results(format!("create table {}", CURSOR_POS).as_str(), "").await;
+    }
+
+    #[tokio::test]
+    async fn completion_in_column_definitions() {
+        let query = format!(r#"create table instruments ( {} )"#, CURSOR_POS);
+        assert_no_complete_results(query.as_str(), "").await;
+    }
 }