use std::collections::BTreeMap;
use datafusion::prelude::SessionContext;
use time::OffsetDateTime;
use crate::error::{Error, Result};
use super::query::QueryResult;
use super::store::{MeterStore, MeterStoreBuilder};
pub struct MeterCatalog {
ctx: SessionContext,
stores: BTreeMap<String, MeterStore>,
}
impl std::fmt::Debug for MeterCatalog {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MeterCatalog")
.field("tables", &self.stores.keys().collect::<Vec<_>>())
.finish_non_exhaustive()
}
}
impl MeterCatalog {
pub fn builder() -> MeterCatalogBuilder {
MeterCatalogBuilder::default()
}
pub fn context(&self) -> &SessionContext {
&self.ctx
}
pub fn table(&self, name: &str) -> Option<&MeterStore> {
self.stores.get(name).or_else(|| {
self.stores
.values()
.find(|s| super::store::resolved_name(s.config().name()) == name)
})
}
pub async fn isolated(&self, name: &str) -> Result<MeterStore> {
let store = self.table(name).ok_or_else(|| {
Error::config(format!(
"this catalog holds no table {name:?}: it holds [{}]",
self.stores
.keys()
.map(String::as_str)
.collect::<Vec<_>>()
.join(", ")
))
})?;
store.in_own_session().await
}
pub async fn scoped(&self, column: &str, value: impl Into<String>) -> Result<Self> {
let value = value.into();
let mut narrowed = BTreeMap::new();
for (name, store) in &self.stores {
let scope = store.narrowed_scope(column, &value).map_err(|e| {
Error::config(format!(
"this catalog cannot be scoped to {column} = {value:?}: {e}. A scope \
that covered the other tables and skipped this one would not be a \
boundary — isolate to one table and scope that, if this table is \
meant to be outside it"
))
})?;
narrowed.insert(name.clone(), scope);
}
self.rebuild(|name, store| {
let builder = store.to_builder();
match narrowed.get(name).and_then(Option::as_ref) {
Some(scope) => builder.row_scope(scope.clone()),
None => builder,
}
})
.await
}
pub async fn as_known_at(&self, at: OffsetDateTime) -> Result<Self> {
self.in_read_mode(crate::planner::ReadMode::AsKnownAt(at))
.await
}
pub async fn in_read_mode(&self, mode: crate::planner::ReadMode) -> Result<Self> {
if let crate::planner::ReadMode::AsOf { .. } = mode {
return Err(Error::config(
"a pinned snapshot belongs to one table: there is no snapshot id that \
means the same moment in two of them, and nothing commits two atomically. \
Pin each table with MeterCatalog::table(..).as_of(..), or use \
MeterCatalog::as_known_at, whose recorded_at axis every table shares",
));
}
self.rebuild(|_, store| store.to_builder().read_mode(mode))
.await
}
async fn rebuild<F>(&self, f: F) -> Result<Self>
where
F: Fn(&str, &MeterStore) -> MeterStoreBuilder,
{
let ctx = SessionContext::new_with_config(
datafusion::prelude::SessionConfig::new().with_information_schema(true),
);
let mut stores = BTreeMap::new();
for (name, store) in &self.stores {
let derived = f(name, store).session(ctx.clone()).build().await?;
stores.insert(name.clone(), derived);
}
Ok(Self { ctx, stores })
}
pub fn tables(&self) -> impl Iterator<Item = &MeterStore> {
self.stores.values()
}
pub fn len(&self) -> usize {
self.stores.len()
}
pub fn is_empty(&self) -> bool {
self.stores.is_empty()
}
pub async fn query(&self, sql: &str) -> Result<QueryResult> {
self.query_with_params(sql, Vec::new()).await
}
pub async fn query_with_params(
&self,
sql: &str,
params: Vec<datafusion::scalar::ScalarValue>,
) -> Result<QueryResult> {
let (first, watermarks) = self.watermarks_for(sql).await?;
first.run(sql, params, watermarks).await
}
async fn watermarks_for(
&self,
sql: &str,
) -> Result<(
&MeterStore,
Vec<(String, crate::watermark::TieringWatermark)>,
)> {
let first = self
.stores
.values()
.next()
.ok_or_else(|| Error::config("catalog hosts no tables"))?;
let plan = self
.ctx
.state()
.create_logical_plan(sql)
.await
.map_err(Error::from)?;
let scanned = scanned_relations(&plan);
let mut watermarks = Vec::with_capacity(self.stores.len());
for (name, store) in &self.stores {
let touched =
scanned.contains(&store.raw_table()) || scanned.contains(&store.resolved_table());
if touched {
watermarks.push((name.clone(), store.watermark().await?));
}
}
Ok((first, watermarks))
}
pub async fn describe(&self, sql: &str) -> Result<super::QueryDescription> {
let (first, watermarks) = self.watermarks_for(sql).await?;
first.describe_with(sql, watermarks).await
}
pub async fn stream(
&self,
sql: &str,
) -> Result<(
super::QueryDescription,
datafusion::execution::SendableRecordBatchStream,
)> {
self.stream_with_params(sql, Vec::new()).await
}
pub async fn stream_with_params(
&self,
sql: &str,
params: Vec<datafusion::scalar::ScalarValue>,
) -> Result<(
super::QueryDescription,
datafusion::execution::SendableRecordBatchStream,
)> {
let (first, watermarks) = self.watermarks_for(sql).await?;
first.stream_at(sql, params, watermarks).await
}
pub fn maintenance(&self) -> super::Maintenance {
super::Maintenance::over(self.stores.values().cloned().collect())
}
pub async fn create_tables(&self) -> Result<()> {
for store in self.stores.values() {
store.create_tables().await?;
}
Ok(())
}
pub async fn refresh_system_tables(&self, now: OffsetDateTime) -> Result<()> {
let views: Vec<super::system::SystemTables<'_>> = self
.stores
.values()
.map(|s| super::system::SystemTables::new(s.hot_store(), s.cold_store(), s.config()))
.collect();
super::system::register_all(&self.ctx, &views, now).await
}
pub async fn status(&self, now: OffsetDateTime) -> Result<Vec<super::system::TableStatus>> {
let mut out = Vec::with_capacity(self.stores.len());
for store in self.stores.values() {
out.push(store.status(now).await?);
}
Ok(out)
}
pub async fn archive_all(
&self,
now: OffsetDateTime,
max_windows: usize,
) -> Result<Vec<(String, Vec<crate::tiering::archive::ArchivalOutcome>)>> {
let mut out = Vec::with_capacity(self.stores.len());
for (name, store) in &self.stores {
let outcome = store.archive(now, max_windows).await.inspect_err(|e| {
tracing::error!(table = %name, error = %e, "archiving a catalog table failed");
})?;
out.push((name.clone(), outcome));
}
Ok(out)
}
pub async fn verify_invariant(&self) -> Result<Vec<(String, String)>> {
let mut out = Vec::new();
for (name, store) in &self.stores {
if let Err(e) = store.verify_invariant().await {
match e {
Error::InvariantViolated { detail, .. } => out.push((name.clone(), detail)),
other => return Err(other),
}
}
}
Ok(out)
}
pub async fn anonymise_before(
&self,
cutoff: OffsetDateTime,
reason: &str,
actor: &str,
now: OffsetDateTime,
) -> Result<Vec<crate::erasure::ErasureRecord>> {
anonymise_across(self.stores.values(), cutoff, reason, actor, now).await
}
#[must_use]
pub fn subject_registry(&self) -> Option<&crate::erasure::SubjectRegistry> {
self.stores.values().find_map(MeterStore::subject_registry)
}
pub async fn erase_subject_by_id(
&self,
natural_id: &str,
reason: &str,
actor: &str,
now: OffsetDateTime,
) -> Result<Vec<crate::erasure::ErasureRecord>> {
self.require_registry()?
.erase_all(natural_id, reason, actor, now)
.await
}
pub async fn subject_epochs(&self, natural_id: &str) -> Result<Vec<i32>> {
self.require_registry()?.epochs(natural_id).await
}
pub async fn subject_registrations(
&self,
natural_id: &str,
) -> Result<Vec<crate::erasure::SubjectRegistration>> {
self.require_registry()?.registrations(natural_id).await
}
pub async fn is_subject_suppressed(&self, natural_id: &str) -> Result<bool> {
self.require_registry()?.is_suppressed(natural_id).await
}
pub async fn lift_subject_suppression(
&self,
natural_id: &str,
reason: &str,
actor: &str,
now: OffsetDateTime,
) -> Result<bool> {
self.require_registry()?
.lift_suppression(natural_id, reason, actor, now)
.await
}
pub async fn erasures(
&self,
query: &crate::erasure::ErasureQuery,
) -> Result<Vec<crate::erasure::ErasureRecord>> {
self.require_registry()?.erasures(query).await
}
fn require_registry(&self) -> Result<&crate::erasure::SubjectRegistry> {
self.subject_registry().ok_or_else(|| {
Error::config(
"no table in this catalog declares a subject column, so this \
deployment holds no subject mapping: there is nothing to enumerate \
and nothing to erase",
)
})
}
}
pub(crate) async fn anonymise_across<'a>(
stores: impl Iterator<Item = &'a MeterStore>,
cutoff: OffsetDateTime,
reason: &str,
actor: &str,
now: OffsetDateTime,
) -> Result<Vec<crate::erasure::ErasureRecord>> {
let mut registry = None;
let mut subject_tables = 0usize;
for store in stores {
if store.config().subject_column().is_none() {
continue;
}
subject_tables += 1;
registry = registry.or_else(|| store.subject_registry());
}
if subject_tables == 0 {
return Err(Error::config(
"no table in this catalog declares a subject column, so there is no \
linkage to destroy: without one the stored readings carry no reference \
to a person and § 60 Abs. 6 has nothing to act on here",
));
}
let registry = registry.ok_or_else(|| {
Error::config(
"a table declares a subject column but no SubjectRegistry is configured, \
so the references it stores resolve to nothing this process can erase",
)
})?;
registry
.expire_epochs_before(cutoff, reason, actor, now)
.await
}
fn scanned_relations(
plan: &datafusion::logical_expr::LogicalPlan,
) -> std::collections::HashSet<String> {
use datafusion::common::tree_node::TreeNodeRecursion;
use datafusion::logical_expr::LogicalPlan;
let mut found = std::collections::HashSet::new();
let _ = plan.apply_with_subqueries(|node| {
if let LogicalPlan::TableScan(scan) = node {
found.insert(scan.table_name.table().to_string());
}
Ok(TreeNodeRecursion::Continue)
});
found
}
#[async_trait::async_trait]
impl super::SqlSurface for MeterCatalog {
fn label(&self) -> String {
self.stores
.values()
.map(MeterStore::resolved_table)
.collect::<Vec<_>>()
.join(", ")
}
async fn describe_sql(&self, sql: &str) -> Result<super::QueryDescription> {
self.describe(sql).await
}
async fn stream_sql(
&self,
sql: &str,
params: Vec<datafusion::scalar::ScalarValue>,
) -> Result<(
super::QueryDescription,
datafusion::execution::SendableRecordBatchStream,
)> {
self.stream_with_params(sql, params).await
}
}
#[derive(Default)]
pub struct MeterCatalogBuilder {
tables: Vec<MeterStoreBuilder>,
}
impl std::fmt::Debug for MeterCatalogBuilder {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MeterCatalogBuilder")
.field("tables", &self.tables.len())
.finish()
}
}
impl MeterCatalogBuilder {
pub fn table(mut self, table: MeterStoreBuilder) -> Self {
self.tables.push(table);
self
}
pub async fn build(self) -> Result<MeterCatalog> {
if self.tables.is_empty() {
return Err(Error::config(
"a catalog needs at least one table: an empty one can answer no query \
and hides the missing configuration until the first request",
));
}
let ctx = SessionContext::new_with_config(
datafusion::prelude::SessionConfig::new().with_information_schema(true),
);
let mut seen: BTreeMap<String, String> = BTreeMap::new();
for builder in &self.tables {
let configured = builder.table_name().ok_or_else(|| {
Error::config("every table in a catalog needs a table configuration")
})?;
let (raw, resolved) = builder
.registered_names()
.expect("a builder with a table name has registered names");
for relation in [raw, resolved] {
if let Some(owner) = seen.get(&relation) {
return Err(Error::config(format!(
"tables {owner:?} and {configured:?} both register the relation \
{relation:?}: only one of them could answer to it, and a query \
naming it would silently read whichever won"
)));
}
seen.insert(relation, configured.to_string());
}
}
let mut stores = BTreeMap::new();
for builder in self.tables {
let store = builder.session(ctx.clone()).build().await?;
stores.insert(store.config().name().to_string(), store);
}
let modes: Vec<(&str, _)> = stores
.values()
.map(|s| (s.config().name(), s.read_mode()))
.collect();
let uniform = modes.windows(2).all(|w| w[0].1 == w[1].1);
if !uniform {
return Err(Error::config(format!(
"a catalog's tables must share one read mode, but these are {modes:?}: \
a result reports the single mode its statement ran under, and mixing \
them would attribute a figure half-computed from the mutable tier to \
a mode that excludes it"
)));
}
Ok(MeterCatalog { ctx, stores })
}
}