add test

Author: juleswritescode

Cargo.lock

@@ -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"] }

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/Cargo.toml

@@ -1,6 +1,7 @@
 use crate::schema_cache::SchemaCacheItem;
 
-enum ColumnClassKind {
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum ColumnClassKind {
     OrdinaryTable,
     View,
     MaterializedView,
@@ -30,19 +31,23 @@ impl From<String> for ColumnClassKind {
     }
 }
 
+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: bool,
+    pub class_kind: ColumnClassKind,
 
     pub schema_name: String,
     pub type_id: i64,
-
     pub is_nullable: bool,
 
     pub is_primary_key: bool,
@@ -52,8 +57,7 @@ pub struct Column {
     pub default_expr: Option<String>,
 
     pub varchar_length: Option<i32>,
-    // /// None if the column is not a foreign key.
-    // pub foreign_key: Option<ForeignKeyReference>,
+
     /// Comment inserted via `COMMENT ON COLUMN my_table.my_comment '...'`, if present.
     pub comment: Option<String>,
 }
@@ -74,3 +78,125 @@ impl SchemaCacheItem for Column {
             .await
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use pg_test_utils::test_database::get_new_test_db;
+    use sqlx::Executor;
+
+    use crate::{columns::ColumnClassKind, schema_cache, 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/columns.rs

@@ -31,19 +31,18 @@ with
 select
   atts.attname as name,
   ts.table_name,
-  ts.table_oid,
-  ts.class_kind,
+  ts.table_oid :: int8 as "table_oid!",
+  ts.class_kind :: char as "class_kind!",
   ts.schema_name,
-  atts.attnum,
-  atts.atttypid as type_id,
-  not atts.attnotnull as is_nullable,
+  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,
+  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
@@ -58,4 +57,4 @@ where
 order by
   schema_name desc,
   table_name,
-  attnum;
+  atts.attnum;

crates/pg_schema_cache/src/queries/columns.sql

@@ -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,8 +62,12 @@ 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_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> {
@@ -78,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");
     }
 }