uqa-execution 0.4.6

Volcano physical operators with row-batch pipelines
//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

use super::*;
use std::{
    cell::{Cell, RefCell},
    collections::BTreeMap,
};
use uqa_sql::catalog::roles::RoleReference;
use uqa_sql::{
    ast::{DatabasePrivilege, DatabaseRevokeBehavior},
    catalog::roles::{
        guards::{RoleDefinitionRead, RoleMembershipRead},
        RoleDefinition, RoleMembership, RoleMembershipKey,
    },
};

struct DatabaseCatalog {
    security: RefCell<BoundDatabaseSecurity>,
    roles: RefCell<BTreeMap<String, RoleDefinition>>,
    memberships: RefCell<BTreeMap<RoleMembershipKey, RoleMembership>>,
    fail_persistence: Cell<bool>,
    epoch: Cell<usize>,
    publications: RefCell<Vec<&'static str>>,
}

impl DatabaseCatalog {
    fn context(&self) -> DatabasePrivilegeContext<'_> {
        DatabasePrivilegeContext {
            names: self,
            locks: self,
            roles: self,
            registry: self,
            publication: self,
        }
    }

    fn assert_authorization_is_retained(&self) {
        assert!(self.roles.try_borrow_mut().is_err());
        assert!(self.memberships.try_borrow_mut().is_err());
    }

    fn assert_guards_are_released(&self) {
        assert!(self.security.try_borrow_mut().is_ok());
        assert!(self.roles.try_borrow_mut().is_ok());
        assert!(self.memberships.try_borrow_mut().is_ok());
    }
}

impl RoleReferenceNames for DatabaseCatalog {
    fn outer_role(&self) -> uqa_sql::catalog::roles::RoleReference {
        self.current_role()
    }
    fn current_role(&self) -> RoleReference {
        "uqa".into()
    }

    fn session_role(&self) -> RoleReference {
        "uqa".into()
    }
}

impl RoleCatalogGuards for DatabaseCatalog {
    fn role_definitions(&self) -> RoleDefinitionRead<'_> {
        Box::new(self.roles.borrow())
    }

    fn role_memberships(&self) -> RoleMembershipRead<'_> {
        Box::new(self.memberships.borrow())
    }
}

impl DatabaseSecurityRegistry for DatabaseCatalog {
    fn security_read(&self) -> DatabaseSecurityRead<'_> {
        Box::new(self.security.borrow())
    }

    fn security_write(&self) -> DatabaseSecurityWrite<'_> {
        self.assert_authorization_is_retained();
        self.publications.borrow_mut().push("publish");
        Box::new(self.security.borrow_mut())
    }
}

impl DatabasePrivilegePublication for DatabaseCatalog {
    fn prepare_writer(&self) -> Result<(), SQLError> {
        self.assert_guards_are_released();
        Ok(())
    }

    fn refresh_catalog(&self) -> StorageBackendResult<()> {
        Ok(())
    }

    fn persist_security(&self, json: &str) -> Result<(), SQLError> {
        self.assert_authorization_is_retained();
        assert!(self.security.try_borrow_mut().is_ok());
        let stored: serde_json::Value = serde_json::from_str(json).unwrap();
        assert_eq!(stored["database_security_format"], 1);
        let bound: uqa_sql::catalog::security::database::binding::BoundDatabaseSecurity =
            serde_json::from_value(stored["security"].clone()).unwrap();
        assert_ne!(bound, *self.security.borrow());
        self.publications.borrow_mut().push("persist");
        if self.fail_persistence.get() {
            return Err(SQLError::Internal("simulated persistence failure".into()));
        }
        Ok(())
    }

    fn catalog_changed(&self) {
        self.assert_authorization_is_retained();
        assert!(self.security.try_borrow_mut().is_ok());
        self.publications.borrow_mut().push("epoch");
        self.epoch.set(self.epoch.get() + 1);
    }

    fn notice(&self, _: &str, _: &str) {
        self.assert_guards_are_released();
        panic!("the database owner can grant CREATE without a warning");
    }
}

impl SharedObjectLockSession for DatabaseCatalog {
    fn acquire_shared_catalog(
        &self,
        _: crate::row_locks::shared_objects::SharedCatalogLock<'_>,
        _: crate::row_locks::RelationLockMode,
    ) -> Result<crate::row_locks::ScopedRelationLock<'_>, SQLError> {
        panic!("PUBLIC must not acquire a role dependency lock")
    }

    fn refresh_shared_catalog(&self) -> Result<(), SQLError> {
        panic!("PUBLIC has no shared role dependency")
    }
}

#[test]
fn database_acl_persistence_failure_leaves_security_and_epoch_unchanged() {
    let catalog = DatabaseCatalog {
        security: RefCell::new(BoundDatabaseSecurity::bootstrap()),
        roles: RefCell::new(BTreeMap::from([(
            "uqa".into(),
            RoleDefinition::bootstrap(),
        )])),
        memberships: RefCell::new(BTreeMap::new()),
        fail_persistence: Cell::new(true),
        epoch: Cell::new(0),
        publications: RefCell::new(Vec::new()),
    };
    let statement = GrantDatabaseStmt {
        is_grant: true,
        grant_option: false,
        grant_option_only: false,
        privileges: vec![DatabasePrivilege::Create],
        databases: vec!["uqa".into()],
        grantees: vec![uqa_sql::ast::AclRoleSpecification::Public],
        grantor: None,
        revoke_behavior: DatabaseRevokeBehavior::Restrict,
    };
    let initial = catalog.security.borrow().clone();
    let result = grant_database_privileges(&catalog.context(), &statement);
    assert!(
        matches!(result, Err(SQLError::Internal(message)) if message == "simulated persistence failure")
    );
    assert_eq!(*catalog.security.borrow(), initial);
    assert_eq!(catalog.epoch.get(), 0);
    assert_eq!(*catalog.publications.borrow(), ["persist"]);
    catalog.assert_guards_are_released();

    catalog.fail_persistence.set(false);
    catalog.publications.borrow_mut().clear();
    grant_database_privileges(&catalog.context(), &statement).unwrap();
    assert_ne!(*catalog.security.borrow(), initial);
    assert_eq!(catalog.epoch.get(), 1);
    assert_eq!(
        *catalog.publications.borrow(),
        ["persist", "publish", "epoch"]
    );
    catalog.assert_guards_are_released();
}