use super::locking::{RoleBinding, RoleDependencyLocks, RoleLockContext};
use std::collections::{BTreeMap, BTreeSet};
use uqa_sql::{
catalog::roles::{
guards::{RoleDefinitionRead, RoleMembershipRead},
RoleDefinition, RoleMembership, RoleMembershipKey,
},
SQLError,
};
pub struct RoleDependencyCandidate<'a, T> {
pub value: T,
pub memberships: RoleMembershipRead<'a>,
pub roles: RoleDefinitionRead<'a>,
pub dependencies: BTreeSet<String>,
}
#[cfg(test)]
mod tests;
pub fn prepare_role_dependencies<'a, T>(
context: &RoleLockContext<'_>,
prepare_writer: impl FnMut() -> Result<(), SQLError>,
prepare: impl FnMut() -> Result<RoleDependencyCandidate<'a, T>, SQLError>,
) -> Result<RoleDependencyCandidate<'a, T>, SQLError> {
prepare_dependencies(context, prepare_writer, prepare, |_| true)
}
pub fn retain_created_owner(
context: RoleLockContext<'_>,
owner: &RoleBinding,
) -> Result<(), SQLError> {
drop(prepare_role_owner(
context,
owner,
|| Ok(()),
|_, _, _| Ok(Some(())),
)?);
Ok(())
}
pub fn prepare_role_owner<'a, T>(
context: RoleLockContext<'a>,
owner: &RoleBinding,
prepare_writer: impl FnMut() -> Result<(), SQLError>,
mut prepare: impl FnMut(
&BTreeMap<String, RoleDefinition>,
&BTreeMap<RoleMembershipKey, RoleMembership>,
&str,
) -> Result<Option<T>, SQLError>,
) -> Result<RoleDependencyCandidate<'a, Option<T>>, SQLError> {
prepare_dependencies(
&context,
prepare_writer,
|| {
let roles = context.roles.role_definitions();
owner.revalidate(&roles)?;
let memberships = context.roles.role_memberships();
let name = owner.require_name(&roles)?;
let value = prepare(&roles, &memberships, name)?;
let dependencies = if value.is_some() {
BTreeSet::from([name.to_owned()])
} else {
BTreeSet::new()
};
Ok(RoleDependencyCandidate {
value,
memberships,
roles,
dependencies,
})
},
Option::is_some,
)
}
fn prepare_dependencies<'a, T>(
context: &RoleLockContext<'_>,
mut prepare_writer: impl FnMut() -> Result<(), SQLError>,
mut prepare: impl FnMut() -> Result<RoleDependencyCandidate<'a, T>, SQLError>,
needs_writer: impl Fn(&T) -> bool,
) -> Result<RoleDependencyCandidate<'a, T>, SQLError> {
let mut held = RoleDependencyLocks::default();
let mut writer_prepared = false;
loop {
let candidate = prepare()?;
let pending = held.missing(&candidate.roles, &candidate.dependencies)?;
if !pending.is_empty() {
drop(candidate);
held.acquire(context, pending)?;
writer_prepared = false;
} else if writer_prepared || !needs_writer(&candidate.value) {
return Ok(candidate);
} else {
drop(candidate);
prepare_writer()?;
writer_prepared = true;
}
}
}