check tests

Author: juleswritescode

crates/pgt_schema_cache/src/policies.rs

@@ -0,0 +1,212 @@
+use crate::schema_cache::SchemaCacheItem;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum PolicyCommand {
+    Select,
+    Insert,
+    Update,
+    Delete,
+    All,
+}
+
+impl From<&str> for PolicyCommand {
+    fn from(value: &str) -> Self {
+        match value {
+            "SELECT" => PolicyCommand::Select,
+            "INSERT" => PolicyCommand::Insert,
+            "UPDATE" => PolicyCommand::Update,
+            "DELETE" => PolicyCommand::Delete,
+            "ALL" => PolicyCommand::All,
+            _ => panic!("Invalid Policy Command {}", value),
+        }
+    }
+}
+impl From<String> for PolicyCommand {
+    fn from(value: String) -> Self {
+        PolicyCommand::from(value.as_str())
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct PolicyQueried {
+    name: String,
+    table_name: String,
+    schema_name: String,
+    is_permissive: String,
+    command: String,
+    role_names: Option<Vec<String>>,
+    security_qualification: Option<String>,
+    with_check: Option<String>,
+}
+
+impl From<PolicyQueried> for Policy {
+    fn from(value: PolicyQueried) -> Self {
+        Self {
+            name: value.name,
+            table_name: value.table_name,
+            schema_name: value.schema_name,
+            is_permissive: value.is_permissive == "PERMISSIVE",
+            command: PolicyCommand::from(value.command),
+            role_names: value.role_names.unwrap_or(vec![]),
+            security_qualification: value.security_qualification,
+            with_check: value.with_check,
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Policy {
+    name: String,
+    table_name: String,
+    schema_name: String,
+    is_permissive: bool,
+    command: PolicyCommand,
+    role_names: Vec<String>,
+    security_qualification: Option<String>,
+    with_check: Option<String>,
+}
+
+impl SchemaCacheItem for Policy {
+    type Item = Policy;
+
+    async fn load(pool: &sqlx::PgPool) -> Result<Vec<Self::Item>, sqlx::Error> {
+        let policies = sqlx::query_file_as!(PolicyQueried, "src/queries/policies.sql")
+            .fetch_all(pool)
+            .await?;
+
+        Ok(policies.into_iter().map(Policy::from).collect())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use pgt_test_utils::test_database::get_new_test_db;
+    use sqlx::Executor;
+
+    use crate::{SchemaCache, policies::PolicyCommand};
+
+    #[tokio::test]
+    async fn loads_policies() {
+        let test_db = get_new_test_db().await;
+
+        let setup = r#"
+            do $$
+            begin
+                if not exists (
+                    select from pg_catalog.pg_roles
+                    where rolname = 'admin'
+                ) then
+                    create role admin;
+                end if;
+            end $$;
+
+
+            create table public.users (
+                id serial primary key,
+                name varchar(255) not null
+            );
+
+            create policy public_policy
+                on public.users
+                for select
+                to public
+                using (true);
+
+            create policy admin_policy
+                on public.users
+                for all
+                to admin
+                with check (true);
+
+            do $$
+            begin
+                if not exists (
+                    select from pg_catalog.pg_roles
+                    where rolname = 'owner'
+                ) then
+                    create role owner;
+                end if;
+            end $$;
+
+            create schema real_estate;
+
+            create table real_estate.properties (
+                id serial primary key,
+                owner_id int not null
+            );
+
+            create policy owner_policy
+                on real_estate.properties
+                for update
+                to owner
+                using (owner_id = current_user::int);
+        "#;
+
+        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_policies = cache
+            .policies
+            .iter()
+            .filter(|p| p.schema_name == "public")
+            .count();
+
+        assert_eq!(public_policies, 2);
+
+        let real_estate_policies = cache
+            .policies
+            .iter()
+            .filter(|p| p.schema_name == "real_estate")
+            .count();
+
+        assert_eq!(real_estate_policies, 1);
+
+        let public_policy = cache
+            .policies
+            .iter()
+            .find(|p| p.name == "public_policy")
+            .unwrap();
+        assert_eq!(public_policy.table_name, "users");
+        assert_eq!(public_policy.schema_name, "public");
+        assert_eq!(public_policy.is_permissive, true);
+        assert_eq!(public_policy.command, PolicyCommand::Select);
+        assert_eq!(public_policy.role_names, vec!["public"]);
+        assert_eq!(public_policy.security_qualification, Some("true".into()));
+        assert_eq!(public_policy.with_check, None);
+
+        let admin_policy = cache
+            .policies
+            .iter()
+            .find(|p| p.name == "admin_policy")
+            .unwrap();
+        assert_eq!(admin_policy.table_name, "users");
+        assert_eq!(admin_policy.schema_name, "public");
+        assert_eq!(admin_policy.is_permissive, true);
+        assert_eq!(admin_policy.command, PolicyCommand::All);
+        assert_eq!(admin_policy.role_names, vec!["admin"]);
+        assert_eq!(admin_policy.security_qualification, None);
+        assert_eq!(admin_policy.with_check, Some("true".into()));
+
+        let owner_policy = cache
+            .policies
+            .iter()
+            .find(|p| p.name == "owner_policy")
+            .unwrap();
+        assert_eq!(owner_policy.table_name, "properties");
+        assert_eq!(owner_policy.schema_name, "real_estate");
+        assert_eq!(owner_policy.is_permissive, true);
+        assert_eq!(owner_policy.command, PolicyCommand::Update);
+        assert_eq!(owner_policy.role_names, vec!["owner"]);
+        assert_eq!(
+            owner_policy.security_qualification,
+            Some("(owner_id = (CURRENT_USER)::integer)".into())
+        );
+        assert_eq!(owner_policy.with_check, None);
+    }
+}

crates/pgt_schema_cache/src/queries/policies.sql

@@ -0,0 +1,11 @@
+select 
+  schemaname as "schema_name!", 
+  tablename as "table_name!", 
+  policyname as "name!", 
+  permissive as "is_permissive!", 
+  roles as "role_names!", 
+  cmd as "command!", 
+  qual as "security_qualification", 
+  with_check
+from 
+ pg_catalog.pg_policies;