use super::{coordination::after_wait, identity::reopen};
use crate::{
tests::relation_lock_support::{sessions, sql},
Engine,
};
use uqa_core::Value;
use uqa_execution::{
catalog::security::roles::locking::ROLE_CATALOG_CLASS_ID,
row_locks::{
shared_objects::{SharedCatalogLock, SharedObjectLockSession},
RelationLockMode,
},
};
struct Target {
setup: &'static str,
grant: &'static str,
access: Access,
}
enum Access {
Inquiry(&'static str),
Execute(&'static str),
}
const TARGETS: &[Target] = &[
Target { setup: "", grant: "GRANT SELECT ON t TO dependent", access: Access::Inquiry("has_table_privilege('dependent', 't', 'SELECT')") },
Target { setup: "", grant: "GRANT SELECT(v) ON t TO dependent", access: Access::Inquiry("has_column_privilege('dependent', 't', 'v', 'SELECT')") },
Target { setup: "", grant: "GRANT UPDATE ON pg_class TO dependent WITH GRANT OPTION", access: Access::Inquiry("has_table_privilege('dependent', 'pg_class', 'UPDATE WITH GRANT OPTION')") },
Target { setup: "", grant: "GRANT UPDATE(relname) ON pg_class TO dependent", access: Access::Inquiry("has_column_privilege('dependent', 'pg_class', 'relname', 'UPDATE')") },
Target { setup: "CREATE VIEW role_view AS SELECT v FROM t", grant: "GRANT SELECT ON role_view TO dependent", access: Access::Inquiry("has_table_privilege('dependent', 'role_view', 'SELECT')") },
Target { setup: "CREATE MATERIALIZED VIEW role_view AS SELECT v FROM t", grant: "GRANT SELECT ON role_view TO dependent", access: Access::Inquiry("has_table_privilege('dependent', 'role_view', 'SELECT')") },
Target { setup: "CREATE SERVER role_remote FOREIGN DATA WRAPPER memory_fdw; CREATE FOREIGN TABLE role_foreign(v integer) SERVER role_remote", grant: "GRANT SELECT ON role_foreign TO dependent", access: Access::Inquiry("has_table_privilege('dependent', 'role_foreign', 'SELECT')") },
Target { setup: "", grant: "GRANT CREATE ON DATABASE uqa TO dependent", access: Access::Inquiry("has_database_privilege('dependent', 'uqa', 'CREATE')") },
Target { setup: "CREATE SCHEMA role_schema", grant: "GRANT USAGE ON SCHEMA role_schema TO dependent", access: Access::Inquiry("has_schema_privilege('dependent', 'role_schema', 'USAGE')") },
Target { setup: "CREATE SEQUENCE role_sequence", grant: "GRANT USAGE ON SEQUENCE role_sequence TO dependent", access: Access::Inquiry("has_sequence_privilege('dependent', 'role_sequence', 'USAGE')") },
Target { setup: "CREATE SEQUENCE role_sequence", grant: "GRANT SELECT ON TABLE role_sequence TO dependent", access: Access::Inquiry("has_sequence_privilege('dependent', 'role_sequence', 'SELECT')") },
Target { setup: "CREATE FUNCTION role_function() RETURNS integer LANGUAGE SQL AS 'SELECT 1'; REVOKE EXECUTE ON FUNCTION role_function() FROM PUBLIC", grant: "GRANT EXECUTE ON FUNCTION role_function() TO dependent", access: Access::Execute("SELECT role_function()") },
Target { setup: "CREATE PROCEDURE role_procedure() LANGUAGE SQL AS 'SELECT 1'; REVOKE EXECUTE ON PROCEDURE role_procedure() FROM PUBLIC", grant: "GRANT EXECUTE ON PROCEDURE role_procedure() TO dependent", access: Access::Execute("CALL role_procedure()") },
];
fn role_lock(engine: &Engine) -> SharedCatalogLock<'static> {
SharedCatalogLock::Object {
class_id: ROLE_CATALOG_CLASS_ID,
oid: u32::try_from(engine.durable.roles.read()["dependent"].oid).unwrap(),
}
}
fn role_exists(engine: &Engine) -> bool {
sql(
engine,
"SELECT count(*) AS n FROM pg_roles WHERE rolname = 'dependent'",
)
.rows[0]["n"]
== Value::Int(1)
}
fn allowed(engine: &Engine, target: &Target) -> bool {
allowed_as(engine, target, "dependent")
}
fn allowed_as(engine: &Engine, target: &Target, role: &str) -> bool {
match target.access {
Access::Inquiry(inquiry) => {
let inquiry = inquiry.replace("'dependent'", &format!("'{role}'"));
sql(engine, &format!("SELECT {inquiry} AS allowed")).rows[0]["allowed"]
== Value::Bool(true)
}
Access::Execute(statement) => {
sql(engine, &format!("SET ROLE {role}"));
let result = engine.sql(statement, &[]);
sql(engine, "RESET ROLE");
match result {
Ok(_) => true,
Err(error) => {
assert_eq!(error.sqlstate(), Some("42501"), "{statement}: {error}");
false
}
}
}
}
}
#[rstest::rstest]
#[case::native_sqlite(0)]
#[case::sqlite_key_value(1)]
#[case::redb(2)]
fn acl_waits_keep_the_original_recipient_after_rename_and_name_reuse(
#[case] provider: usize,
#[values("READ COMMITTED", "REPEATABLE READ", "SERIALIZABLE")] isolation: &str,
#[values(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12)] target_index: usize,
) {
let target = &TARGETS[target_index];
let (directory, first, second) = sessions(provider);
sql(&first, "CREATE ROLE dependent");
if !target.setup.is_empty() {
sql(&first, target.setup);
}
let original = first.durable.roles.read()["dependent"].identity();
let lock = role_lock(&first);
sql(&first, "BEGIN");
first
.acquire_shared_catalog(lock, RelationLockMode::AccessExclusive)
.unwrap()
.retain();
sql(
&second,
&format!("BEGIN ISOLATION LEVEL {isolation}; SELECT 1"),
);
let (second, result) = after_wait(
&first,
second,
target.grant,
lock,
"ALTER ROLE dependent RENAME TO renamed; CREATE ROLE dependent; COMMIT",
);
result.unwrap();
sql(&second, "COMMIT");
assert!(
allowed_as(&first, target, "renamed"),
"{provider}/{isolation}/{}",
target.grant
);
assert!(
!allowed(&first, target),
"{provider}/{isolation}/{}",
target.grant
);
assert_eq!(first.durable.roles.read()["renamed"].identity(), original);
drop(second);
drop(first);
let restored = reopen(provider, &directory.path().join("table-locks.db"));
assert!(allowed_as(&restored, target, "renamed"));
assert!(!allowed(&restored, target));
}
#[rstest::rstest]
#[case::native_sqlite(0)]
#[case::sqlite_key_value(1)]
#[case::redb(2)]
fn role_drop_waits_for_added_acl_dependencies(
#[case] provider: usize,
#[values("READ COMMITTED", "REPEATABLE READ", "SERIALIZABLE")] isolation: &str,
#[values(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12)] target_index: usize,
#[values("COMMIT", "ROLLBACK", "ROLLBACK TO before_grant; COMMIT")] finish: &str,
) {
let target = &TARGETS[target_index];
let (directory, first, second) = sessions(provider);
sql(&first, "CREATE ROLE dependent");
if !target.setup.is_empty() {
sql(&first, target.setup);
}
let lock = role_lock(&first);
sql(&first, "BEGIN; SAVEPOINT before_grant");
sql(&first, target.grant);
sql(
&second,
&format!("BEGIN ISOLATION LEVEL {isolation}; SELECT 1"),
);
let (second, result) = after_wait(&first, second, "DROP ROLE dependent", lock, finish);
if finish == "COMMIT" {
assert_eq!(result.unwrap_err().sqlstate(), Some("2BP01"));
sql(&second, "ROLLBACK");
assert!(
allowed(&first, target),
"provider {provider}, {isolation}, {}, {finish}",
target.grant
);
} else {
result.unwrap();
sql(&second, "COMMIT");
}
assert_eq!(role_exists(&first), finish == "COMMIT");
drop(second);
drop(first);
let restored = reopen(provider, &directory.path().join("table-locks.db"));
assert_eq!(role_exists(&restored), finish == "COMMIT");
}
#[rstest::rstest]
#[case::native_sqlite(0)]
#[case::sqlite_key_value(1)]
#[case::redb(2)]
fn acl_grants_wait_for_role_drop_and_never_rebind_a_recreated_name(
#[case] provider: usize,
#[values("READ COMMITTED", "REPEATABLE READ", "SERIALIZABLE")] isolation: &str,
#[values(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12)] target_index: usize,
#[values("COMMIT", "ROLLBACK", "ROLLBACK TO before_drop; COMMIT", "RECREATE")] finish: &str,
) {
let target = &TARGETS[target_index];
let (directory, first, second) = sessions(provider);
sql(&first, "CREATE ROLE dependent");
if !target.setup.is_empty() {
sql(&first, target.setup);
}
let lock = role_lock(&first);
let original = first.durable.roles.read()["dependent"].object_id;
sql(&first, "BEGIN; SAVEPOINT before_drop; DROP ROLE dependent");
if finish == "RECREATE" {
sql(&first, "CREATE ROLE dependent");
assert_ne!(first.durable.roles.read()["dependent"].object_id, original);
}
sql(
&second,
&format!("BEGIN ISOLATION LEVEL {isolation}; INSERT INTO t VALUES (2)"),
);
let (second, result) = after_wait(
&first,
second,
target.grant,
lock,
if finish == "RECREATE" {
"COMMIT"
} else {
finish
},
);
let granted = finish != "COMMIT" && finish != "RECREATE";
if granted {
result.unwrap();
sql(&second, "COMMIT");
assert!(
allowed(&first, target),
"provider {provider}, {isolation}, {}, {finish}",
target.grant
);
} else {
assert_eq!(result.unwrap_err().sqlstate(), Some("42704"));
sql(&second, "ROLLBACK");
if finish == "RECREATE" {
assert!(!allowed(&first, target));
}
}
assert_eq!(
sql(&first, "SELECT v FROM t").rows.len(),
if granted { 2 } else { 1 }
);
drop(second);
drop(first);
let restored = reopen(provider, &directory.path().join("table-locks.db"));
assert_eq!(role_exists(&restored), finish != "COMMIT");
if finish != "COMMIT" {
assert_eq!(allowed(&restored, target), granted);
}
}
#[rstest::rstest]
#[case::native_sqlite(0)]
#[case::sqlite_key_value(1)]
#[case::redb(2)]
fn unchanged_acl_roles_and_revoke_do_not_acquire_new_dependency_locks(
#[case] provider: usize,
#[values(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12)] target_index: usize,
) {
let target = &TARGETS[target_index];
let (_directory, first, second) = sessions(provider);
sql(&first, "CREATE ROLE dependent");
if !target.setup.is_empty() {
sql(&first, target.setup);
}
sql(&first, target.grant);
let revoke = target
.grant
.replacen("GRANT ", "REVOKE ", 1)
.replace(" TO dependent", " FROM dependent")
.replace(" WITH GRANT OPTION", "");
let key = first.row_locks.shared_catalog_key(role_lock(&first));
for statement in [target.grant, revoke.as_str()] {
sql(&first, "BEGIN");
sql(&first, statement);
assert!(
first
.row_locks
.try_acquire_relation(
second.session_id,
key,
RelationLockMode::AccessExclusive,
0,
&second.runtime.cancellation,
)
.unwrap(),
"unexpected dependency lock for {statement}"
);
first.row_locks.release_session(second.session_id);
sql(&first, "ROLLBACK");
}
}