pgevolve_core/diff/
grants.rs1use std::collections::BTreeSet;
9
10use crate::identifier::Identifier;
11use crate::ir::catalog::Catalog;
12use crate::ir::grant::{Grant, GrantTarget};
13
14#[must_use]
23pub fn diff_grants(
24 target: &[Grant],
25 source: &[Grant],
26 managed_roles: &BTreeSet<Identifier>,
27) -> (Vec<Grant>, Vec<Grant>, Vec<Grant>) {
28 let target_set: BTreeSet<&Grant> = target.iter().collect();
29 let source_set: BTreeSet<&Grant> = source.iter().collect();
30
31 let to_add: Vec<Grant> = source_set
32 .difference(&target_set)
33 .map(|g| (*g).clone())
34 .collect();
35
36 let mut to_revoke = Vec::new();
37 let mut unmanaged_observed = Vec::new();
38 for g in target_set.difference(&source_set) {
39 if grantee_is_managed(&g.grantee, managed_roles) {
40 to_revoke.push((*g).clone());
41 } else {
42 unmanaged_observed.push((*g).clone());
43 }
44 }
45 (to_add, to_revoke, unmanaged_observed)
46}
47
48fn grantee_is_managed(target: &GrantTarget, managed_roles: &BTreeSet<Identifier>) -> bool {
49 match target {
50 GrantTarget::Public => true,
51 GrantTarget::Role(name) => managed_roles.contains(name),
52 }
53}
54
55#[must_use]
59pub fn collect_managed_roles(cat: &Catalog) -> BTreeSet<Identifier> {
60 let mut out = BTreeSet::new();
61
62 for s in &cat.schemas {
63 collect_grants_into(&s.grants, &mut out);
64 if let Some(o) = &s.owner {
65 out.insert(o.clone());
66 }
67 }
68 for s in &cat.sequences {
69 collect_grants_into(&s.grants, &mut out);
70 if let Some(o) = &s.owner {
71 out.insert(o.clone());
72 }
73 }
74 for t in &cat.tables {
75 collect_grants_into(&t.grants, &mut out);
76 if let Some(o) = &t.owner {
77 out.insert(o.clone());
78 }
79 }
80 for v in &cat.views {
81 collect_grants_into(&v.grants, &mut out);
82 if let Some(o) = &v.owner {
83 out.insert(o.clone());
84 }
85 }
86 for m in &cat.materialized_views {
87 collect_grants_into(&m.grants, &mut out);
88 if let Some(o) = &m.owner {
89 out.insert(o.clone());
90 }
91 }
92 for f in &cat.functions {
93 collect_grants_into(&f.grants, &mut out);
94 if let Some(o) = &f.owner {
95 out.insert(o.clone());
96 }
97 }
98 for p in &cat.procedures {
99 collect_grants_into(&p.grants, &mut out);
100 if let Some(o) = &p.owner {
101 out.insert(o.clone());
102 }
103 }
104 for t in &cat.types {
105 collect_grants_into(&t.grants, &mut out);
106 if let Some(o) = &t.owner {
107 out.insert(o.clone());
108 }
109 }
110 for r in &cat.default_privileges {
111 out.insert(r.target_role.clone());
112 collect_grants_into(&r.grants, &mut out);
113 }
114
115 out
116}
117
118fn collect_grants_into(grants: &[Grant], set: &mut BTreeSet<Identifier>) {
120 for g in grants {
121 if let GrantTarget::Role(name) = &g.grantee {
122 set.insert(name.clone());
123 }
124 }
125}
126
127#[cfg(test)]
128mod tests {
129 use super::*;
130 use crate::ir::grant::{Grant, GrantTarget, Privilege};
131
132 fn id(s: &str) -> Identifier {
133 Identifier::from_unquoted(s).unwrap()
134 }
135
136 fn grant_to_role(role: &str, priv_: Privilege) -> Grant {
137 Grant {
138 grantee: GrantTarget::Role(id(role)),
139 privilege: priv_,
140 with_grant_option: false,
141 columns: None,
142 }
143 }
144
145 fn grant_public(priv_: Privilege) -> Grant {
146 Grant {
147 grantee: GrantTarget::Public,
148 privilege: priv_,
149 with_grant_option: false,
150 columns: None,
151 }
152 }
153
154 fn managed(roles: &[&str]) -> BTreeSet<Identifier> {
155 roles.iter().map(|r| id(r)).collect()
156 }
157
158 #[test]
159 fn empty_lists_yield_empty_diff() {
160 let (add, rev, unmanaged) = diff_grants(&[], &[], &managed(&[]));
161 assert!(add.is_empty());
162 assert!(rev.is_empty());
163 assert!(unmanaged.is_empty());
164 }
165
166 #[test]
167 fn add_only_grant_in_source_not_in_target() {
168 let g = grant_to_role("alice", Privilege::Select);
169 let (add, rev, unmanaged) =
170 diff_grants(&[], std::slice::from_ref(&g), &managed(&["alice"]));
171 assert_eq!(add, vec![g]);
172 assert!(rev.is_empty());
173 assert!(unmanaged.is_empty());
174 }
175
176 #[test]
177 fn revoke_managed_grantee() {
178 let g = grant_to_role("alice", Privilege::Select);
180 let (add, rev, unmanaged) =
181 diff_grants(std::slice::from_ref(&g), &[], &managed(&["alice"]));
182 assert!(add.is_empty());
183 assert_eq!(rev, vec![g]);
184 assert!(unmanaged.is_empty());
185 }
186
187 #[test]
188 fn ignore_unmanaged_grantee() {
189 let g = grant_to_role("external_role", Privilege::Select);
191 let (add, rev, unmanaged) = diff_grants(std::slice::from_ref(&g), &[], &managed(&[]));
192 assert!(add.is_empty());
193 assert!(rev.is_empty());
194 assert_eq!(unmanaged, vec![g]);
195 }
196
197 #[test]
198 fn public_always_managed() {
199 let g = grant_public(Privilege::Usage);
201 let (add, rev, unmanaged) = diff_grants(std::slice::from_ref(&g), &[], &managed(&[]));
202 assert!(add.is_empty());
203 assert_eq!(rev, vec![g]);
204 assert!(unmanaged.is_empty());
205 }
206}