draft: query columns type

Author: juleswritescode

crates/pg_completions/src/providers/columns.rs

@@ -0,0 +1,20 @@
+use crate::{
+    builder::CompletionBuilder, context::CompletionContext, relevance::CompletionRelevanceData,
+    CompletionItem, CompletionItemKind,
+};
+
+pub fn complete_functions(ctx: &CompletionContext, builder: &mut CompletionBuilder) {
+    let available_functions = &ctx.schema_cache.functions;
+
+    for foo in available_functions {
+        let item = CompletionItem {
+            label: foo.name.clone(),
+            score: CompletionRelevanceData::Function(foo).get_score(ctx),
+            description: format!("Schema: {}", foo.schema),
+            preselected: false,
+            kind: CompletionItemKind::Function,
+        };
+
+        builder.add_item(item);
+    }
+}

crates/pg_completions/src/providers/mod.rs

@@ -1,5 +1,7 @@
+mod columns;
 mod functions;
 mod tables;
 
+pub use columns::*;
 pub use functions::*;
 pub use tables::*;

crates/pg_schema_cache/src/columns.rs

@@ -0,0 +1,73 @@
+use crate::schema_cache::SchemaCacheItem;
+
+enum ColumnClassKind {
+    OrdinaryTable,
+    View,
+    MaterializedView,
+    ForeignTable,
+    PartitionedTable,
+}
+
+impl From<&str> for ColumnClassKind {
+    fn from(value: &str) -> Self {
+        match value {
+            "r" => ColumnClassKind::OrdinaryTable,
+            "v" => ColumnClassKind::View,
+            "m" => ColumnClassKind::MaterializedView,
+            "f" => ColumnClassKind::ForeignTable,
+            "p" => ColumnClassKind::PartitionedTable,
+            _ => panic!(
+                "Columns belonging to a class with pg_class.relkind = '{}' should be filtered out in the query.",
+                value
+            ),
+        }
+    }
+}
+
+impl From<String> for ColumnClassKind {
+    fn from(value: String) -> Self {
+        ColumnClassKind::from(value.as_str())
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Column {
+    pub name: String,
+
+    pub table_name: String,
+    pub table_oid: i64,
+
+    /// What type of class does this column belong to?
+    pub class_kind: bool,
+
+    pub schema_name: String,
+    pub type_id: i64,
+    // pub is_nullable: bool,
+    // pub is_primary_key: bool,
+
+    // pub default_value: Option<String>,
+
+    // pub varchar_length: Option<i32>,
+
+    // /// None if the column is not a foreign key.
+    // pub foreign_key: Option<ForeignKeyReference>,
+
+    // pub comment: Option<String>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct ForeignKeyReference {
+    pub schema: Option<String>,
+    pub table: String,
+    pub column: String,
+}
+
+impl SchemaCacheItem for Column {
+    type Item = Column;
+
+    async fn load(pool: &sqlx::PgPool) -> Result<Vec<Self::Item>, sqlx::Error> {
+        sqlx::query_file_as!(Column, "src/queries/columns.sql")
+            .fetch_all(pool)
+            .await
+    }
+}

crates/pg_schema_cache/src/lib.rs

@@ -2,6 +2,7 @@
 
 #![allow(dead_code)]
 
+mod columns;
 mod functions;
 mod schema_cache;
 mod schemas;

crates/pg_schema_cache/src/queries/columns.sql

@@ -0,0 +1,75 @@
+select
+  tbl.schemaname as schema,
+  tbl.tablename as table,
+  tbl.quoted_name,
+  tbl.is_table,
+  json_agg(a) as columns
+from
+  (
+    select
+      n.nspname as schemaname,
+      c.relname as tablename,
+      (
+        quote_ident(n.nspname) || '.' || quote_ident(c.relname)
+      ) as quoted_name,
+      true as is_table
+    from
+      pg_catalog.pg_class c
+      join pg_catalog.pg_namespace n on n.oid = c.relnamespace
+    where
+      c.relkind = 'r'
+      and n.nspname != 'pg_toast'
+      and n.nspname not like 'pg_temp_%'
+      and n.nspname not like 'pg_toast_temp_%'
+      and has_schema_privilege(n.oid, 'USAGE') = true
+      and has_table_privilege(
+        quote_ident(n.nspname) || '.' || quote_ident(c.relname),
+        'SELECT, INSERT, UPDATE, DELETE, TRUNCATE, REFERENCES, TRIGGER'
+      ) = true
+    union
+    all
+    select
+      n.nspname as schemaname,
+      c.relname as tablename,
+      (
+        quote_ident(n.nspname) || '.' || quote_ident(c.relname)
+      ) as quoted_name,
+      false as is_table
+    from
+      pg_catalog.pg_class c
+      join pg_catalog.pg_namespace n on n.oid = c.relnamespace
+    where
+      c.relkind in ('v', 'm')
+      and n.nspname != 'pg_toast'
+      and n.nspname not like 'pg_temp_%'
+      and n.nspname not like 'pg_toast_temp_%'
+      and has_schema_privilege(n.oid, 'USAGE') = true
+      and has_table_privilege(
+        quote_ident(n.nspname) || '.' || quote_ident(c.relname),
+        'SELECT, INSERT, UPDATE, DELETE, TRUNCATE, REFERENCES, TRIGGER'
+      ) = true
+  ) as tbl
+  left join (
+    select
+      attrelid,
+      attname,
+      format_type(atttypid, atttypmod) as data_type,
+      attnum,
+      attisdropped
+    from
+      pg_attribute
+  ) as a on (
+    a.attrelid = tbl.quoted_name :: regclass
+    and a.attnum > 0
+    and not a.attisdropped
+    and has_column_privilege(
+      tbl.quoted_name,
+      a.attname,
+      'SELECT, INSERT, UPDATE, REFERENCES'
+    )
+  )
+group by
+  schemaname,
+  tablename,
+  quoted_name,
+  is_table;

crates/pg_schema_cache/src/schema_cache.rs

@@ -58,6 +58,10 @@ impl SchemaCache {
             .find(|t| t.name == name && schema.is_none() || Some(t.schema.as_str()) == schema)
     }
 
+    pub fn find_type_by_id(&self, type_id: i64) -> Option<&PostgresType> {
+        self.types.iter().find(|t| t.id == type_id)
+    }
+
     pub fn find_types(&self, name: &str, schema: Option<&str>) -> Vec<&PostgresType> {
         self.types
             .iter()