test: update, insert, copy completions

crates/pgls_completions/src/complete.rs

@@ -54,3 +54,160 @@ pub fn complete(params: CompletionParams) -> Vec<CompletionItem> {
 
     builder.finish()
 }
+
+#[cfg(test)]
+mod tests {
+
+    use sqlx::PgPool;
+
+    use crate::test_helper::{TestCompletionsCase, TestCompletionsSuite};
+
+    #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")]
+    async fn completions_in_update_statements(pool: PgPool) {
+        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
+            );
+        "#;
+
+        TestCompletionsSuite::new(&pool, Some(setup))
+            .with_case(
+                TestCompletionsCase::new()
+                    .type_sql("update instruments set name = 'new' where id = 1;"),
+            )
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement("update instruments as i <sql> where i.id = 1;")
+                    .type_sql("set name = 'x', z = 'y'"),
+            )
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement("update instruments set name = 'x' <sql>")
+                    .type_sql("returning id, name;"),
+            )
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement("update instruments <sql> where id = 1;")
+                    .type_sql("set name = default, z = 'y'"),
+            )
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement("update instruments set name = o.a <sql>")
+                    .type_sql("from others o where instruments.id = o.id returning name, z;"),
+            )
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement("update instruments set name = (<sql>) where id = 1;")
+                    .type_sql("select a from others where id = 1"),
+            )
+            .snapshot("completions_in_update_statements")
+            .await;
+    }
+
+    #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")]
+    async fn completions_in_insert_statements(pool: PgPool) {
+        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
+            );
+        "#;
+
+        TestCompletionsSuite::new(&pool, Some(setup))
+            // basic VALUES (full statement typed once)
+            .with_case(
+                TestCompletionsCase::new()
+                    .type_sql("insert into instruments (id, name) values (1, 'bass');"),
+            )
+            // RETURNING clause
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement(
+                        "insert into instruments (id, name) values (1, 'x') <sql>",
+                    )
+                    .type_sql("returning id, name;"),
+            )
+            // multi-row VALUES with DEFAULT
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement("insert into instruments (id, name, z) <sql>")
+                    .type_sql("values (1, 'a', 'b'), (2, 'c', default);"),
+            )
+            // INSERT SELECT
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement("insert into instruments (id, name) <sql>")
+                    .type_sql("select id, a from others;"),
+            )
+            // schema-qualified table
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement("insert into <sql> (id, name) values (1, 'x');")
+                    .type_sql("public.instruments"),
+            )
+            .snapshot("completions_in_insert_statements")
+            .await;
+    }
+
+    #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")]
+    async fn completions_in_copy_statements(pool: PgPool) {
+        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
+            );
+        "#;
+
+        TestCompletionsSuite::new(&pool, Some(setup))
+            .with_case(TestCompletionsCase::new().type_sql("copy instruments to stdout;"))
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement("copy instruments <sql>")
+                    .type_sql("from stdin;"),
+            )
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement("copy instruments (<sql>) to stdout;")
+                    .type_sql("id, name"),
+            )
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement("copy instruments to stdout <sql>")
+                    .type_sql("with (format csv, header);"),
+            )
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement("copy instruments from stdin <sql>")
+                    .type_sql("with (format csv, header);"),
+            )
+            .with_case(
+                TestCompletionsCase::new()
+                    .inside_static_statement("copy (<sql>) to stdout;")
+                    .type_sql("select * from instruments where id > 0"),
+            )
+            .snapshot("completions_in_copy_statements")
+            .await;
+    }
+}

crates/pgls_completions/src/providers/columns.rs

@@ -51,9 +51,12 @@ fn get_completion_text(ctx: &TreesitterContext, col: &Column) -> CompletionText
 #[cfg(test)]
 mod tests {
 
+    use pgls_test_utils::QueryWithCursorPosition;
     use sqlx::PgPool;
 
-    use crate::test_helper::{TestCompletionsCase, TestCompletionsSuite};
+    use crate::test_helper::{
+        CompletionAssertion, TestCompletionsCase, TestCompletionsSuite, assert_complete_results,
+    };
 
     #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")]
     async fn handles_nested_queries(pool: PgPool) {
@@ -597,4 +600,41 @@ mod tests {
             .snapshot("completes_quoted_columns_with_aliases")
             .await;
     }
+
+    #[sqlx::test(migrator = "pgls_test_utils::MIGRATIONS")]
+    async fn suggest_column_names_in_returning(pool: PgPool) {
+        let setup = r#"
+            create schema if not exists private;
+
+            create table private.users (
+                id serial primary key,
+                email text unique not null,
+                name text not null
+            );
+        "#;
+
+        let query = format!(
+            r#"insert into private.users (email, name) values ('test@ema.il', 'hello') returning {}"#,
+            QueryWithCursorPosition::cursor_marker()
+        );
+
+        assert_complete_results(
+            query.as_str(),
+            vec![
+                CompletionAssertion::LabelAndKind(
+                    "email".into(),
+                    crate::CompletionItemKind::Column,
+                ),
+                CompletionAssertion::LabelAndKind("id".into(), crate::CompletionItemKind::Column),
+                CompletionAssertion::LabelAndKind("name".into(), crate::CompletionItemKind::Column),
+                CompletionAssertion::LabelAndKind(
+                    "with".into(),
+                    crate::CompletionItemKind::Keyword,
+                ),
+            ],
+            Some(setup),
+            &pool,
+        )
+        .await;
+    }
 }