use std::sync::Arc;
use std::time::Duration;
use crate::core::{LifecycleState, OperationGuard};
use crate::database::Database;
use crate::error::{MongrelError, Result};
use crate::service_principal::ServicePrincipalDefinition;
pub struct SecretString(zeroize::Zeroizing<String>);
impl SecretString {
pub fn new(secret: impl Into<String>) -> Self {
Self(zeroize::Zeroizing::new(secret.into()))
}
pub(crate) fn expose(&self) -> &str {
self.0.as_str()
}
}
impl std::fmt::Debug for SecretString {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("SecretString([REDACTED])")
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum HandleIdentity {
Credentialless,
CatalogUser {
username: String,
user_id: u64,
created_version: u64,
},
ServicePrincipal {
token_id: String,
principal_id: [u8; 16],
creation_version: u64,
},
}
#[derive(Debug)]
pub enum OpenIdentity {
Credentialless,
ServiceCredentials {
token_id: String,
secret: SecretString,
},
CatalogCredentials {
username: String,
password: SecretString,
},
}
#[derive(Debug, Clone)]
pub(crate) struct InternalServiceCapability {
pub principal_id: [u8; 16],
pub permissions: Vec<crate::auth::Permission>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct HandleAccess {
read_only: bool,
}
impl HandleAccess {
pub fn read_write() -> Self {
Self { read_only: false }
}
pub fn read_only() -> Self {
Self { read_only: true }
}
pub fn is_read_only(&self) -> bool {
self.read_only
}
}
pub struct DatabaseHandle {
database: Database,
identity: HandleIdentity,
access: HandleAccess,
internal_capability: Option<InternalServiceCapability>,
}
impl DatabaseHandle {
pub(crate) fn new(
database: Database,
identity: HandleIdentity,
access: HandleAccess,
internal_capability: Option<InternalServiceCapability>,
) -> Self {
Self {
database,
identity,
access,
internal_capability,
}
}
#[allow(dead_code)] pub(crate) fn with_internal_capability(
database: Database,
capability: InternalServiceCapability,
access: HandleAccess,
) -> Self {
let identity = HandleIdentity::ServicePrincipal {
token_id: format!("internal:{:02x?}", capability.principal_id),
principal_id: capability.principal_id,
creation_version: 0,
};
Self::new(database, identity, access, Some(capability))
}
pub fn identity(&self) -> &HandleIdentity {
&self.identity
}
pub fn access(&self) -> HandleAccess {
self.access
}
pub fn shares_core_with(&self, other: &Self) -> bool {
Arc::ptr_eq(&self.database.core(), &other.database.core())
}
pub fn lifecycle_state(&self) -> LifecycleState {
self.database.lifecycle_state()
}
pub fn operation_guard(&self) -> Result<OperationGuard> {
self.database.operation_guard()
}
fn authorize_write(&self, operation: &'static str) -> Result<()> {
if self.access.is_read_only() {
Err(MongrelError::ReadOnlyHandle { operation })
} else {
Ok(())
}
}
fn resolve_service_authority(&self) -> Result<()> {
if self.internal_capability.is_some() {
return Ok(());
}
let HandleIdentity::ServicePrincipal {
token_id,
principal_id,
creation_version,
} = &self.identity
else {
return Ok(());
};
let def = self.database.core().service_principals.resolve_live(
token_id,
*principal_id,
*creation_version,
)?;
self.database.rebind_service_principal(&def);
Ok(())
}
fn authorize_permission(&self, permission: &crate::auth::Permission) -> Result<()> {
if let Some(capability) = &self.internal_capability {
if !capability
.permissions
.iter()
.any(|granted| granted.satisfies(permission))
{
return Err(MongrelError::PermissionDenied {
required: permission.clone(),
principal: format!("service:{:02x?}", capability.principal_id),
});
}
return Ok(());
}
if let HandleIdentity::ServicePrincipal {
token_id,
principal_id,
creation_version,
} = &self.identity
{
let def = self.database.core().service_principals.resolve_live(
token_id,
*principal_id,
*creation_version,
)?;
self.database.rebind_service_principal(&def);
if !def
.permissions
.iter()
.any(|granted| granted.satisfies(permission))
{
return Err(MongrelError::PermissionDenied {
required: permission.clone(),
principal: format!("service:{token_id}"),
});
}
return Ok(());
}
self.database.require(permission)
}
pub fn query(
&self,
table: &str,
query: &crate::query::Query,
projection: Option<&[u16]>,
) -> Result<Vec<crate::memtable::Row>> {
self.resolve_service_authority()?;
self.database
.query_for_current_principal(table, query, projection)
}
pub fn count(&self, table: &str) -> Result<u64> {
self.resolve_service_authority()?;
self.database
.count_for(table, self.database.principal().as_ref())
}
pub fn rows(&self, table: &str) -> Result<Vec<crate::memtable::Row>> {
self.resolve_service_authority()?;
self.database
.rows_for(table, self.database.principal().as_ref())
}
pub fn create_table(&self, name: &str, schema: crate::schema::Schema) -> Result<u64> {
self.authorize_write("create table")?;
self.authorize_permission(&crate::auth::Permission::Ddl)?;
self.database.create_table(name, schema)
}
pub fn put(
&self,
table: &str,
cells: Vec<(u16, crate::memtable::Value)>,
) -> Result<Option<i64>> {
self.authorize_write("put")?;
self.resolve_service_authority()?;
self.database
.transaction_for_current_principal(|transaction| transaction.put(table, cells))
}
pub fn put_batch(
&self,
table: &str,
rows: Vec<Vec<(u16, crate::memtable::Value)>>,
) -> Result<Vec<Option<i64>>> {
self.authorize_write("put_batch")?;
self.resolve_service_authority()?;
self.database
.transaction_for_current_principal(|transaction| transaction.put_batch(table, rows))
}
pub fn update(
&self,
table: &str,
row_id: crate::RowId,
cells: Vec<(u16, crate::memtable::Value)>,
) -> Result<crate::txn::OwnedRow> {
self.authorize_write("update")?;
self.resolve_service_authority()?;
self.database
.transaction_for_current_principal(|transaction| {
let mut images = transaction.update_many(table, vec![(row_id, cells)])?;
images.pop().ok_or_else(|| {
MongrelError::NotFound(format!("row {row_id:?} not found for update"))
})
})
}
pub fn session(&self) -> Result<AuthorizedMongrelSession<'_>> {
self.resolve_service_authority()?;
Ok(AuthorizedMongrelSession { handle: self })
}
pub fn begin(&self) -> Result<AuthorizedTransaction<'_>> {
self.authorize_write("begin")?;
self.resolve_service_authority()?;
let principal = self.database.principal();
let txn = self.database.begin_as(principal);
Ok(AuthorizedTransaction {
handle: self,
txn: Some(txn),
})
}
pub fn delete(&self, table: &str, row_id: crate::RowId) -> Result<()> {
self.authorize_write("delete")?;
self.resolve_service_authority()?;
self.database
.transaction_for_current_principal(|transaction| transaction.delete(table, row_id))
}
pub fn create_index(&self, table: &str, definition: crate::schema::IndexDef) -> Result<u64> {
self.authorize_write("create index")?;
self.authorize_permission(&crate::auth::Permission::Ddl)?;
self.database.create_index(table, definition)
}
pub fn drop_index(&self, table: &str, name: &str) -> Result<()> {
self.authorize_write("drop index")?;
self.authorize_permission(&crate::auth::Permission::Ddl)?;
self.database.drop_index(table, name)
}
pub fn create_procedure(
&self,
procedure: crate::procedure::StoredProcedure,
) -> Result<crate::procedure::StoredProcedure> {
self.authorize_write("create procedure")?;
self.authorize_permission(&crate::auth::Permission::Ddl)?;
self.resolve_service_authority()?;
self.database.create_procedure(procedure)
}
pub fn drop_procedure(&self, name: &str) -> Result<()> {
self.authorize_write("drop procedure")?;
self.authorize_permission(&crate::auth::Permission::Ddl)?;
self.resolve_service_authority()?;
self.database.drop_procedure(name)
}
pub fn call_procedure(
&self,
name: &str,
args: std::collections::HashMap<String, crate::memtable::Value>,
) -> Result<crate::procedure::ProcedureCallResult> {
self.resolve_service_authority()?;
self.database.call_procedure(name, args)
}
pub fn create_trigger(
&self,
trigger: crate::trigger::StoredTrigger,
) -> Result<crate::trigger::StoredTrigger> {
self.authorize_write("create trigger")?;
self.authorize_permission(&crate::auth::Permission::Ddl)?;
self.resolve_service_authority()?;
self.database.create_trigger(trigger)
}
pub fn drop_trigger(&self, name: &str) -> Result<()> {
self.authorize_write("drop trigger")?;
self.authorize_permission(&crate::auth::Permission::Ddl)?;
self.resolve_service_authority()?;
self.database.drop_trigger(name)
}
pub fn rows_at_epoch(
&self,
table: &str,
snapshot: crate::epoch::Snapshot,
) -> Result<Vec<crate::memtable::Row>> {
self.resolve_service_authority()?;
self.database
.rows_at_epoch_for_current_principal(table, snapshot)
}
pub fn create_user(&self, username: &str, password: &str) -> Result<crate::auth::UserEntry> {
self.authorize_write("create user")?;
self.authorize_permission(&crate::auth::Permission::Admin)?;
self.database.create_user(username, password)
}
pub fn drop_user(&self, username: &str) -> Result<()> {
self.authorize_write("drop user")?;
self.authorize_permission(&crate::auth::Permission::Admin)?;
self.database.drop_user(username)
}
pub fn create_role(&self, role: &str) -> Result<crate::auth::RoleEntry> {
self.authorize_write("create role")?;
self.authorize_permission(&crate::auth::Permission::Admin)?;
self.database.create_role(role)
}
pub fn grant_role(&self, username: &str, role: &str) -> Result<()> {
self.authorize_write("grant role")?;
self.authorize_permission(&crate::auth::Permission::Admin)?;
self.database.grant_role(username, role)
}
pub fn revoke_role(&self, username: &str, role: &str) -> Result<()> {
self.authorize_write("revoke role")?;
self.authorize_permission(&crate::auth::Permission::Admin)?;
self.database.revoke_role(username, role)
}
pub fn grant_permission(&self, role: &str, permission: crate::auth::Permission) -> Result<()> {
self.authorize_write("grant permission")?;
self.authorize_permission(&crate::auth::Permission::Admin)?;
self.database.grant_permission(role, permission)
}
pub fn revoke_permission(&self, role: &str, permission: crate::auth::Permission) -> Result<()> {
self.authorize_write("revoke permission")?;
self.authorize_permission(&crate::auth::Permission::Admin)?;
self.database.revoke_permission(role, permission)
}
pub fn register_service_principal(
&self,
token_id: impl Into<String>,
principal_id: [u8; 16],
permissions: Vec<crate::auth::Permission>,
raw_secret: &str,
expires_unix: u64,
) -> Result<ServicePrincipalDefinition> {
self.authorize_write("register service principal")?;
self.authorize_permission(&crate::auth::Permission::Admin)?;
self.database.core().service_principals.register(
token_id,
principal_id,
permissions,
raw_secret,
expires_unix,
)
}
pub fn revoke_service_principal(&self, token_id: &str) -> Result<()> {
self.authorize_write("revoke service principal")?;
self.authorize_permission(&crate::auth::Permission::Admin)?;
self.database.core().service_principals.revoke(token_id)
}
pub fn set_service_principal_permissions(
&self,
token_id: &str,
permissions: Vec<crate::auth::Permission>,
) -> Result<()> {
self.authorize_write("set service principal permissions")?;
self.authorize_permission(&crate::auth::Permission::Admin)?;
self.database
.core()
.service_principals
.set_permissions(token_id, permissions)
}
pub fn shutdown(&self, drain_deadline: Duration) -> Result<()> {
self.authorize_write("shutdown")?;
self.authorize_permission(&crate::auth::Permission::Admin)?;
self.database.core().shutdown(drain_deadline)
}
}
pub struct AuthorizedMongrelSession<'a> {
handle: &'a DatabaseHandle,
}
pub type AuthorizedSession<'a> = AuthorizedMongrelSession<'a>;
impl<'a> AuthorizedMongrelSession<'a> {
pub fn begin(&self) -> Result<AuthorizedTransaction<'a>> {
self.handle.begin()
}
pub fn put(
&self,
table: &str,
cells: Vec<(u16, crate::memtable::Value)>,
) -> Result<Option<i64>> {
self.handle.put(table, cells)
}
pub fn update(
&self,
table: &str,
row_id: crate::RowId,
cells: Vec<(u16, crate::memtable::Value)>,
) -> Result<crate::txn::OwnedRow> {
self.handle.update(table, row_id, cells)
}
pub fn delete(&self, table: &str, row_id: crate::RowId) -> Result<()> {
self.handle.delete(table, row_id)
}
pub fn query(
&self,
table: &str,
query: &crate::query::Query,
projection: Option<&[u16]>,
) -> Result<Vec<crate::memtable::Row>> {
self.handle.query(table, query, projection)
}
pub fn count(&self, table: &str) -> Result<u64> {
self.handle.count(table)
}
pub fn put_batch(
&self,
table: &str,
rows: Vec<Vec<(u16, crate::memtable::Value)>>,
) -> Result<Vec<Option<i64>>> {
self.handle.put_batch(table, rows)
}
pub fn create_index(&self, table: &str, definition: crate::schema::IndexDef) -> Result<u64> {
self.handle.create_index(table, definition)
}
pub fn drop_index(&self, table: &str, name: &str) -> Result<()> {
self.handle.drop_index(table, name)
}
pub fn create_procedure(
&self,
procedure: crate::procedure::StoredProcedure,
) -> Result<crate::procedure::StoredProcedure> {
self.handle.create_procedure(procedure)
}
pub fn drop_procedure(&self, name: &str) -> Result<()> {
self.handle.drop_procedure(name)
}
pub fn call_procedure(
&self,
name: &str,
args: std::collections::HashMap<String, crate::memtable::Value>,
) -> Result<crate::procedure::ProcedureCallResult> {
self.handle.call_procedure(name, args)
}
pub fn create_trigger(
&self,
trigger: crate::trigger::StoredTrigger,
) -> Result<crate::trigger::StoredTrigger> {
self.handle.create_trigger(trigger)
}
pub fn drop_trigger(&self, name: &str) -> Result<()> {
self.handle.drop_trigger(name)
}
pub fn rows_at_epoch(
&self,
table: &str,
snapshot: crate::epoch::Snapshot,
) -> Result<Vec<crate::memtable::Row>> {
self.handle.rows_at_epoch(table, snapshot)
}
pub fn aggregate_count(&self, table: &str) -> Result<u64> {
self.handle.count(table)
}
}
pub struct AuthorizedTransaction<'a> {
handle: &'a DatabaseHandle,
txn: Option<crate::txn::Transaction<'a>>,
}
impl<'a> AuthorizedTransaction<'a> {
fn txn_mut(&mut self) -> Result<&mut crate::txn::Transaction<'a>> {
self.txn
.as_mut()
.ok_or_else(|| MongrelError::InvalidArgument("transaction already finished".into()))
}
pub fn put(
&mut self,
table: &str,
cells: Vec<(u16, crate::memtable::Value)>,
) -> Result<Option<i64>> {
self.handle.authorize_write("put")?;
self.handle.resolve_service_authority()?;
self.txn_mut()?.put(table, cells)
}
pub fn put_batch(
&mut self,
table: &str,
rows: Vec<Vec<(u16, crate::memtable::Value)>>,
) -> Result<Vec<Option<i64>>> {
self.handle.authorize_write("put_batch")?;
self.handle.resolve_service_authority()?;
self.txn_mut()?.put_batch(table, rows)
}
pub fn update(
&mut self,
table: &str,
row_id: crate::RowId,
cells: Vec<(u16, crate::memtable::Value)>,
) -> Result<crate::txn::OwnedRow> {
self.handle.authorize_write("update")?;
self.handle.resolve_service_authority()?;
let mut images = self.txn_mut()?.update_many(table, vec![(row_id, cells)])?;
images
.pop()
.ok_or_else(|| MongrelError::NotFound(format!("row {row_id:?} not found for update")))
}
pub fn delete(&mut self, table: &str, row_id: crate::RowId) -> Result<()> {
self.handle.authorize_write("delete")?;
self.handle.resolve_service_authority()?;
self.txn_mut()?.delete(table, row_id)
}
pub fn commit(mut self) -> Result<crate::epoch::Epoch> {
self.handle.authorize_write("commit")?;
let txn = self
.txn
.take()
.ok_or_else(|| MongrelError::InvalidArgument("transaction already finished".into()))?;
txn.commit()
}
pub fn rollback(mut self) {
if let Some(txn) = self.txn.take() {
txn.rollback();
}
}
}
impl Drop for AuthorizedTransaction<'_> {
fn drop(&mut self) {
if let Some(txn) = self.txn.take() {
txn.rollback();
}
}
}
impl std::fmt::Debug for DatabaseHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("DatabaseHandle")
.field("identity", &self.identity)
.field("access", &self.access)
.field("database", &self.database)
.finish()
}
}