use crate::{Engine, TableState};
use std::sync::Arc;
use std::{cell::Cell, collections::BTreeMap};
use uqa_core::RelationIdentity;
use uqa_execution::catalog::security::{
table_grants::context::TableSecurityWrite,
table_inquiry::{TableColumnsRead, TablePrivilegeState},
table_ownership::{TableOwnerRegistry, TableOwnerSchema, TableOwnerState},
};
use uqa_execution::schema::sequences::role_ownership::{
OwnedSequenceSecurityCatalog, OwnedSequenceSecurityRead,
};
use uqa_sql::{ast::RelationPersistence, catalog::security::TableSecurity};
use uqa_sql::{catalog::security::SequenceSecurity, SQLError};
use uqa_storage::StorageBackendError;
use uqa_storage::StorageBackendResult;
#[derive(Debug, PartialEq, Clone)]
struct SecuritySnapshot {
table: TableSecurity,
sequences: BTreeMap<RelationIdentity, SequenceSecurity>,
}
fn snapshot(engine: &Engine) -> SecuritySnapshot {
SecuritySnapshot {
table: engine.storage.tables.read()[&RelationIdentity::new("public", "items")].security(),
sequences: engine.durable.sequence_security.read().clone(),
}
}
fn schema(engine: &Engine) -> serde_json::Value {
let table =
TableOwnerRegistry::table(engine, &RelationIdentity::new("public", "items")).unwrap();
let schema = table.schema();
serde_json::to_value((&schema.columns, &schema.constraints)).unwrap()
}
fn setup(engine: &Engine) {
engine.sql("CREATE ROLE target; CREATE ROLE reader; CREATE TABLE items(id serial PRIMARY KEY, generated_id bigint GENERATED BY DEFAULT AS IDENTITY, value integer DEFAULT 7 CHECK(value > 0), calculated integer GENERATED ALWAYS AS (value + 1) STORED, CONSTRAINT bounded CHECK(value < 100), UNIQUE(value)); GRANT SELECT, UPDATE(value) ON items TO reader WITH GRANT OPTION; GRANT USAGE ON SEQUENCE items_id_seq, items_generated_id_seq TO reader", &[]).unwrap();
}
struct FailedTables<'a> {
engine: &'a Engine,
before: SecuritySnapshot,
schema_before: serde_json::Value,
reached: Cell<bool>,
}
struct FailedTable<'a> {
table: Box<dyn TableOwnerState + 'a>,
observer: &'a FailedTables<'a>,
}
impl TableOwnerRegistry for FailedTables<'_> {
fn table(&self, relation: &RelationIdentity) -> Option<Box<dyn TableOwnerState + '_>> {
TableOwnerRegistry::table(self.engine, relation).map(|table| {
Box::new(FailedTable {
table,
observer: self,
}) as Box<dyn TableOwnerState>
})
}
}
impl TablePrivilegeState for FailedTable<'_> {
fn role_owner(&self) -> String {
self.table.role_owner()
}
fn security(&self) -> TableSecurity {
self.table.security()
}
fn columns(&self) -> TableColumnsRead<'_> {
self.table.columns()
}
fn column_names(&self) -> Vec<String> {
self.table.column_names()
}
}
impl TableOwnerState for FailedTable<'_> {
fn object_id(&self) -> [u8; 16] {
self.table.object_id()
}
fn persistence(&self) -> RelationPersistence {
self.table.persistence()
}
fn schema(&self) -> TableOwnerSchema {
self.table.schema()
}
fn security_write(&self) -> TableSecurityWrite<'_> {
self.table.security_write()
}
fn persist_schema(
&self,
_name: &str,
schema: &TableOwnerSchema,
security: &TableSecurity,
) -> StorageBackendResult<()> {
assert_eq!(security.role_owner, "target");
assert_eq!(
serde_json::to_value((&schema.columns, &schema.constraints)).unwrap(),
self.observer.schema_before
);
assert_eq!(
snapshot(self.observer.engine),
self.observer.before,
"no loaded state may be published before table persistence succeeds"
);
let catalog = self.observer.engine.storage.catalog.as_ref().unwrap();
let sequences = catalog.load_sequence_rows().unwrap();
assert_eq!(sequences.len(), 2);
assert!(
sequences.iter().all(|row| row.role_owner == "target"),
"all owned sequences must be saved before the table"
);
assert_eq!(catalog.load_tables().unwrap()[0].role_owner, "uqa");
self.observer.reached.set(true);
Err(StorageBackendError::Other(
"injected table ownership persistence failure".into(),
))
}
}
#[test]
fn failed_table_owner_save_rolls_back_prior_sequence_writes_without_publishing_loaded_security() {
let directory = tempfile::tempdir().unwrap();
let path = directory.path().join("ownership.db");
let engine = Engine::open(&path).unwrap();
setup(&engine);
let before = snapshot(&engine);
let schema_before = schema(&engine);
{
let tables = FailedTables {
engine: &engine,
before: before.clone(),
schema_before: schema_before.clone(),
reached: Cell::new(false),
};
let error = engine
.with_implicit_transaction(|engine| {
let mut context = engine.table_ownership_context();
context.tables = &tables;
context.alter_table_role_owner("public.items", "target")
})
.unwrap_err();
assert!(
matches!(&error, SQLError::Internal(message) if message.starts_with("persist table owner: ") && message.contains("injected table ownership persistence failure"))
);
assert!(tables.reached.get());
assert_eq!(snapshot(&engine), before);
assert_eq!(schema(&engine), schema_before);
assert!(engine
.storage
.catalog
.as_ref()
.unwrap()
.load_sequence_rows()
.unwrap()
.iter()
.all(|row| row.role_owner == "uqa"));
}
drop(engine);
let reopened = Engine::open(&path).unwrap();
assert_eq!(snapshot(&reopened), before);
assert_eq!(schema(&reopened), schema_before);
}
#[test]
fn table_owner_transfer_preserves_schema_and_durable_table_column_and_sequence_grants() {
let directory = tempfile::tempdir().unwrap();
let path = directory.path().join("grants.db");
let engine = Engine::open(&path).unwrap();
setup(&engine);
let before = snapshot(&engine);
let schema_before = schema(&engine);
engine
.sql("ALTER TABLE items OWNER TO target", &[])
.unwrap();
let after = snapshot(&engine);
assert_eq!(after.table.role_owner, "target");
assert_eq!(after.sequences.len(), 2);
for entry in after.table.acl.as_ref().unwrap() {
assert_eq!(entry.grantor.as_deref(), Some("target"));
}
for entries in after.table.column_acls.values() {
for entry in entries {
assert_eq!(entry.grantor.as_deref(), Some("target"));
}
}
for security in after.sequences.values() {
assert_eq!(security.role_owner, "target");
assert!(security
.acl
.as_ref()
.unwrap()
.iter()
.all(|entry| entry.grantor.as_deref() == Some("target")));
}
assert_eq!(
after.table.acl.as_ref().unwrap().len(),
before.table.acl.as_ref().unwrap().len()
);
assert_eq!(
after.table.column_acls.keys().collect::<Vec<_>>(),
before.table.column_acls.keys().collect::<Vec<_>>()
);
assert_eq!(schema(&engine), schema_before);
engine
.sql("INSERT INTO items(value) VALUES(9)", &[])
.unwrap();
engine.sql("SET ROLE reader; SELECT id, generated_id, value, calculated FROM items; UPDATE items SET value = 10; RESET ROLE", &[]).unwrap();
assert_eq!(
engine
.sql("SELECT calculated FROM items", &[])
.unwrap()
.rows[0]["calculated"],
uqa_core::Value::Int(11)
);
drop(engine);
let reopened = Engine::open(&path).unwrap();
assert_eq!(snapshot(&reopened), after);
assert_eq!(schema(&reopened), schema_before);
}
struct ReplacingTables<'a> {
engine: &'a Engine,
replacement: Arc<TableState>,
}
impl TableOwnerRegistry for ReplacingTables<'_> {
fn table(&self, relation: &RelationIdentity) -> Option<Box<dyn TableOwnerState + '_>> {
let selected = TableOwnerRegistry::table(self.engine, relation);
self.engine
.storage
.tables
.try_write()
.expect("binding releases the real registry guard")
.insert(relation.clone(), self.replacement.clone());
selected
}
}
#[test]
fn owner_authorization_rebinds_after_retaining_the_table_generation() {
let engine = Engine::new();
engine.sql("CREATE ROLE target; CREATE TABLE items(id integer); CREATE TABLE replacement(other text); ALTER TABLE replacement OWNER TO target", &[]).unwrap();
let retained = engine.storage.tables.read()[&RelationIdentity::new("public", "items")].clone();
let replacement =
engine.storage.tables.read()[&RelationIdentity::new("public", "replacement")].clone();
let tables = ReplacingTables {
engine: &engine,
replacement,
};
let mut context = engine.table_ownership_context();
context.tables = &tables;
context
.alter_table_role_owner("public.items", "target")
.unwrap();
assert_eq!(retained.role_owner(), "uqa", "fresh owner authorization observes the replacement and returns before mutating the retained generation");
assert_eq!(
engine.storage.tables.read()[&RelationIdentity::new("public", "items")].role_owner(),
"target"
);
}
struct CountedSequences<'a> {
engine: &'a Engine,
reads: Cell<usize>,
}
impl OwnedSequenceSecurityCatalog for CountedSequences<'_> {
fn owned_sequences(&self, table_object_id: [u8; 16]) -> Vec<RelationIdentity> {
self.reads.set(self.reads.get() + 1);
OwnedSequenceSecurityCatalog::owned_sequences(self.engine, table_object_id)
}
fn security_registry(&self) -> OwnedSequenceSecurityRead<'_> {
self.reads.set(self.reads.get() + 1);
OwnedSequenceSecurityCatalog::security_registry(self.engine)
}
}
#[test]
fn unchanged_table_owner_skips_owned_sequence_catalogs_after_role_validation() {
let engine = Engine::new();
setup(&engine);
let before = snapshot(&engine);
let sequences = CountedSequences {
engine: &engine,
reads: Cell::new(0),
};
let mut context = engine.table_ownership_context();
context.owned_sequences = &sequences;
context
.alter_table_role_owner("public.items", "CURRENT_USER")
.unwrap();
assert_eq!(sequences.reads.get(), 0);
assert_eq!(snapshot(&engine), before);
let error: SQLError = context
.alter_table_role_owner("public.items", "missing_owner")
.unwrap_err();
assert_eq!(error.sqlstate(), Some("42704"));
assert_eq!(sequences.reads.get(), 0);
}