diff --git a/Cargo.lock b/Cargo.lock index f6a2c53a..09684985 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2570,6 +2570,7 @@ dependencies = [ "serde", "serde_json", "sqlx", + "tokio", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 0f4a4195..d804b35f 100644 --- a/Cargo.toml +++ b/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"] } diff --git a/crates/pg_schema_cache/Cargo.toml b/crates/pg_schema_cache/Cargo.toml index 534a6c16..f06520b8 100644 --- a/crates/pg_schema_cache/Cargo.toml +++ b/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 diff --git a/crates/pg_schema_cache/src/columns.rs b/crates/pg_schema_cache/src/columns.rs new file mode 100644 index 00000000..1ff4a412 --- /dev/null +++ b/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 for ColumnClassKind { + fn from(value: String) -> Self { + ColumnClassKind::from(value.as_str()) + } +} + +impl From 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, + + pub varchar_length: Option, + + /// Comment inserted via `COMMENT ON COLUMN my_table.my_comment '...'`, if present. + pub comment: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ForeignKeyReference { + pub schema: Option, + pub table: String, + pub column: String, +} + +impl SchemaCacheItem for Column { + type Item = Column; + + async fn load(pool: &sqlx::PgPool) -> Result, 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); + } +} diff --git a/crates/pg_schema_cache/src/lib.rs b/crates/pg_schema_cache/src/lib.rs index ec1596c3..719da404 100644 --- a/crates/pg_schema_cache/src/lib.rs +++ b/crates/pg_schema_cache/src/lib.rs @@ -2,6 +2,7 @@ #![allow(dead_code)] +mod columns; mod functions; mod schema_cache; mod schemas; diff --git a/crates/pg_schema_cache/src/queries/columns.sql b/crates/pg_schema_cache/src/queries/columns.sql new file mode 100644 index 00000000..86df7cf4 --- /dev/null +++ b/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; \ No newline at end of file diff --git a/crates/pg_schema_cache/src/schema_cache.rs b/crates/pg_schema_cache/src/schema_cache.rs index b9762122..8d73e631 100644 --- a/crates/pg_schema_cache/src/schema_cache.rs +++ b/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, pub types: Vec, pub versions: Vec, + pub columns: Vec, } impl SchemaCache { @@ -21,12 +23,13 @@ impl SchemaCache { } pub async fn load(pool: &PgPool) -> Result { - 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"); } }