diff --git a/crates/pgls_completions/src/complete.rs b/crates/pgls_completions/src/complete.rs index 617cdc91b..b75334ba4 100644 --- a/crates/pgls_completions/src/complete.rs +++ b/crates/pgls_completions/src/complete.rs @@ -54,3 +54,160 @@ pub fn complete(params: CompletionParams) -> Vec { 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 where i.id = 1;") + .type_sql("set name = 'x', z = 'y'"), + ) + .with_case( + TestCompletionsCase::new() + .inside_static_statement("update instruments set name = 'x' ") + .type_sql("returning id, name;"), + ) + .with_case( + TestCompletionsCase::new() + .inside_static_statement("update instruments where id = 1;") + .type_sql("set name = default, z = 'y'"), + ) + .with_case( + TestCompletionsCase::new() + .inside_static_statement("update instruments set name = o.a ") + .type_sql("from others o where instruments.id = o.id returning name, z;"), + ) + .with_case( + TestCompletionsCase::new() + .inside_static_statement("update instruments set name = () 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') ", + ) + .type_sql("returning id, name;"), + ) + // multi-row VALUES with DEFAULT + .with_case( + TestCompletionsCase::new() + .inside_static_statement("insert into instruments (id, name, z) ") + .type_sql("values (1, 'a', 'b'), (2, 'c', default);"), + ) + // INSERT SELECT + .with_case( + TestCompletionsCase::new() + .inside_static_statement("insert into instruments (id, name) ") + .type_sql("select id, a from others;"), + ) + // schema-qualified table + .with_case( + TestCompletionsCase::new() + .inside_static_statement("insert into (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 ") + .type_sql("from stdin;"), + ) + .with_case( + TestCompletionsCase::new() + .inside_static_statement("copy instruments () to stdout;") + .type_sql("id, name"), + ) + .with_case( + TestCompletionsCase::new() + .inside_static_statement("copy instruments to stdout ") + .type_sql("with (format csv, header);"), + ) + .with_case( + TestCompletionsCase::new() + .inside_static_statement("copy instruments from stdin ") + .type_sql("with (format csv, header);"), + ) + .with_case( + TestCompletionsCase::new() + .inside_static_statement("copy () to stdout;") + .type_sql("select * from instruments where id > 0"), + ) + .snapshot("completions_in_copy_statements") + .await; + } +} diff --git a/crates/pgls_completions/src/providers/columns.rs b/crates/pgls_completions/src/providers/columns.rs index d0908d0e0..b3b842f7a 100644 --- a/crates/pgls_completions/src/providers/columns.rs +++ b/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; + } } diff --git a/crates/pgls_completions/src/snapshots/pgls_completions__test_helper__completions_in_copy_statements.snap b/crates/pgls_completions/src/snapshots/pgls_completions__test_helper__completions_in_copy_statements.snap new file mode 100644 index 000000000..c64bbeed9 --- /dev/null +++ b/crates/pgls_completions/src/snapshots/pgls_completions__test_helper__completions_in_copy_statements.snap @@ -0,0 +1,406 @@ +--- +source: crates/pgls_completions/src/test_helper.rs +expression: final_snapshot +--- +***Setup*** + +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 +); + + +-------------- + +***Case 1:*** + +c| + +Results: +copy - copy (Keyword) +create - create (Keyword) +select - select (Keyword) +truncate - truncate (Keyword) + +-------------- + +copy | + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +_sqlx_migrations - public._sqlx_migrations (Table) +information_schema - information_schema (Schema) +pg_catalog - pg_catalog (Schema) + +-------------- + +copy i| + +Results: +instruments - public.instruments (Table) +information_schema - information_schema (Schema) +information_schema.information_schema_catalog_name - information_schema.information_schema_catalog_name (Table) +public - public (Schema) +pg_catalog.pg_ident_file_mappings - pg_catalog.pg_ident_file_mappings (Table) + +-------------- + +copy instruments | + +Results: +from - from (Keyword) +to - to (Keyword) + +-------------- + +copy instruments t| + +Results: +to - to (Keyword) + +-------------- + +copy instruments to | + +Results: +program - program (Keyword) +stdin - stdin (Keyword) +stdout - stdout (Keyword) + +-------------- + +copy instruments to s| + +Results: +stdin - stdin (Keyword) +stdout - stdout (Keyword) + +-------------- + +copy instruments to stdout; | + +Results: +copy - copy (Keyword) +create - create (Keyword) +drop - drop (Keyword) +insert - insert (Keyword) +reset - reset (Keyword) + +-------------- + + + + + +***Case 2:*** + +copy instruments | + +Results: +from - from (Keyword) +to - to (Keyword) + +-------------- + +copy instruments f| + +Results: +from - from (Keyword) + +-------------- + +copy instruments from | + +Results: +program - program (Keyword) +stdin - stdin (Keyword) +stdout - stdout (Keyword) + +-------------- + +copy instruments from s| + +Results: +stdin - stdin (Keyword) +stdout - stdout (Keyword) + +-------------- + +copy instruments from stdin; | + + + + +***Case 3:*** + +copy instruments (|) to stdout; + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +a - public.others.a (Column) +b - public.others.b (Column) + +-------------- + +copy instruments (i|) to stdout; + +Results: +id - public.instruments.id (Column) +ident - pg_catalog.pg_backend_memory_contexts.ident (Column) +identity_cycle - information_schema.columns.identity_cycle (Column) +identity_generation - information_schema.columns.identity_generation (Column) +identity_increment - information_schema.columns.identity_increment (Column) + +-------------- + +copy instruments (id, |) to stdout; + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +a - public.others.a (Column) +b - public.others.b (Column) + +-------------- + +copy instruments (id, n|) to stdout; + +Results: +name - public.instruments.name (Column) +n_dead_tup - pg_catalog.pg_stat_all_tables.n_dead_tup (Column) +n_distinct - pg_catalog.pg_stats.n_distinct (Column) +n_ins_since_vacuum - pg_catalog.pg_stat_all_tables.n_ins_since_vacuum (Column) +n_live_tup - pg_catalog.pg_stat_all_tables.n_live_tup (Column) + +-------------- + +copy instruments (id, name |) to stdout; + + + + +***Case 4:*** + +copy instruments to stdout | + +Results: +binary - binary (Keyword) +csv - csv (Keyword) +delimiter - delimiter (Keyword) +with - with (Keyword) + +-------------- + +copy instruments to stdout w| + +Results: +with - with (Keyword) + +-------------- + +copy instruments to stdout with | + +Results: +binary - binary (Keyword) +csv - csv (Keyword) +delimiter - delimiter (Keyword) + +-------------- + +copy instruments to stdout with (| +copy instruments to stdout with (|) + +Results: +default - default (Keyword) +delimiter - delimiter (Keyword) +encoding - encoding (Keyword) +escape - escape (Keyword) +force_not_null - force_not_null (Keyword) + +-------------- + +copy instruments to stdout with (f|) + +Results: +force_not_null - force_not_null (Keyword) +force_null - force_null (Keyword) +force_quote - force_quote (Keyword) +format - format (Keyword) +freeze - freeze (Keyword) + +-------------- + +copy instruments to stdout with (format |) +copy instruments to stdout with (format c|) + +Results: +csv - csv (Keyword) + +-------------- + +copy instruments to stdout with (format csv, |) +copy instruments to stdout with (format csv, h|) +copy instruments to stdout with (format csv, header); | + +Results: +copy - copy (Keyword) +create - create (Keyword) +drop - drop (Keyword) +insert - insert (Keyword) +reset - reset (Keyword) + +-------------- + + + + + +***Case 5:*** + +copy instruments from stdin | + +Results: +binary - binary (Keyword) +csv - csv (Keyword) +delimiter - delimiter (Keyword) +where - where (Keyword) +with - with (Keyword) + +-------------- + +copy instruments from stdin w| + +Results: +where - where (Keyword) +with - with (Keyword) + +-------------- + +copy instruments from stdin with | + +Results: +binary - binary (Keyword) +csv - csv (Keyword) +delimiter - delimiter (Keyword) + +-------------- + +copy instruments from stdin with (| +copy instruments from stdin with (|) + +Results: +default - default (Keyword) +delimiter - delimiter (Keyword) +encoding - encoding (Keyword) +escape - escape (Keyword) +force_not_null - force_not_null (Keyword) + +-------------- + +copy instruments from stdin with (f|) + +Results: +force_not_null - force_not_null (Keyword) +force_null - force_null (Keyword) +force_quote - force_quote (Keyword) +format - format (Keyword) +freeze - freeze (Keyword) + +-------------- + +copy instruments from stdin with (format |) +copy instruments from stdin with (format c|) + +Results: +csv - csv (Keyword) + +-------------- + +copy instruments from stdin with (format csv, |) +copy instruments from stdin with (format csv, h|) +copy instruments from stdin with (format csv, header); | + + + + +***Case 6:*** + +copy (|) to stdout; + +Results: +select - select (Keyword) +table - table (Keyword) +tables - tables (Keyword) +tablespace - tablespace (Keyword) +values - values (Keyword) + +-------------- + +copy (s|) to stdout; + +Results: +select - select (Keyword) +tables - tables (Keyword) +tablespace - tablespace (Keyword) +values - values (Keyword) + +-------------- + +copy (select |) to stdout; +copy (select * |) to stdout; +copy (select * f|) to stdout; + +Results: +for - for (Keyword) +from - from (Keyword) +offset - offset (Keyword) + +-------------- + +copy (select * from |) to stdout; +copy (select * from i|) to stdout; + +Results: +instruments - public.instruments (Table) +information_schema - information_schema (Schema) +public - public (Schema) +information_schema.information_schema_catalog_name - information_schema.information_schema_catalog_name (Table) +pg_catalog.pg_ident_file_mappings - pg_catalog.pg_ident_file_mappings (Table) + +-------------- + +copy (select * from instruments |) to stdout; +copy (select * from instruments w|) to stdout; + +Results: +where - where (Keyword) +window - window (Keyword) + +-------------- + +copy (select * from instruments where |) to stdout; +copy (select * from instruments where i|) to stdout; + +Results: +id - public.instruments.id (Column) +ident - pg_catalog.pg_backend_memory_contexts.ident (Column) +identity_cycle - information_schema.columns.identity_cycle (Column) +identity_generation - information_schema.columns.identity_generation (Column) +identity_increment - information_schema.columns.identity_increment (Column) + +-------------- + +copy (select * from instruments where id |) to stdout; +copy (select * from instruments where id > |) to stdout; +copy (select * from instruments where id > 0 |) to stdout; diff --git a/crates/pgls_completions/src/snapshots/pgls_completions__test_helper__completions_in_insert_statements.snap b/crates/pgls_completions/src/snapshots/pgls_completions__test_helper__completions_in_insert_statements.snap new file mode 100644 index 000000000..4f314d83e --- /dev/null +++ b/crates/pgls_completions/src/snapshots/pgls_completions__test_helper__completions_in_insert_statements.snap @@ -0,0 +1,523 @@ +--- +source: crates/pgls_completions/src/test_helper.rs +expression: final_snapshot +--- +***Setup*** + +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 +); + + +-------------- + +***Case 1:*** + +i| + +Results: +insert - insert (Keyword) +with - with (Keyword) + +-------------- + +insert | + +Results: +into - into (Keyword) + +-------------- + +insert i| + +Results: +into - into (Keyword) + +-------------- + +insert into | + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +_sqlx_migrations - public._sqlx_migrations (Table) +information_schema - information_schema (Schema) +pg_catalog - pg_catalog (Schema) + +-------------- + +insert into i| + +Results: +instruments - public.instruments (Table) +information_schema - information_schema (Schema) +information_schema.information_schema_catalog_name - information_schema.information_schema_catalog_name (Table) +public - public (Schema) +pg_catalog.pg_ident_file_mappings - pg_catalog.pg_ident_file_mappings (Table) + +-------------- + +insert into instruments | + +Results: +default - default (Keyword) +select - select (Keyword) +values - values (Keyword) + +-------------- + +insert into instruments (| +insert into instruments (|) + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +select - select (Keyword) +a - public.others.a (Column) + +-------------- + +insert into instruments (i|) + +Results: +id - public.instruments.id (Column) +ident - pg_catalog.pg_backend_memory_contexts.ident (Column) +identity_cycle - information_schema.columns.identity_cycle (Column) +identity_generation - information_schema.columns.identity_generation (Column) +identity_increment - information_schema.columns.identity_increment (Column) + +-------------- + +insert into instruments (id, |) + +Results: +name - public.instruments.name (Column) +z - public.instruments.z (Column) +id - public.instruments.id (Column) +a - public.others.a (Column) +b - public.others.b (Column) + +-------------- + +insert into instruments (id, n|) + +Results: +name - public.instruments.name (Column) +n_dead_tup - pg_catalog.pg_stat_all_tables.n_dead_tup (Column) +n_distinct - pg_catalog.pg_stats.n_distinct (Column) +n_ins_since_vacuum - pg_catalog.pg_stat_all_tables.n_ins_since_vacuum (Column) +n_live_tup - pg_catalog.pg_stat_all_tables.n_live_tup (Column) + +-------------- + +insert into instruments (id, name) | + +Results: +default - default (Keyword) +select - select (Keyword) +values - values (Keyword) + +-------------- + +insert into instruments (id, name) v| + +Results: +values - values (Keyword) + +-------------- + +insert into instruments (id, name) values | +insert into instruments (id, name) values (| +insert into instruments (id, name) values (|) + +Results: +z - public.instruments.z (Column) +id - public.instruments.id (Column) +name - public.instruments.name (Column) +default - default (Keyword) +a - public.others.a (Column) + +-------------- + +insert into instruments (id, name) values (1|) +insert into instruments (id, name) values (1, |) + +Results: +z - public.instruments.z (Column) +id - public.instruments.id (Column) +name - public.instruments.name (Column) +default - default (Keyword) +a - public.others.a (Column) + +-------------- + +insert into instruments (id, name) values (1, '|) +insert into instruments (id, name) values (1, 'bass'); | + +Results: +copy - copy (Keyword) +create - create (Keyword) +drop - drop (Keyword) +insert - insert (Keyword) +reset - reset (Keyword) + +-------------- + + + + + +***Case 2:*** + +insert into instruments (id, name) values (1, 'x') | + +Results: +returning - returning (Keyword) + +-------------- + +insert into instruments (id, name) values (1, 'x') r| + +Results: +returning - returning (Keyword) + +-------------- + +insert into instruments (id, name) values (1, 'x') returning | + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +with - with (Keyword) +a - public.others.a (Column) + +-------------- + +insert into instruments (id, name) values (1, 'x') returning i| + +Results: +id - public.instruments.id (Column) +interval - interval (Keyword) +id - public.others.id (Column) +with - with (Keyword) +ident - pg_catalog.pg_backend_memory_contexts.ident (Column) + +-------------- + +insert into instruments (id, name) values (1, 'x') returning id, | + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +a - public.others.a (Column) +b - public.others.b (Column) + +-------------- + +insert into instruments (id, name) values (1, 'x') returning id, n| + +Results: +name - public.instruments.name (Column) +null - null (Keyword) +n_dead_tup - pg_catalog.pg_stat_all_tables.n_dead_tup (Column) +n_distinct - pg_catalog.pg_stats.n_distinct (Column) +n_ins_since_vacuum - pg_catalog.pg_stat_all_tables.n_ins_since_vacuum (Column) + +-------------- + +insert into instruments (id, name) values (1, 'x') returning id, name; | + +Results: +copy - copy (Keyword) +create - create (Keyword) +drop - drop (Keyword) +insert - insert (Keyword) +reset - reset (Keyword) + +-------------- + + + + + +***Case 3:*** + +insert into instruments (id, name, z) | + +Results: +default - default (Keyword) +select - select (Keyword) +values - values (Keyword) + +-------------- + +insert into instruments (id, name, z) v| + +Results: +values - values (Keyword) + +-------------- + +insert into instruments (id, name, z) values | +insert into instruments (id, name, z) values (| +insert into instruments (id, name, z) values (|) + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +default - default (Keyword) +a - public.others.a (Column) + +-------------- + +insert into instruments (id, name, z) values (1|) +insert into instruments (id, name, z) values (1, |) + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +default - default (Keyword) +a - public.others.a (Column) + +-------------- + +insert into instruments (id, name, z) values (1, '|) +insert into instruments (id, name, z) values (1, 'a', |) + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +default - default (Keyword) +a - public.others.a (Column) + +-------------- + +insert into instruments (id, name, z) values (1, 'a', '|) +insert into instruments (id, name, z) values (1, 'a', 'b'), |) +insert into instruments (id, name, z) values (1, 'a', 'b'), (| +insert into instruments (id, name, z) values (1, 'a', 'b'), (|) + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +default - default (Keyword) +a - public.others.a (Column) + +-------------- + +insert into instruments (id, name, z) values (1, 'a', 'b'), (2|) +insert into instruments (id, name, z) values (1, 'a', 'b'), (2, |) + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +default - default (Keyword) +a - public.others.a (Column) + +-------------- + +insert into instruments (id, name, z) values (1, 'a', 'b'), (2, '|) +insert into instruments (id, name, z) values (1, 'a', 'b'), (2, 'c', |) + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +default - default (Keyword) +a - public.others.a (Column) + +-------------- + +insert into instruments (id, name, z) values (1, 'a', 'b'), (2, 'c', d|) + +Results: +default - default (Keyword) +id - public.instruments.id (Column) +data - pg_catalog.pg_largeobject.data (Column) +data_type - pg_catalog.pg_sequences.data_type (Column) +database - pg_catalog.pg_hba_file_rules.database (Column) + +-------------- + +insert into instruments (id, name, z) values (1, 'a', 'b'), (2, 'c', default); | + +Results: +copy - copy (Keyword) +create - create (Keyword) +drop - drop (Keyword) +insert - insert (Keyword) +reset - reset (Keyword) + +-------------- + + + + + +***Case 4:*** + +insert into instruments (id, name) | + +Results: +default - default (Keyword) +select - select (Keyword) +values - values (Keyword) + +-------------- + +insert into instruments (id, name) s| + +Results: +select - select (Keyword) +values - values (Keyword) + +-------------- + +insert into instruments (id, name) select | + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +a - public.others.a (Column) +b - public.others.b (Column) + +-------------- + +insert into instruments (id, name) select i| + +Results: +id - public.instruments.id (Column) +ident - pg_catalog.pg_backend_memory_contexts.ident (Column) +identity_cycle - information_schema.columns.identity_cycle (Column) +identity_generation - information_schema.columns.identity_generation (Column) +identity_increment - information_schema.columns.identity_increment (Column) + +-------------- + +insert into instruments (id, name) select id, | + +Results: +name - public.instruments.name (Column) +z - public.instruments.z (Column) +id - public.instruments.id (Column) +a - public.others.a (Column) +b - public.others.b (Column) + +-------------- + +insert into instruments (id, name) select id, a | + +Results: +for - for (Keyword) +from - from (Keyword) +group - group (Keyword) +limit - limit (Keyword) +offset - offset (Keyword) + +-------------- + +insert into instruments (id, name) select id, a f| + +Results: +for - for (Keyword) +from - from (Keyword) +offset - offset (Keyword) + +-------------- + +insert into instruments (id, name) select id, a from | + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +information_schema - information_schema (Schema) +pg_catalog - pg_catalog (Schema) +pg_toast - pg_toast (Schema) + +-------------- + +insert into instruments (id, name) select id, a from o| + +Results: +others - public.others (Table) +only - only (Keyword) +information_schema - information_schema (Schema) +pg_catalog - pg_catalog (Schema) +pg_toast - pg_toast (Schema) + +-------------- + +insert into instruments (id, name) select id, a from others; | + +Results: +copy - copy (Keyword) +create - create (Keyword) +drop - drop (Keyword) +insert - insert (Keyword) +reset - reset (Keyword) + +-------------- + + + + + +***Case 5:*** + +insert into | (id, name) values (1, 'x'); + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +_sqlx_migrations - public._sqlx_migrations (Table) +information_schema - information_schema (Schema) +pg_catalog - pg_catalog (Schema) + +-------------- + +insert into p| (id, name) values (1, 'x'); + +Results: +public - public (Schema) +information_schema.parameters - information_schema.parameters (Table) +pg_catalog.pg_aggregate - pg_catalog.pg_aggregate (Table) +pg_catalog.pg_am - pg_catalog.pg_am (Table) +pg_catalog.pg_amop - pg_catalog.pg_amop (Table) + +-------------- + +insert into public.| (id, name) values (1, 'x'); + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +_sqlx_migrations - public._sqlx_migrations (Table) + +-------------- + +insert into public.i| (id, name) values (1, 'x'); + +Results: +instruments - public.instruments (Table) +_sqlx_migrations - public._sqlx_migrations (Table) + +-------------- + +insert into public.instruments | (id, name) values (1, 'x'); diff --git a/crates/pgls_completions/src/snapshots/pgls_completions__test_helper__completions_in_update_statements.snap.new b/crates/pgls_completions/src/snapshots/pgls_completions__test_helper__completions_in_update_statements.snap.new new file mode 100644 index 000000000..2f6375257 --- /dev/null +++ b/crates/pgls_completions/src/snapshots/pgls_completions__test_helper__completions_in_update_statements.snap.new @@ -0,0 +1,719 @@ +--- +source: crates/pgls_completions/src/test_helper.rs +assertion_line: 646 +expression: final_snapshot +--- +***Setup*** + +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 +); + + +-------------- + +***Case 1:*** + +u| + +Results: +update - update (Keyword) +truncate - truncate (Keyword) + +-------------- + +update | + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +_sqlx_migrations - public._sqlx_migrations (Table) +information_schema - information_schema (Schema) +pg_catalog - pg_catalog (Schema) + +-------------- + +update i| + +Results: +instruments - public.instruments (Table) +information_schema - information_schema (Schema) +public - public (Schema) +information_schema.information_schema_catalog_name - information_schema.information_schema_catalog_name (Table) +id - public.instruments.id (Column) + +-------------- + +update instruments | + +Results: +set - set (Keyword) + +-------------- + +update instruments s| + +Results: +set - set (Keyword) + +-------------- + +update instruments set | + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +public - public (Schema) +information_schema - information_schema (Schema) + +-------------- + +update instruments set n| + +Results: +name - public.instruments.name (Column) +information_schema - information_schema (Schema) +n_dead_tup - pg_catalog.pg_stat_all_tables.n_dead_tup (Column) +n_distinct - pg_catalog.pg_stats.n_distinct (Column) +n_ins_since_vacuum - pg_catalog.pg_stat_all_tables.n_ins_since_vacuum (Column) + +-------------- + +update instruments set name | +update instruments set name = | + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) + +-------------- + +update instruments set name = '| +update instruments set name = 'new' | + +Results: +from - from (Keyword) +where - where (Keyword) + +-------------- + +update instruments set name = 'new' w| + +Results: +where - where (Keyword) + +-------------- + +update instruments set name = 'new' where | + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +a - public.others.a (Column) +b - public.others.b (Column) + +-------------- + +update instruments set name = 'new' where i| + +Results: +id - public.instruments.id (Column) +ident - pg_catalog.pg_backend_memory_contexts.ident (Column) +identity_cycle - information_schema.columns.identity_cycle (Column) +identity_generation - information_schema.columns.identity_generation (Column) +identity_increment - information_schema.columns.identity_increment (Column) + +-------------- + +update instruments set name = 'new' where id | +update instruments set name = 'new' where id = | + +Results: +name - public.instruments.name (Column) +z - public.instruments.z (Column) +id - public.instruments.id (Column) +a - public.others.a (Column) +b - public.others.b (Column) + +-------------- + +update instruments set name = 'new' where id = 1| +update instruments set name = 'new' where id = 1; | + +Results: +copy - copy (Keyword) +create - create (Keyword) +drop - drop (Keyword) +insert - insert (Keyword) +reset - reset (Keyword) + +-------------- + + + + + +***Case 2:*** + +update instruments as i | where i.id = 1; + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) + +-------------- + +update instruments as i s| where i.id = 1; +update instruments as i set | where i.id = 1; + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +public - public (Schema) +information_schema - information_schema (Schema) + +-------------- + +update instruments as i set n| where i.id = 1; + +Results: +name - public.instruments.name (Column) +information_schema - information_schema (Schema) +n_dead_tup - pg_catalog.pg_stat_all_tables.n_dead_tup (Column) +n_distinct - pg_catalog.pg_stats.n_distinct (Column) +n_ins_since_vacuum - pg_catalog.pg_stat_all_tables.n_ins_since_vacuum (Column) + +-------------- + +update instruments as i set name | where i.id = 1; +update instruments as i set name = | where i.id = 1; + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) + +-------------- + +update instruments as i set name = '| where i.id = 1; +update instruments as i set name = 'x', | where i.id = 1; + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +public - public (Schema) +information_schema - information_schema (Schema) + +-------------- + +update instruments as i set name = 'x', z | where i.id = 1; +update instruments as i set name = 'x', z = | where i.id = 1; + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) + +-------------- + +update instruments as i set name = 'x', z = '| where i.id = 1; +update instruments as i set name = 'x', z = 'y' | where i.id = 1; + +Results: +from - from (Keyword) + +-------------- + + + + + +***Case 3:*** + +update instruments set name = 'x' | + +Results: +from - from (Keyword) +where - where (Keyword) + +-------------- + +update instruments set name = 'x' r| + +Results: +from - from (Keyword) +where - where (Keyword) + +-------------- + +update instruments set name = 'x' returning | + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +a - public.others.a (Column) +b - public.others.b (Column) + +-------------- + +update instruments set name = 'x' returning i| + +Results: +id - public.instruments.id (Column) +ident - pg_catalog.pg_backend_memory_contexts.ident (Column) +identity_cycle - information_schema.columns.identity_cycle (Column) +identity_generation - information_schema.columns.identity_generation (Column) +identity_increment - information_schema.columns.identity_increment (Column) + +-------------- + +update instruments set name = 'x' returning id, | + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +a - public.others.a (Column) +b - public.others.b (Column) + +-------------- + +update instruments set name = 'x' returning id, n| + +Results: +name - public.instruments.name (Column) +n_dead_tup - pg_catalog.pg_stat_all_tables.n_dead_tup (Column) +n_distinct - pg_catalog.pg_stats.n_distinct (Column) +n_ins_since_vacuum - pg_catalog.pg_stat_all_tables.n_ins_since_vacuum (Column) +n_live_tup - pg_catalog.pg_stat_all_tables.n_live_tup (Column) + +-------------- + +update instruments set name = 'x' returning id, name; | + +Results: +copy - copy (Keyword) +create - create (Keyword) +drop - drop (Keyword) +insert - insert (Keyword) +reset - reset (Keyword) + +-------------- + + + + + +***Case 4:*** + +update instruments | where id = 1; + +Results: +set - set (Keyword) + +-------------- + +update instruments s| where id = 1; + +Results: +set - set (Keyword) + +-------------- + +update instruments set | where id = 1; + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +public - public (Schema) +information_schema - information_schema (Schema) + +-------------- + +update instruments set n| where id = 1; + +Results: +name - public.instruments.name (Column) +information_schema - information_schema (Schema) +n_dead_tup - pg_catalog.pg_stat_all_tables.n_dead_tup (Column) +n_distinct - pg_catalog.pg_stats.n_distinct (Column) +n_ins_since_vacuum - pg_catalog.pg_stat_all_tables.n_ins_since_vacuum (Column) + +-------------- + +update instruments set name | where id = 1; +update instruments set name = | where id = 1; + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) + +-------------- + +update instruments set name = d| where id = 1; + +Results: +id - public.instruments.id (Column) +information_schema.data_type_privileges - information_schema.data_type_privileges (Table) +default - default (Keyword) +information_schema.domain_constraints - information_schema.domain_constraints (Table) +information_schema.domain_udt_usage - information_schema.domain_udt_usage (Table) + +-------------- + +update instruments set name = default, | where id = 1; + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +public - public (Schema) +information_schema - information_schema (Schema) + +-------------- + +update instruments set name = default, z | where id = 1; +update instruments set name = default, z = | where id = 1; + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) + +-------------- + +update instruments set name = default, z = '| where id = 1; +update instruments set name = default, z = 'y' | where id = 1; + +Results: +from - from (Keyword) + +-------------- + + + + + +***Case 5:*** + +update instruments set name = o.a | + +Results: +from - from (Keyword) +where - where (Keyword) + +-------------- + +update instruments set name = o.a f| + +Results: +from - from (Keyword) + +-------------- + +update instruments set name = o.a from | + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) + +-------------- + +update instruments set name = o.a from o| + +Results: +others - public.others (Table) +information_schema - information_schema (Schema) +pg_catalog - pg_catalog (Schema) +pg_toast - pg_toast (Schema) +information_schema.column_options - information_schema.column_options (Table) + +-------------- + +update instruments set name = o.a from others | + +Results: +cross - cross (Keyword) +full - full (Keyword) +inner - inner (Keyword) +join - join (Keyword) +left - left (Keyword) + +-------------- + +update instruments set name = o.a from others o | + +Results: +cross - cross (Keyword) +full - full (Keyword) +inner - inner (Keyword) +join - join (Keyword) +left - left (Keyword) + +-------------- + +update instruments set name = o.a from others o w| + +Results: +where - where (Keyword) + +-------------- + +update instruments set name = o.a from others o where | + +Results: +o.a - public.others.a (Column) +o.b - public.others.b (Column) +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) + +-------------- + +update instruments set name = o.a from others o where i| + +Results: +id - public.instruments.id (Column) +ident - pg_catalog.pg_backend_memory_contexts.ident (Column) +identity_cycle - information_schema.columns.identity_cycle (Column) +identity_generation - information_schema.columns.identity_generation (Column) +identity_increment - information_schema.columns.identity_increment (Column) + +-------------- + +update instruments set name = o.a from others o where instruments.| + +Results: +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) + +-------------- + +update instruments set name = o.a from others o where instruments.i| + +Results: +id - public.instruments.id (Column) + +-------------- + +update instruments set name = o.a from others o where instruments.id | + +Results: +returning - returning (Keyword) + +-------------- + +update instruments set name = o.a from others o where instruments.id = | + +Results: +o.a - public.others.a (Column) +o.b - public.others.b (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) +id - public.instruments.id (Column) + +-------------- + +update instruments set name = o.a from others o where instruments.id = o.| + +Results: +a - public.others.a (Column) +b - public.others.b (Column) +id - public.others.id (Column) + +-------------- + +update instruments set name = o.a from others o where instruments.id = o.i| + +Results: +id - public.others.id (Column) + +-------------- + +update instruments set name = o.a from others o where instruments.id = o.id | + +Results: +returning - returning (Keyword) + +-------------- + +update instruments set name = o.a from others o where instruments.id = o.id r| + +Results: +returning - returning (Keyword) + +-------------- + +update instruments set name = o.a from others o where instruments.id = o.id returning | + +Results: +o.a - public.others.a (Column) +o.b - public.others.b (Column) +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) + +-------------- + +update instruments set name = o.a from others o where instruments.id = o.id returning n| + +Results: +name - public.instruments.name (Column) +n_dead_tup - pg_catalog.pg_stat_all_tables.n_dead_tup (Column) +n_distinct - pg_catalog.pg_stats.n_distinct (Column) +n_ins_since_vacuum - pg_catalog.pg_stat_all_tables.n_ins_since_vacuum (Column) +n_live_tup - pg_catalog.pg_stat_all_tables.n_live_tup (Column) + +-------------- + +update instruments set name = o.a from others o where instruments.id = o.id returning name, | + +Results: +o.a - public.others.a (Column) +o.b - public.others.b (Column) +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) + +-------------- + +update instruments set name = o.a from others o where instruments.id = o.id returning name, z| + +Results: +z - public.instruments.z (Column) +blks_zeroed - pg_catalog.pg_stat_slru.blks_zeroed (Column) +allocated_size - pg_catalog.pg_shmem_allocations.allocated_size (Column) +analyze_count - pg_catalog.pg_stat_all_tables.analyze_count (Column) +authorization_identifier - information_schema._pg_foreign_data_wrappers.authorization_identifier (Column) + +-------------- + +update instruments set name = o.a from others o where instruments.id = o.id returning name, z; | + +Results: +copy - copy (Keyword) +create - create (Keyword) +drop - drop (Keyword) +insert - insert (Keyword) +reset - reset (Keyword) + +-------------- + + + + + +***Case 6:*** + +update instruments set name = (|) where id = 1; + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +id - public.instruments.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) + +-------------- + +update instruments set name = (s|) where id = 1; + +Results: +instruments - public.instruments (Table) +others - public.others (Table) +information_schema.schemata - information_schema.schemata (Table) +select - select (Keyword) +information_schema.sequences - information_schema.sequences (Table) + +-------------- + +update instruments set name = (select |) where id = 1; +update instruments set name = (select a |) where id = 1; +update instruments set name = (select a f|) where id = 1; + +Results: +for - for (Keyword) +from - from (Keyword) +offset - offset (Keyword) + +-------------- + +update instruments set name = (select a from |) where id = 1; +update instruments set name = (select a from o|) where id = 1; + +Results: +others - public.others (Table) +only - only (Keyword) +information_schema - information_schema (Schema) +pg_catalog - pg_catalog (Schema) +pg_toast - pg_toast (Schema) + +-------------- + +update instruments set name = (select a from others |) where id = 1; +update instruments set name = (select a from others w|) where id = 1; + +Results: +where - where (Keyword) +window - window (Keyword) + +-------------- + +update instruments set name = (select a from others where |) where id = 1; +update instruments set name = (select a from others where i|) where id = 1; + +Results: +id - public.others.id (Column) +ident - pg_catalog.pg_backend_memory_contexts.ident (Column) +identity_cycle - information_schema.columns.identity_cycle (Column) +identity_generation - information_schema.columns.identity_generation (Column) +identity_increment - information_schema.columns.identity_increment (Column) + +-------------- + +update instruments set name = (select a from others where id |) where id = 1; +update instruments set name = (select a from others where id = |) where id = 1; + +Results: +a - public.others.a (Column) +b - public.others.b (Column) +id - public.others.id (Column) +name - public.instruments.name (Column) +z - public.instruments.z (Column) + +-------------- + +update instruments set name = (select a from others where id = 1 |) where id = 1; diff --git a/crates/pgls_treesitter/src/context/mod.rs b/crates/pgls_treesitter/src/context/mod.rs index 8819ca78f..c8d81fa22 100644 --- a/crates/pgls_treesitter/src/context/mod.rs +++ b/crates/pgls_treesitter/src/context/mod.rs @@ -28,6 +28,7 @@ pub enum WrappingClause<'a> { AlterTable, DropTable, DropColumn, + Returning, AlterColumn, RenameColumn, SetStatement, @@ -35,6 +36,7 @@ pub enum WrappingClause<'a> { DropRole, RevokeStatement, GrantStatement, + CopyStatement, CreatePolicy, AlterPolicy, @@ -394,12 +396,14 @@ impl<'a> TreesitterContext<'a> { "alter_role" => Some(WrappingClause::AlterRole), "drop_role" => Some(WrappingClause::DropRole), "drop_column" => Some(WrappingClause::DropColumn), + "returning" => Some(WrappingClause::Returning), "alter_column" => Some(WrappingClause::AlterColumn), "rename_column" => Some(WrappingClause::RenameColumn), "alter_table" => Some(WrappingClause::AlterTable), "set_statement" => Some(WrappingClause::SetStatement), "revoke_statement" => Some(WrappingClause::RevokeStatement), "grant_statement" => Some(WrappingClause::GrantStatement), + "copy_statement" => Some(WrappingClause::CopyStatement), "column_definitions" => Some(WrappingClause::ColumnDefinitions), "create_policy" => Some(WrappingClause::CreatePolicy), "alter_policy" => Some(WrappingClause::AlterPolicy), diff --git a/justfile b/justfile index a8317f4f9..b4fd3996a 100644 --- a/justfile +++ b/justfile @@ -70,7 +70,7 @@ test: # Run tests for the crate passed as argument e.g. just test-create pgls_cli test-crate name: - cargo test run -p {{name}} --no-fail-fast + cargo test -p {{name}} --no-fail-fast # Run doc tests test-doc: