#![cfg_attr(not(feature = "gql"), allow(dead_code))]
use std::collections::HashMap;
use std::sync::Arc;
use super::pipeline::{
FieldState, build_field_state, compute_fields_for_value, filter_fields_by_permission,
};
use crate::catalog::providers::TableProvider;
use crate::exec::permission::{
PhysicalPermission, check_permission_for_value, convert_permission_to_physical_runtime,
resolve_select_permission, should_check_perms,
};
use crate::exec::{ControlFlowExt, ExecutionContext};
use crate::expr::ControlFlow;
use crate::iam::Action;
use crate::val::{RecordId, TableName, Value};
struct TableFetchState {
permission: PhysicalPermission,
field_state: FieldState,
}
pub(crate) struct FetchFieldStateCache {
tables: HashMap<TableName, TableFetchState>,
check_perms: Option<bool>,
}
impl FetchFieldStateCache {
pub(crate) fn new() -> Self {
Self {
tables: HashMap::new(),
check_perms: None,
}
}
pub(crate) fn check_perms(&mut self, ctx: &ExecutionContext) -> Result<bool, ControlFlow> {
self.resolve_check_perms(ctx)
}
fn resolve_check_perms(&mut self, ctx: &ExecutionContext) -> Result<bool, ControlFlow> {
if let Some(flag) = self.check_perms {
return Ok(flag);
}
let db_ctx = ctx.database().context("MATCH binding fetch requires database context")?;
let flag = should_check_perms(db_ctx, Action::View)
.map_err(|e| ControlFlow::Err(anyhow::Error::new(e)))?;
self.check_perms = Some(flag);
Ok(flag)
}
async fn state_for<'a>(
&'a mut self,
ctx: &ExecutionContext,
table: &TableName,
check_perms: bool,
) -> Result<&'a TableFetchState, ControlFlow> {
if !self.tables.contains_key(table) {
let permission = if check_perms {
let db_ctx =
ctx.database().context("MATCH binding fetch requires database context")?;
let version = ctx.version_stamp();
let table_def = db_ctx
.get_table_def(table, version)
.await
.context("Failed to get table definition")?;
let catalog_perm = resolve_select_permission(table_def.as_deref());
convert_permission_to_physical_runtime(catalog_perm, ctx.ctx())
.await
.context("Failed to convert permission")?
} else {
PhysicalPermission::Allow
};
let field_state = build_field_state(ctx, table, check_perms, None).await?;
self.tables.insert(
table.clone(),
TableFetchState {
permission,
field_state,
},
);
}
Ok(self.tables.get(table).expect("table state just inserted"))
}
}
impl Default for FetchFieldStateCache {
fn default() -> Self {
Self::new()
}
}
pub(crate) async fn resolve_with_field_state(
ctx: &ExecutionContext,
cache: &mut FetchFieldStateCache,
rids: &[RecordId],
) -> Result<Vec<Option<Value>>, ControlFlow> {
if rids.is_empty() {
return Ok(Vec::new());
}
let check_perms = cache.resolve_check_perms(ctx)?;
let version = ctx.version_stamp();
for rid in rids {
if !cache.tables.contains_key(&rid.table) {
cache.state_for(ctx, &rid.table, check_perms).await?;
}
}
let db_ctx = ctx.database().context("MATCH binding fetch requires database context")?;
let txn = ctx.txn();
let ns_id = db_ctx.ns_ctx.ns.namespace_id;
let db_id = db_ctx.db.database_id;
let records = txn
.get_records(ns_id, db_id, rids, version, crate::kvs::CachePolicy::ReadWrite)
.await
.context("Failed to fetch records")?;
let mut out: Vec<Option<Value>> = Vec::with_capacity(rids.len());
for (rid, record) in rids.iter().zip(records) {
if record.data.is_none() {
out.push(None);
continue;
}
let table_state = cache
.tables
.get(&rid.table)
.expect("table state pre-warmed above for every batch table");
if check_perms {
let allowed =
check_permission_for_value(&table_state.permission, &record.data, None, ctx)
.await
.context("Failed to check table permission")?;
if !allowed {
out.push(None);
continue;
}
}
let mut value = match Arc::try_unwrap(record) {
Ok(rec) => rec.data,
Err(arc) => arc.data.clone(),
};
compute_fields_for_value(ctx, &table_state.field_state, &mut value, false).await?;
if check_perms {
filter_fields_by_permission(ctx, &table_state.field_state, &mut value).await?;
}
out.push(Some(value));
}
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn empty_cache_starts_uninitialised() {
let cache = FetchFieldStateCache::new();
assert!(cache.tables.is_empty());
assert!(cache.check_perms.is_none());
}
#[test]
fn default_matches_new() {
let a = FetchFieldStateCache::default();
assert!(a.tables.is_empty());
assert!(a.check_perms.is_none());
}
}