feat(schema_cache): column type + query

Cargo.toml

@@ -36,7 +36,7 @@ sqlx                     = { version = "0.8.2", features = ["runtime-async-std",
 syn                      = "1.0.109"
 termcolor                = "1.4.1"
 text-size                = "1.1.1"
-tokio                    = "1.40.0"
+tokio                    = { version = "1.40.0", features = ["full"] }
 toml                     = "0.8.19"
 tower-lsp                = "0.20.0"
 tracing                  = { version = "0.1.40", default-features = false, features = ["std"] }

crates/pg_schema_cache/Cargo.toml

@@ -20,6 +20,7 @@ serde_json.workspace = true
 pg_diagnostics.workspace = true
 pg_console.workspace = true
 sqlx.workspace = true
+tokio.workspace = true
 
 [dev-dependencies]
 pg_test_utils.workspace = true

crates/pg_schema_cache/src/columns.rs

@@ -0,0 +1,202 @@
+use crate::schema_cache::SchemaCacheItem;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub 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())
+    }
+}
+
+impl From<char> for ColumnClassKind {
+    fn from(value: char) -> Self {
+        ColumnClassKind::from(String::from(value))
+    }
+}
+
+#[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: ColumnClassKind,
+
+    pub schema_name: String,
+    pub type_id: i64,
+    pub is_nullable: bool,
+
+    pub is_primary_key: bool,
+    pub is_unique: bool,
+
+    /// The Default "value" of the column. Might be a function call, hence "_expr".
+    pub default_expr: Option<String>,
+
+    pub varchar_length: Option<i32>,
+
+    /// Comment inserted via `COMMENT ON COLUMN my_table.my_comment '...'`, if present.
+    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
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use pg_test_utils::test_database::get_new_test_db;
+    use sqlx::Executor;
+
+    use crate::{columns::ColumnClassKind, SchemaCache};
+
+    #[tokio::test]
+    async fn loads_columns() {
+        let test_db = get_new_test_db().await;
+
+        let setup = r#"
+            create table public.users (
+                id serial primary key,
+                name varchar(255) not null,
+                is_vegetarian bool default false,
+                middle_name varchar(255)
+            );
+
+            create schema real_estate;
+
+            create table real_estate.addresses (
+                user_id serial references users(id),
+                postal_code smallint not null,
+                street text,
+                city text
+            );
+
+            create table real_estate.properties (
+                id serial primary key,
+                owner_id int references users(id),
+                square_meters smallint not null
+            );
+
+            comment on column real_estate.properties.owner_id is 'users might own many houses';
+        "#;
+
+        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 public_schema_columns = cache
+            .columns
+            .iter()
+            .filter(|c| c.schema_name.as_str() == "public")
+            .count();
+
+        assert_eq!(public_schema_columns, 4);
+
+        let real_estate_schema_columns = cache
+            .columns
+            .iter()
+            .filter(|c| c.schema_name.as_str() == "real_estate")
+            .count();
+
+        assert_eq!(real_estate_schema_columns, 7);
+
+        let user_id_col = cache.find_col("id", "users", None).unwrap();
+        assert_eq!(user_id_col.class_kind, ColumnClassKind::OrdinaryTable);
+        assert_eq!(user_id_col.comment, None);
+        assert_eq!(
+            user_id_col.default_expr,
+            Some("nextval('users_id_seq'::regclass)".into())
+        );
+        assert_eq!(user_id_col.is_nullable, false);
+        assert_eq!(user_id_col.is_primary_key, true);
+        assert_eq!(user_id_col.is_unique, true);
+        assert_eq!(user_id_col.varchar_length, None);
+
+        let user_name_col = cache.find_col("name", "users", None).unwrap();
+        assert_eq!(user_name_col.class_kind, ColumnClassKind::OrdinaryTable);
+        assert_eq!(user_name_col.comment, None);
+        assert_eq!(user_name_col.default_expr, None);
+        assert_eq!(user_name_col.is_nullable, false);
+        assert_eq!(user_name_col.is_primary_key, false);
+        assert_eq!(user_name_col.is_unique, false);
+        assert_eq!(user_name_col.varchar_length, Some(255));
+
+        let user_is_veg_col = cache.find_col("is_vegetarian", "users", None).unwrap();
+        assert_eq!(user_is_veg_col.class_kind, ColumnClassKind::OrdinaryTable);
+        assert_eq!(user_is_veg_col.comment, None);
+        assert_eq!(user_is_veg_col.default_expr, Some("false".into()));
+        assert_eq!(user_is_veg_col.is_nullable, true);
+        assert_eq!(user_is_veg_col.is_primary_key, false);
+        assert_eq!(user_is_veg_col.is_unique, false);
+        assert_eq!(user_is_veg_col.varchar_length, None);
+
+        let user_middle_name_col = cache.find_col("middle_name", "users", None).unwrap();
+        assert_eq!(
+            user_middle_name_col.class_kind,
+            ColumnClassKind::OrdinaryTable
+        );
+        assert_eq!(user_middle_name_col.comment, None);
+        assert_eq!(user_middle_name_col.default_expr, None);
+        assert_eq!(user_middle_name_col.is_nullable, true);
+        assert_eq!(user_middle_name_col.is_primary_key, false);
+        assert_eq!(user_middle_name_col.is_unique, false);
+        assert_eq!(user_middle_name_col.varchar_length, Some(255));
+
+        let properties_owner_id_col = cache
+            .find_col("owner_id", "properties", Some("real_estate"))
+            .unwrap();
+        assert_eq!(
+            properties_owner_id_col.class_kind,
+            ColumnClassKind::OrdinaryTable
+        );
+        assert_eq!(
+            properties_owner_id_col.comment,
+            Some("users might own many houses".into())
+        );
+        assert_eq!(properties_owner_id_col.is_nullable, true);
+        assert_eq!(properties_owner_id_col.is_primary_key, false);
+        assert_eq!(properties_owner_id_col.is_unique, false);
+        assert_eq!(properties_owner_id_col.varchar_length, None);
+    }
+}

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,60 @@
+with
+  available_tables as (
+    select
+      c.relname as table_name,
+      c.oid as table_oid,
+      c.relkind as class_kind,
+      n.nspname as schema_name
+    from
+      pg_catalog.pg_class c
+      join pg_catalog.pg_namespace n on n.oid = c.relnamespace
+    where
+      -- r: normal tables
+      -- v: views
+      -- m: materialized views
+      -- f: foreign tables
+      -- p: partitioned tables
+      c.relkind in ('r', 'v', 'm', 'f', 'p')
+  ),
+  available_indexes as (
+    select
+      unnest (ix.indkey) as attnum,
+      ix.indisprimary as is_primary,
+      ix.indisunique as is_unique,
+      ix.indrelid as table_oid
+    from
+      pg_catalog.pg_class c
+      join pg_catalog.pg_index ix on c.oid = ix.indexrelid
+    where
+      c.relkind = 'i'
+  )
+select
+  atts.attname as name,
+  ts.table_name,
+  ts.table_oid :: int8 as "table_oid!",
+  ts.class_kind :: char as "class_kind!",
+  ts.schema_name,
+  atts.atttypid :: int8 as "type_id!",
+  not atts.attnotnull as "is_nullable!",
+  nullif(
+    information_schema._pg_char_max_length (atts.atttypid, atts.atttypmod),
+    -1
+  ) as varchar_length,
+  pg_get_expr (def.adbin, def.adrelid) as default_expr,
+  coalesce(ix.is_primary, false) as "is_primary_key!",
+  coalesce(ix.is_unique, false) as "is_unique!",
+  pg_catalog.col_description (ts.table_oid, atts.attnum) as comment
+from
+  pg_catalog.pg_attribute atts
+  join available_tables ts on atts.attrelid = ts.table_oid
+  left join available_indexes ix on atts.attrelid = ix.table_oid
+  and atts.attnum = ix.attnum
+  left join pg_catalog.pg_attrdef def on atts.attrelid = def.adrelid
+  and atts.attnum = def.adnum
+where
+  -- system columns, such as `cmax` or `tableoid`, have negative `attnum`s
+  atts.attnum >= 0
+order by
+  schema_name desc,
+  table_name,
+  atts.attnum;

