chore(schema_cache): add query for roles (#404)

Author: juleswritescode

.sqlx/query-5e12c1d242ea9fcc68c20807b72300a7e131d2fb17fc74bd7f40a60b68df56c0.json

@@ -5,6 +5,7 @@
 mod columns;
 mod functions;
 mod policies;
+mod roles;
 mod schema_cache;
 mod schemas;
 mod tables;
@@ -15,6 +16,7 @@ mod versions;
 pub use columns::*;
 pub use functions::{Behavior, Function, FunctionArg, FunctionArgs};
 pub use policies::{Policy, PolicyCommand};
+pub use roles::*;
 pub use schema_cache::SchemaCache;
 pub use schemas::Schema;
 pub use tables::{ReplicaIdentity, Table};

crates/pgt_schema_cache/src/lib.rs

@@ -0,0 +1,7 @@
+select 
+  rolname as "name!", 
+  rolsuper as "is_super_user!", 
+  rolcreatedb as "can_create_db!", 
+  rolcanlogin  as "can_login!",
+  rolbypassrls as "can_bypass_rls!"
+from pg_catalog.pg_roles;

crates/pgt_schema_cache/src/queries/roles.sql

@@ -0,0 +1,85 @@
+use crate::schema_cache::SchemaCacheItem;
+
+#[derive(Debug, PartialEq, Eq)]
+pub struct Role {
+    pub name: String,
+    pub is_super_user: bool,
+    pub can_create_db: bool,
+    pub can_login: bool,
+    pub can_bypass_rls: bool,
+}
+
+impl SchemaCacheItem for Role {
+    type Item = Role;
+
+    async fn load(pool: &sqlx::PgPool) -> Result<Vec<Self::Item>, sqlx::Error> {
+        sqlx::query_file_as!(Role, "src/queries/roles.sql")
+            .fetch_all(pool)
+            .await
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::SchemaCache;
+    use pgt_test_utils::test_database::get_new_test_db;
+    use sqlx::Executor;
+
+    #[tokio::test]
+    async fn loads_roles() {
+        let test_db = get_new_test_db().await;
+
+        let setup = r#"
+            do $$
+            begin
+                if not exists (
+                    select from pg_catalog.pg_roles
+                    where rolname = 'test_super'
+                ) then
+                    create role test_super superuser createdb login bypassrls;
+                end if;
+                if not exists (
+                    select from pg_catalog.pg_roles
+                    where rolname = 'test_nologin'
+                ) then
+                    create role test_nologin;
+                end if;
+                if not exists (
+                    select from pg_catalog.pg_roles
+                    where rolname = 'test_login'
+                ) then
+                    create role test_login login;
+                end if;
+            end $$;
+        "#;
+
+        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 roles = &cache.roles;
+
+        let super_role = roles.iter().find(|r| r.name == "test_super").unwrap();
+        assert!(super_role.is_super_user);
+        assert!(super_role.can_create_db);
+        assert!(super_role.can_login);
+        assert!(super_role.can_bypass_rls);
+
+        let nologin_role = roles.iter().find(|r| r.name == "test_nologin").unwrap();
+        assert!(!nologin_role.is_super_user);
+        assert!(!nologin_role.can_create_db);
+        assert!(!nologin_role.can_login);
+        assert!(!nologin_role.can_bypass_rls);
+
+        let login_role = roles.iter().find(|r| r.name == "test_login").unwrap();
+        assert!(!login_role.is_super_user);
+        assert!(!login_role.can_create_db);
+        assert!(login_role.can_login);
+        assert!(!login_role.can_bypass_rls);
+    }
+}

crates/pgt_schema_cache/src/roles.rs

@@ -1,13 +1,13 @@
 use sqlx::postgres::PgPool;
 
-use crate::Trigger;
 use crate::columns::Column;
 use crate::functions::Function;
 use crate::policies::Policy;
 use crate::schemas::Schema;
 use crate::tables::Table;
 use crate::types::PostgresType;
 use crate::versions::Version;
+use crate::{Role, Trigger};
 
 #[derive(Debug, Default)]
 pub struct SchemaCache {
@@ -19,11 +19,12 @@ pub struct SchemaCache {
     pub columns: Vec<Column>,
     pub policies: Vec<Policy>,
     pub triggers: Vec<Trigger>,
+    pub roles: Vec<Role>,
 }
 
 impl SchemaCache {
     pub async fn load(pool: &PgPool) -> Result<SchemaCache, sqlx::Error> {
-        let (schemas, tables, functions, types, versions, columns, policies, triggers) = futures_util::try_join!(
+        let (schemas, tables, functions, types, versions, columns, policies, triggers, roles) = futures_util::try_join!(
             Schema::load(pool),
             Table::load(pool),
             Function::load(pool),
@@ -32,6 +33,7 @@ impl SchemaCache {
             Column::load(pool),
             Policy::load(pool),
             Trigger::load(pool),
+            Role::load(pool)
         )?;
 
         Ok(SchemaCache {
@@ -43,6 +45,7 @@ impl SchemaCache {
             columns,
             policies,
             triggers,
+            roles,
         })
     }