From b52b7961591d90703b9f80967d7dcdcc62dadc47 Mon Sep 17 00:00:00 2001 From: Julian Date: Tue, 27 May 2025 08:57:45 +0200 Subject: [PATCH 1/5] feat(completions): complete (materialized) views --- .../pgt_completions/src/providers/columns.rs | 2 +- .../pgt_completions/src/providers/tables.rs | 12 +- crates/pgt_schema_cache/src/lib.rs | 2 +- .../pgt_schema_cache/src/queries/tables.sql | 3 +- crates/pgt_schema_cache/src/tables.rs | 108 ++++++++++++++++++ 5 files changed, 122 insertions(+), 5 deletions(-) diff --git a/crates/pgt_completions/src/providers/columns.rs b/crates/pgt_completions/src/providers/columns.rs index d4767f14..dd209d50 100644 --- a/crates/pgt_completions/src/providers/columns.rs +++ b/crates/pgt_completions/src/providers/columns.rs @@ -17,7 +17,7 @@ pub fn complete_columns<'a>(ctx: &CompletionContext<'a>, builder: &mut Completio label: col.name.clone(), score: CompletionScore::from(relevance.clone()), filter: CompletionFilter::from(relevance), - description: format!("Table: {}.{}", col.schema_name, col.table_name), + description: format!("{}.{}", col.schema_name, col.table_name), kind: CompletionItemKind::Column, completion_text: None, detail: col.type_name.as_ref().map(|t| t.to_string()), diff --git a/crates/pgt_completions/src/providers/tables.rs b/crates/pgt_completions/src/providers/tables.rs index 6ed3760e..46c3950b 100644 --- a/crates/pgt_completions/src/providers/tables.rs +++ b/crates/pgt_completions/src/providers/tables.rs @@ -13,13 +13,21 @@ pub fn complete_tables<'a>(ctx: &'a CompletionContext, builder: &mut CompletionB for table in available_tables { let relevance = CompletionRelevanceData::Table(table); + let detail: Option = match table.table_kind { + pgt_schema_cache::TableKind::Ordinary | pgt_schema_cache::TableKind::Partitioned => { + None + } + pgt_schema_cache::TableKind::View => Some("View".into()), + pgt_schema_cache::TableKind::MaterializedView => Some("MView".into()), + }; + let item = PossibleCompletionItem { label: table.name.clone(), score: CompletionScore::from(relevance.clone()), filter: CompletionFilter::from(relevance), - description: format!("Schema: {}", table.schema), + description: format!("{}", table.schema), kind: CompletionItemKind::Table, - detail: None, + detail, completion_text: get_completion_text_with_schema_or_alias( ctx, &table.name, diff --git a/crates/pgt_schema_cache/src/lib.rs b/crates/pgt_schema_cache/src/lib.rs index 186fbdb9..9beb2f8a 100644 --- a/crates/pgt_schema_cache/src/lib.rs +++ b/crates/pgt_schema_cache/src/lib.rs @@ -19,6 +19,6 @@ pub use policies::{Policy, PolicyCommand}; pub use roles::*; pub use schema_cache::SchemaCache; pub use schemas::Schema; -pub use tables::{ReplicaIdentity, Table}; +pub use tables::{ReplicaIdentity, Table, TableKind}; pub use triggers::{Trigger, TriggerAffected, TriggerEvent}; pub use types::{PostgresType, PostgresTypeAttribute}; diff --git a/crates/pgt_schema_cache/src/queries/tables.sql b/crates/pgt_schema_cache/src/queries/tables.sql index bcce4fcc..6e6865a2 100644 --- a/crates/pgt_schema_cache/src/queries/tables.sql +++ b/crates/pgt_schema_cache/src/queries/tables.sql @@ -2,6 +2,7 @@ select c.oid :: int8 as "id!", nc.nspname as schema, c.relname as name, + c.relkind as table_kind, c.relrowsecurity as rls_enabled, c.relforcerowsecurity as rls_forced, case @@ -21,7 +22,7 @@ from pg_namespace nc join pg_class c on nc.oid = c.relnamespace where - c.relkind in ('r', 'p') + c.relkind in ('r', 'p', 'v', 'm') and not pg_is_other_temp_schema(nc.oid) and ( pg_has_role(c.relowner, 'USAGE') diff --git a/crates/pgt_schema_cache/src/tables.rs b/crates/pgt_schema_cache/src/tables.rs index 99061384..798a03a6 100644 --- a/crates/pgt_schema_cache/src/tables.rs +++ b/crates/pgt_schema_cache/src/tables.rs @@ -23,6 +23,40 @@ impl From for ReplicaIdentity { } } +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub enum TableKind { + #[default] + Ordinary, + View, + MaterializedView, + Partitioned, +} + +impl From for TableKind { + fn from(s: char) -> Self { + match s { + 'r' => Self::Ordinary, + 'p' => Self::Partitioned, + 'v' => Self::View, + 'm' => Self::MaterializedView, + _ => panic!("Invalid table kind"), + } + } +} + +impl From for TableKind { + fn from(s: i8) -> Self { + let c = char::from(u8::try_from(s).unwrap()); + match c { + 'r' => Self::Ordinary, + 'p' => Self::Partitioned, + 'v' => Self::View, + 'm' => Self::MaterializedView, + _ => panic!("Invalid table kind"), + } + } +} + #[derive(Debug, Default, PartialEq, Eq)] pub struct Table { pub id: i64, @@ -31,6 +65,7 @@ pub struct Table { pub rls_enabled: bool, pub rls_forced: bool, pub replica_identity: ReplicaIdentity, + pub table_kind: TableKind, pub bytes: i64, pub size: String, pub live_rows_estimate: i64, @@ -47,3 +82,76 @@ impl SchemaCacheItem for Table { .await } } + +#[cfg(test)] +mod tests { + use crate::{SchemaCache, tables::TableKind}; + use pgt_test_utils::test_database::get_new_test_db; + use sqlx::Executor; + + #[tokio::test] + async fn includes_views_in_query() { + let test_db = get_new_test_db().await; + + let setup = r#" + create table public.base_table ( + id serial primary key, + value text + ); + + create view public.my_view as + select * from public.base_table; + "#; + + test_db + .execute(setup) + .await + .expect("Failed to setup test database"); + + let cache = SchemaCache::load(&test_db) + .await + .expect("Failed to load Schema Cache"); + + let view = cache + .tables + .iter() + .find(|t| t.name == "my_view") + .expect("View not found"); + + assert_eq!(view.table_kind, TableKind::View); + assert_eq!(view.schema, "public"); + } + + #[tokio::test] + async fn includes_materialized_views_in_query() { + let test_db = get_new_test_db().await; + + let setup = r#" + create table public.base_table ( + id serial primary key, + value text + ); + + create materialized view public.my_mat_view as + select * from public.base_table; + "#; + + test_db + .execute(setup) + .await + .expect("Failed to setup test database"); + + let cache = SchemaCache::load(&test_db) + .await + .expect("Failed to load Schema Cache"); + + let mat_view = cache + .tables + .iter() + .find(|t| t.name == "my_mat_view") + .expect("Materialized view not found"); + + assert_eq!(mat_view.table_kind, TableKind::MaterializedView); + assert_eq!(mat_view.schema, "public"); + } +} From 52e2d6f45250e6dae30313895b55bad05004a1ac Mon Sep 17 00:00:00 2001 From: Julian Date: Tue, 27 May 2025 09:00:20 +0200 Subject: [PATCH 2/5] ok --- ...cf4c296b594fe9e6cebbdc382acde73f4fb9.json} | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) rename .sqlx/{query-2a964a12383b977bbbbd6fe7298dfce00358ecbe878952e8d4915c06cc5c9e0f.json => query-66d92238c94b5f1c99fbf068a0b5cf4c296b594fe9e6cebbdc382acde73f4fb9.json} (50%) diff --git a/.sqlx/query-2a964a12383b977bbbbd6fe7298dfce00358ecbe878952e8d4915c06cc5c9e0f.json b/.sqlx/query-66d92238c94b5f1c99fbf068a0b5cf4c296b594fe9e6cebbdc382acde73f4fb9.json similarity index 50% rename from .sqlx/query-2a964a12383b977bbbbd6fe7298dfce00358ecbe878952e8d4915c06cc5c9e0f.json rename to .sqlx/query-66d92238c94b5f1c99fbf068a0b5cf4c296b594fe9e6cebbdc382acde73f4fb9.json index 96439422..447ba93b 100644 --- a/.sqlx/query-2a964a12383b977bbbbd6fe7298dfce00358ecbe878952e8d4915c06cc5c9e0f.json +++ b/.sqlx/query-66d92238c94b5f1c99fbf068a0b5cf4c296b594fe9e6cebbdc382acde73f4fb9.json @@ -1,6 +1,6 @@ { "db_name": "PostgreSQL", - "query": "select\n c.oid :: int8 as \"id!\",\n nc.nspname as schema,\n c.relname as name,\n c.relrowsecurity as rls_enabled,\n c.relforcerowsecurity as rls_forced,\n case\n when c.relreplident = 'd' then 'DEFAULT'\n when c.relreplident = 'i' then 'INDEX'\n when c.relreplident = 'f' then 'FULL'\n else 'NOTHING'\n end as \"replica_identity!\",\n pg_total_relation_size(format('%I.%I', nc.nspname, c.relname)) :: int8 as \"bytes!\",\n pg_size_pretty(\n pg_total_relation_size(format('%I.%I', nc.nspname, c.relname))\n ) as \"size!\",\n pg_stat_get_live_tuples(c.oid) as \"live_rows_estimate!\",\n pg_stat_get_dead_tuples(c.oid) as \"dead_rows_estimate!\",\n obj_description(c.oid) as comment\nfrom\n pg_namespace nc\n join pg_class c on nc.oid = c.relnamespace\nwhere\n c.relkind in ('r', 'p')\n and not pg_is_other_temp_schema(nc.oid)\n and (\n pg_has_role(c.relowner, 'USAGE')\n or has_table_privilege(\n c.oid,\n 'SELECT, INSERT, UPDATE, DELETE, TRUNCATE, REFERENCES, TRIGGER'\n )\n or has_any_column_privilege(c.oid, 'SELECT, INSERT, UPDATE, REFERENCES')\n )\ngroup by\n c.oid,\n c.relname,\n c.relrowsecurity,\n c.relforcerowsecurity,\n c.relreplident,\n nc.nspname;", + "query": "select\n c.oid :: int8 as \"id!\",\n nc.nspname as schema,\n c.relname as name,\n c.relkind as table_kind,\n c.relrowsecurity as rls_enabled,\n c.relforcerowsecurity as rls_forced,\n case\n when c.relreplident = 'd' then 'DEFAULT'\n when c.relreplident = 'i' then 'INDEX'\n when c.relreplident = 'f' then 'FULL'\n else 'NOTHING'\n end as \"replica_identity!\",\n pg_total_relation_size(format('%I.%I', nc.nspname, c.relname)) :: int8 as \"bytes!\",\n pg_size_pretty(\n pg_total_relation_size(format('%I.%I', nc.nspname, c.relname))\n ) as \"size!\",\n pg_stat_get_live_tuples(c.oid) as \"live_rows_estimate!\",\n pg_stat_get_dead_tuples(c.oid) as \"dead_rows_estimate!\",\n obj_description(c.oid) as comment\nfrom\n pg_namespace nc\n join pg_class c on nc.oid = c.relnamespace\nwhere\n c.relkind in ('r', 'p', 'v', 'm')\n and not pg_is_other_temp_schema(nc.oid)\n and (\n pg_has_role(c.relowner, 'USAGE')\n or has_table_privilege(\n c.oid,\n 'SELECT, INSERT, UPDATE, DELETE, TRUNCATE, REFERENCES, TRIGGER'\n )\n or has_any_column_privilege(c.oid, 'SELECT, INSERT, UPDATE, REFERENCES')\n )\ngroup by\n c.oid,\n c.relname,\n c.relrowsecurity,\n c.relforcerowsecurity,\n c.relreplident,\n nc.nspname;", "describe": { "columns": [ { @@ -20,41 +20,46 @@ }, { "ordinal": 3, + "name": "table_kind", + "type_info": "Char" + }, + { + "ordinal": 4, "name": "rls_enabled", "type_info": "Bool" }, { - "ordinal": 4, + "ordinal": 5, "name": "rls_forced", "type_info": "Bool" }, { - "ordinal": 5, + "ordinal": 6, "name": "replica_identity!", "type_info": "Text" }, { - "ordinal": 6, + "ordinal": 7, "name": "bytes!", "type_info": "Int8" }, { - "ordinal": 7, + "ordinal": 8, "name": "size!", "type_info": "Text" }, { - "ordinal": 8, + "ordinal": 9, "name": "live_rows_estimate!", "type_info": "Int8" }, { - "ordinal": 9, + "ordinal": 10, "name": "dead_rows_estimate!", "type_info": "Int8" }, { - "ordinal": 10, + "ordinal": 11, "name": "comment", "type_info": "Text" } @@ -68,6 +73,7 @@ false, false, false, + false, null, null, null, @@ -76,5 +82,5 @@ null ] }, - "hash": "2a964a12383b977bbbbd6fe7298dfce00358ecbe878952e8d4915c06cc5c9e0f" + "hash": "66d92238c94b5f1c99fbf068a0b5cf4c296b594fe9e6cebbdc382acde73f4fb9" } From 22e1e396d65565b1d3312ece3f92d0c99d7de7fb Mon Sep 17 00:00:00 2001 From: Julian Date: Tue, 27 May 2025 09:01:58 +0200 Subject: [PATCH 3/5] ok --- crates/pgt_completions/src/providers/tables.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/pgt_completions/src/providers/tables.rs b/crates/pgt_completions/src/providers/tables.rs index 46c3950b..2102d41c 100644 --- a/crates/pgt_completions/src/providers/tables.rs +++ b/crates/pgt_completions/src/providers/tables.rs @@ -25,7 +25,7 @@ pub fn complete_tables<'a>(ctx: &'a CompletionContext, builder: &mut CompletionB label: table.name.clone(), score: CompletionScore::from(relevance.clone()), filter: CompletionFilter::from(relevance), - description: format!("{}", table.schema), + description: table.schema.to_string(), kind: CompletionItemKind::Table, detail, completion_text: get_completion_text_with_schema_or_alias( From 1e24e22b2ddea0df1b013d40840b18e54aa45881 Mon Sep 17 00:00:00 2001 From: Julian Date: Tue, 27 May 2025 09:18:01 +0200 Subject: [PATCH 4/5] where i come from, this is considered a classical oopsie --- .../pgt_completions/src/providers/columns.rs | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/crates/pgt_completions/src/providers/columns.rs b/crates/pgt_completions/src/providers/columns.rs index dd209d50..da6d23bc 100644 --- a/crates/pgt_completions/src/providers/columns.rs +++ b/crates/pgt_completions/src/providers/columns.rs @@ -92,7 +92,7 @@ mod tests { message: "correctly prefers the columns of present tables", query: format!(r#"select na{} from public.audio_books;"#, CURSOR_POS), label: "narrator", - description: "Table: public.audio_books", + description: "public.audio_books", }, TestCase { message: "correctly handles nested queries", @@ -110,13 +110,13 @@ mod tests { CURSOR_POS ), label: "narrator_id", - description: "Table: private.audio_books", + description: "private.audio_books", }, TestCase { message: "works without a schema", query: format!(r#"select na{} from users;"#, CURSOR_POS), label: "name", - description: "Table: public.users", + description: "public.users", }, ]; @@ -186,10 +186,10 @@ mod tests { .collect(); let expected = vec![ - ("name", "Table: public.users"), - ("narrator", "Table: public.audio_books"), - ("narrator_id", "Table: private.audio_books"), - ("id", "Table: public.audio_books"), + ("name", "public.users"), + ("narrator", "public.audio_books"), + ("narrator_id", "private.audio_books"), + ("id", "public.audio_books"), ("name", "Schema: pg_catalog"), ("nameconcatoid", "Schema: pg_catalog"), ] @@ -559,10 +559,7 @@ mod tests { ) .as_str(), vec![ - CompletionAssertion::LabelAndDesc( - "id".to_string(), - "Table: public.two".to_string(), - ), + CompletionAssertion::LabelAndDesc("id".to_string(), "public.two".to_string()), CompletionAssertion::Label("z".to_string()), ], setup, From 3c2699e08c656dfe20326c41c0120655a9e378f5 Mon Sep 17 00:00:00 2001 From: Julian Date: Tue, 27 May 2025 10:03:06 +0200 Subject: [PATCH 5/5] well --- crates/pgt_schema_cache/src/tables.rs | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/crates/pgt_schema_cache/src/tables.rs b/crates/pgt_schema_cache/src/tables.rs index 798a03a6..a0a40d6a 100644 --- a/crates/pgt_schema_cache/src/tables.rs +++ b/crates/pgt_schema_cache/src/tables.rs @@ -47,13 +47,7 @@ impl From for TableKind { impl From for TableKind { fn from(s: i8) -> Self { let c = char::from(u8::try_from(s).unwrap()); - match c { - 'r' => Self::Ordinary, - 'p' => Self::Partitioned, - 'v' => Self::View, - 'm' => Self::MaterializedView, - _ => panic!("Invalid table kind"), - } + c.into() } }