feat(completions): complete (materialized) views (#409)

Author: juleswritescode

.sqlx/query-2a964a12383b977bbbbd6fe7298dfce00358ecbe878952e8d4915c06cc5c9e0f.json renamed to .sqlx/query-66d92238c94b5f1c99fbf068a0b5cf4c296b594fe9e6cebbdc382acde73f4fb9.json

@@ -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()),
@@ -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,

crates/pgt_completions/src/providers/columns.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<String> = 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: table.schema.to_string(),
             kind: CompletionItemKind::Table,
-            detail: None,
+            detail,
             completion_text: get_completion_text_with_schema_or_alias(
                 ctx,
                 &table.name,

crates/pgt_completions/src/providers/tables.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};

crates/pgt_schema_cache/src/lib.rs

@@ -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')

crates/pgt_schema_cache/src/queries/tables.sql

@@ -23,6 +23,34 @@ impl From<String> for ReplicaIdentity {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Default)]
+pub enum TableKind {
+    #[default]
+    Ordinary,
+    View,
+    MaterializedView,
+    Partitioned,
+}
+
+impl From<char> 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<i8> for TableKind {
+    fn from(s: i8) -> Self {
+        let c = char::from(u8::try_from(s).unwrap());
+        c.into()
+    }
+}
+
 #[derive(Debug, Default, PartialEq, Eq)]
 pub struct Table {
     pub id: i64,
@@ -31,6 +59,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 +76,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");
+    }
+}