crates/pg_schema_cache/src/schema_cache.rs

@@ -1,5 +1,6 @@
 use sqlx::postgres::PgPool;
 
+use crate::columns::Column;
 use crate::functions::Function;
 use crate::schemas::Schema;
 use crate::tables::Table;
@@ -13,6 +14,7 @@ pub struct SchemaCache {
     pub functions: Vec<Function>,
     pub types: Vec<PostgresType>,
     pub versions: Vec<Version>,
+    pub columns: Vec<Column>,
 }
 
 impl SchemaCache {
@@ -21,12 +23,13 @@ impl SchemaCache {
     }
 
     pub async fn load(pool: &PgPool) -> Result<SchemaCache, sqlx::Error> {
-        let (schemas, tables, functions, types, versions) = futures_util::try_join!(
+        let (schemas, tables, functions, types, versions, columns) = futures_util::try_join!(
             Schema::load(pool),
             Table::load(pool),
             Function::load(pool),
             PostgresType::load(pool),
             Version::load(pool),
+            Column::load(pool)
         )?;
 
         Ok(SchemaCache {
@@ -35,6 +38,7 @@ impl SchemaCache {
             functions,
             types,
             versions,
+            columns,
         })
     }
 
@@ -58,6 +62,14 @@ impl SchemaCache {
             .find(|t| t.name == name && schema.is_none() || Some(t.schema.as_str()) == schema)
     }
 
+    pub fn find_col(&self, name: &str, table: &str, schema: Option<&str>) -> Option<&Column> {
+        self.columns.iter().find(|c| {
+            c.name.as_str() == name
+                && c.table_name.as_str() == table
+                && schema.is_none_or(|s| s == c.schema_name.as_str())
+        })
+    }
+
     pub fn find_types(&self, name: &str, schema: Option<&str>) -> Vec<&PostgresType> {
         self.types
             .iter()
@@ -74,17 +86,16 @@ pub trait SchemaCacheItem {
 
 #[cfg(test)]
 mod tests {
-    use async_std::task::block_on;
     use pg_test_utils::test_database::get_new_test_db;
 
     use crate::SchemaCache;
 
-    #[test]
-    fn test_schema_cache() {
-        let test_db = block_on(get_new_test_db());
-
-        block_on(SchemaCache::load(&test_db)).expect("Couldn't load Schema Cache");
+    #[tokio::test]
+    async fn it_loads() {
+        let test_db = get_new_test_db().await;
 
-        assert!(true);
+        SchemaCache::load(&test_db)
+            .await
+            .expect("Couldnt' load Schema Cache");
     }
 }