Skip to content

Commit fd8f91d

Browse files
check tests
1 parent c08ade4 commit fd8f91d

File tree

2 files changed

+223
-0
lines changed

2 files changed

+223
-0
lines changed
Lines changed: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,212 @@
1+
use crate::schema_cache::SchemaCacheItem;
2+
3+
#[derive(Debug, Clone, PartialEq, Eq)]
4+
pub enum PolicyCommand {
5+
Select,
6+
Insert,
7+
Update,
8+
Delete,
9+
All,
10+
}
11+
12+
impl From<&str> for PolicyCommand {
13+
fn from(value: &str) -> Self {
14+
match value {
15+
"SELECT" => PolicyCommand::Select,
16+
"INSERT" => PolicyCommand::Insert,
17+
"UPDATE" => PolicyCommand::Update,
18+
"DELETE" => PolicyCommand::Delete,
19+
"ALL" => PolicyCommand::All,
20+
_ => panic!("Invalid Policy Command {}", value),
21+
}
22+
}
23+
}
24+
impl From<String> for PolicyCommand {
25+
fn from(value: String) -> Self {
26+
PolicyCommand::from(value.as_str())
27+
}
28+
}
29+
30+
#[derive(Debug, Clone, PartialEq, Eq)]
31+
struct PolicyQueried {
32+
name: String,
33+
table_name: String,
34+
schema_name: String,
35+
is_permissive: String,
36+
command: String,
37+
role_names: Option<Vec<String>>,
38+
security_qualification: Option<String>,
39+
with_check: Option<String>,
40+
}
41+
42+
impl From<PolicyQueried> for Policy {
43+
fn from(value: PolicyQueried) -> Self {
44+
Self {
45+
name: value.name,
46+
table_name: value.table_name,
47+
schema_name: value.schema_name,
48+
is_permissive: value.is_permissive == "PERMISSIVE",
49+
command: PolicyCommand::from(value.command),
50+
role_names: value.role_names.unwrap_or(vec![]),
51+
security_qualification: value.security_qualification,
52+
with_check: value.with_check,
53+
}
54+
}
55+
}
56+
57+
#[derive(Debug, Clone, PartialEq, Eq)]
58+
pub struct Policy {
59+
name: String,
60+
table_name: String,
61+
schema_name: String,
62+
is_permissive: bool,
63+
command: PolicyCommand,
64+
role_names: Vec<String>,
65+
security_qualification: Option<String>,
66+
with_check: Option<String>,
67+
}
68+
69+
impl SchemaCacheItem for Policy {
70+
type Item = Policy;
71+
72+
async fn load(pool: &sqlx::PgPool) -> Result<Vec<Self::Item>, sqlx::Error> {
73+
let policies = sqlx::query_file_as!(PolicyQueried, "src/queries/policies.sql")
74+
.fetch_all(pool)
75+
.await?;
76+
77+
Ok(policies.into_iter().map(Policy::from).collect())
78+
}
79+
}
80+
81+
#[cfg(test)]
82+
mod tests {
83+
use pgt_test_utils::test_database::get_new_test_db;
84+
use sqlx::Executor;
85+
86+
use crate::{SchemaCache, policies::PolicyCommand};
87+
88+
#[tokio::test]
89+
async fn loads_policies() {
90+
let test_db = get_new_test_db().await;
91+
92+
let setup = r#"
93+
do $$
94+
begin
95+
if not exists (
96+
select from pg_catalog.pg_roles
97+
where rolname = 'admin'
98+
) then
99+
create role admin;
100+
end if;
101+
end $$;
102+
103+
104+
create table public.users (
105+
id serial primary key,
106+
name varchar(255) not null
107+
);
108+
109+
create policy public_policy
110+
on public.users
111+
for select
112+
to public
113+
using (true);
114+
115+
create policy admin_policy
116+
on public.users
117+
for all
118+
to admin
119+
with check (true);
120+
121+
do $$
122+
begin
123+
if not exists (
124+
select from pg_catalog.pg_roles
125+
where rolname = 'owner'
126+
) then
127+
create role owner;
128+
end if;
129+
end $$;
130+
131+
create schema real_estate;
132+
133+
create table real_estate.properties (
134+
id serial primary key,
135+
owner_id int not null
136+
);
137+
138+
create policy owner_policy
139+
on real_estate.properties
140+
for update
141+
to owner
142+
using (owner_id = current_user::int);
143+
"#;
144+
145+
test_db
146+
.execute(setup)
147+
.await
148+
.expect("Failed to setup test database");
149+
150+
let cache = SchemaCache::load(&test_db)
151+
.await
152+
.expect("Failed to load Schema Cache");
153+
154+
let public_policies = cache
155+
.policies
156+
.iter()
157+
.filter(|p| p.schema_name == "public")
158+
.count();
159+
160+
assert_eq!(public_policies, 2);
161+
162+
let real_estate_policies = cache
163+
.policies
164+
.iter()
165+
.filter(|p| p.schema_name == "real_estate")
166+
.count();
167+
168+
assert_eq!(real_estate_policies, 1);
169+
170+
let public_policy = cache
171+
.policies
172+
.iter()
173+
.find(|p| p.name == "public_policy")
174+
.unwrap();
175+
assert_eq!(public_policy.table_name, "users");
176+
assert_eq!(public_policy.schema_name, "public");
177+
assert_eq!(public_policy.is_permissive, true);
178+
assert_eq!(public_policy.command, PolicyCommand::Select);
179+
assert_eq!(public_policy.role_names, vec!["public"]);
180+
assert_eq!(public_policy.security_qualification, Some("true".into()));
181+
assert_eq!(public_policy.with_check, None);
182+
183+
let admin_policy = cache
184+
.policies
185+
.iter()
186+
.find(|p| p.name == "admin_policy")
187+
.unwrap();
188+
assert_eq!(admin_policy.table_name, "users");
189+
assert_eq!(admin_policy.schema_name, "public");
190+
assert_eq!(admin_policy.is_permissive, true);
191+
assert_eq!(admin_policy.command, PolicyCommand::All);
192+
assert_eq!(admin_policy.role_names, vec!["admin"]);
193+
assert_eq!(admin_policy.security_qualification, None);
194+
assert_eq!(admin_policy.with_check, Some("true".into()));
195+
196+
let owner_policy = cache
197+
.policies
198+
.iter()
199+
.find(|p| p.name == "owner_policy")
200+
.unwrap();
201+
assert_eq!(owner_policy.table_name, "properties");
202+
assert_eq!(owner_policy.schema_name, "real_estate");
203+
assert_eq!(owner_policy.is_permissive, true);
204+
assert_eq!(owner_policy.command, PolicyCommand::Update);
205+
assert_eq!(owner_policy.role_names, vec!["owner"]);
206+
assert_eq!(
207+
owner_policy.security_qualification,
208+
Some("(owner_id = (CURRENT_USER)::integer)".into())
209+
);
210+
assert_eq!(owner_policy.with_check, None);
211+
}
212+
}
Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
select
2+
schemaname as "schema_name!",
3+
tablename as "table_name!",
4+
policyname as "name!",
5+
permissive as "is_permissive!",
6+
roles as "role_names!",
7+
cmd as "command!",
8+
qual as "security_qualification",
9+
with_check
10+
from
11+
pg_catalog.pg_policies;

0 commit comments

Comments
 (0)