#![allow(clippy::manual_async_fn)]
use crate::DjogiError;
use crate::context::DjogiContext;
use crate::ident::check_user_supplied_ident;
use crate::model::Model;
use crate::pg::accumulator::{SqlAccumulator, as_params};
use crate::pg::decode::try_get_scalar;
use crate::query::terminal::auto_set_tenant;
use postgres_types::ToSql;
use std::future::Future;
#[derive(Debug)]
pub struct MaterializeClosureOptions<Pk: ToSql + Sync + Send + 'static> {
pub max_depth: Option<u32>,
pub roots: Option<Vec<Pk>>,
}
impl<Pk: ToSql + Sync + Send + 'static> Default for MaterializeClosureOptions<Pk> {
fn default() -> Self {
Self {
max_depth: None,
roots: None,
}
}
}
impl<Pk: ToSql + Sync + Send + 'static> MaterializeClosureOptions<Pk> {
pub fn with_max_depth(mut self, n: u32) -> Self {
self.max_depth = Some(n);
self
}
pub fn with_roots(mut self, ids: Vec<Pk>) -> Self {
self.roots = Some(ids);
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MaterializeClosureReport {
pub rows_written: u64,
pub sources_visited: usize,
}
pub trait ClosureModel: Model {
type Source: Model;
fn table() -> &'static str {
Self::table_name()
}
fn source_column() -> &'static str;
fn ancestor_column() -> &'static str;
fn depth_column() -> &'static str;
fn path_count_column() -> &'static str;
}
pub(crate) fn materialize_closure_impl<'ctx, T, C>(
ctx: &'ctx mut DjogiContext,
opts: MaterializeClosureOptions<T::Pk>,
) -> impl Future<Output = Result<MaterializeClosureReport, DjogiError>> + Send + 'ctx
where
T: Model,
T::Pk: ToSql + Sync + Send + 'static,
C: ClosureModel<Source = T>,
{
async move {
if matches!(&opts.roots, Some(ids) if ids.is_empty()) {
return Ok(MaterializeClosureReport {
rows_written: 0,
sources_visited: 0,
});
}
let descriptor = T::descriptor();
if descriptor.self_fk_count() == 0 {
return Err(DjogiError::Validation(format!(
"model '{}' has no self-FK; materialize_closure requires at least one",
T::table_name(),
)));
}
validate_closure_metadata_idents::<C>()?;
auto_set_tenant::<T>(ctx).await?;
let acc = build_materialize_closure_sql::<T, C>(opts);
let (sql, binds) = acc.into_parts();
let params = as_params(&binds);
let row = ctx.query_one(&sql, ¶ms).await?;
let rows_written: i64 = try_get_scalar::<i64>(&row, 0)?;
let sources_visited: i64 = try_get_scalar::<i64>(&row, 1)?;
Ok(MaterializeClosureReport {
rows_written: rows_written.max(0) as u64,
sources_visited: sources_visited.max(0) as usize,
})
}
}
pub(crate) fn validate_closure_metadata_idents<C: ClosureModel>() -> Result<(), DjogiError> {
for (label, col) in [
("closure table", C::table()),
("source_column", C::source_column()),
("ancestor_column", C::ancestor_column()),
("depth_column", C::depth_column()),
("path_count_column", C::path_count_column()),
] {
check_user_supplied_ident(col, true).map_err(|e| {
DjogiError::Validation(format!(
"ClosureModel<{}>::{} returned invalid identifier {:?}: {:?}",
std::any::type_name::<C>(),
label,
col,
e,
))
})?;
}
Ok(())
}
pub(crate) fn build_materialize_closure_sql<T, C>(
opts: MaterializeClosureOptions<T::Pk>,
) -> SqlAccumulator
where
T: Model,
T::Pk: ToSql + Sync + Send + 'static,
C: ClosureModel<Source = T>,
{
let mut acc = SqlAccumulator::new("");
acc.push_sql("WITH inserted AS (INSERT INTO ");
acc.push_sql(C::table());
acc.push_sql(" (");
acc.push_sql(C::source_column());
acc.push_sql(", ");
acc.push_sql(C::ancestor_column());
acc.push_sql(", ");
acc.push_sql(C::depth_column());
acc.push_sql(", ");
acc.push_sql(C::path_count_column());
acc.push_sql(") ");
acc.push_sql("WITH RECURSIVE __djogi_closure (source_id, ancestor_id, depth, path) AS (");
acc.push_sql("SELECT s.id, s.id, 0, ARRAY[]::text[] FROM ");
acc.push_sql(T::table_name());
acc.push_sql(" s WHERE ");
let roots = opts.roots;
let max_depth = opts.max_depth;
match roots {
None => {
acc.push_sql("TRUE");
}
Some(ids) => {
debug_assert!(
!ids.is_empty(),
"empty roots Vec must be short-circuited before SQL emission",
);
acc.push_sql("s.id IN (");
acc.push_list_binds(ids);
acc.push_sql(")");
}
}
let edges: Vec<&'static str> = T::descriptor().self_fk_columns().collect();
acc.push_sql(
" UNION ALL SELECT closure.source_id, p.id, closure.depth + 1, \
closure.path || ARRAY[step.label] FROM __djogi_closure closure JOIN ",
);
acc.push_sql(T::table_name());
acc.push_sql(" a ON a.id = closure.ancestor_id CROSS JOIN LATERAL (VALUES ");
for (i, edge) in edges.iter().enumerate() {
if i > 0 {
acc.push_sql(", ");
}
acc.push_sql("(a.");
acc.push_sql(edge);
acc.push_sql(", '");
acc.push_sql(edge);
acc.push_sql("'::text)");
}
acc.push_sql(") AS step(pid, label) JOIN ");
acc.push_sql(T::table_name());
acc.push_sql(" p ON p.id = step.pid");
if let Some(n) = max_depth {
acc.push_sql(" WHERE closure.depth < ");
let n_i32 = i32::try_from(n).unwrap_or(i32::MAX);
acc.push_bind(n_i32);
}
acc.push_sql(") CYCLE source_id, ancestor_id SET is_cycle USING cycle_path");
acc.push_sql(
" SELECT source_id, ancestor_id, depth, COUNT(*) AS path_count \
FROM __djogi_closure WHERE NOT is_cycle \
GROUP BY source_id, ancestor_id, depth",
);
acc.push_sql(" ON CONFLICT (");
acc.push_sql(C::source_column());
acc.push_sql(", ");
acc.push_sql(C::ancestor_column());
acc.push_sql(", ");
acc.push_sql(C::depth_column());
acc.push_sql(") DO UPDATE SET ");
acc.push_sql(C::path_count_column());
acc.push_sql(" = EXCLUDED.");
acc.push_sql(C::path_count_column());
acc.push_sql(" RETURNING ");
acc.push_sql(C::source_column());
acc.push_sql(") SELECT COUNT(*), COUNT(DISTINCT ");
acc.push_sql(C::source_column());
acc.push_sql(") FROM inserted");
acc
}
#[cfg(test)]
mod tests {
use super::*;
use crate::descriptor::{
FieldDescriptor, FieldSqlType, ModelDescriptor, PkType, RelationKind, field_descriptor,
model_descriptor,
};
use crate::pg::decode::FromPgRow;
use crate::types::HeerId;
struct MiniTree;
impl crate::model::__sealed::Sealed for MiniTree {}
impl crate::model::Model for MiniTree {
type Pk = HeerId;
type Fields = ();
fn table_name() -> &'static str {
"mini_trees"
}
fn pk_value(&self) -> &Self::Pk {
unreachable!()
}
fn descriptor() -> &'static ModelDescriptor {
&MINI_TREE_DESC_ONE_EDGE
}
fn get(
_ctx: &mut DjogiContext,
_id: Self::Pk,
) -> impl std::future::Future<Output = Result<Self, DjogiError>> + Send {
async { unreachable!() }
}
fn create(
_ctx: &mut DjogiContext,
_v: Self,
) -> impl std::future::Future<Output = Result<Self, DjogiError>> + Send {
async { unreachable!() }
}
fn save<'ctx>(
&'ctx mut self,
_ctx: &'ctx mut DjogiContext,
) -> impl std::future::Future<Output = Result<(), DjogiError>> + Send + 'ctx {
async { unreachable!() }
}
fn delete(
self,
_ctx: &mut DjogiContext,
) -> impl std::future::Future<Output = Result<(), DjogiError>> + Send {
async { unreachable!() }
}
fn refresh_from_db<'ctx>(
&'ctx self,
_ctx: &'ctx mut DjogiContext,
) -> impl std::future::Future<Output = Result<Self, DjogiError>> + Send + 'ctx {
async { unreachable!() }
}
}
impl FromPgRow for MiniTree {
const COLUMNS: &'static [&'static str] = &["id", "parent_id"];
const COLUMN_LIST: &'static str = "id, parent_id";
fn from_pg_row(_row: &tokio_postgres::Row) -> Result<Self, DjogiError> {
unreachable!()
}
}
struct MiniPedigree;
impl crate::model::__sealed::Sealed for MiniPedigree {}
impl crate::model::Model for MiniPedigree {
type Pk = HeerId;
type Fields = ();
fn table_name() -> &'static str {
"mini_pedigrees"
}
fn pk_value(&self) -> &Self::Pk {
unreachable!()
}
fn descriptor() -> &'static ModelDescriptor {
&MINI_PEDIGREE_DESC_TWO_EDGES
}
fn get(
_ctx: &mut DjogiContext,
_id: Self::Pk,
) -> impl std::future::Future<Output = Result<Self, DjogiError>> + Send {
async { unreachable!() }
}
fn create(
_ctx: &mut DjogiContext,
_v: Self,
) -> impl std::future::Future<Output = Result<Self, DjogiError>> + Send {
async { unreachable!() }
}
fn save<'ctx>(
&'ctx mut self,
_ctx: &'ctx mut DjogiContext,
) -> impl std::future::Future<Output = Result<(), DjogiError>> + Send + 'ctx {
async { unreachable!() }
}
fn delete(
self,
_ctx: &mut DjogiContext,
) -> impl std::future::Future<Output = Result<(), DjogiError>> + Send {
async { unreachable!() }
}
fn refresh_from_db<'ctx>(
&'ctx self,
_ctx: &'ctx mut DjogiContext,
) -> impl std::future::Future<Output = Result<Self, DjogiError>> + Send + 'ctx {
async { unreachable!() }
}
}
impl FromPgRow for MiniPedigree {
const COLUMNS: &'static [&'static str] = &["id", "mother_id", "father_id"];
const COLUMN_LIST: &'static str = "id, mother_id, father_id";
fn from_pg_row(_row: &tokio_postgres::Row) -> Result<Self, DjogiError> {
unreachable!()
}
}
struct MiniClosure;
impl crate::model::__sealed::Sealed for MiniClosure {}
impl crate::model::Model for MiniClosure {
type Pk = HeerId;
type Fields = ();
fn table_name() -> &'static str {
"mini_tree_closures"
}
fn pk_value(&self) -> &Self::Pk {
unreachable!()
}
fn descriptor() -> &'static ModelDescriptor {
unreachable!()
}
fn get(
_ctx: &mut DjogiContext,
_id: Self::Pk,
) -> impl std::future::Future<Output = Result<Self, DjogiError>> + Send {
async { unreachable!() }
}
fn create(
_ctx: &mut DjogiContext,
_v: Self,
) -> impl std::future::Future<Output = Result<Self, DjogiError>> + Send {
async { unreachable!() }
}
fn save<'ctx>(
&'ctx mut self,
_ctx: &'ctx mut DjogiContext,
) -> impl std::future::Future<Output = Result<(), DjogiError>> + Send + 'ctx {
async { unreachable!() }
}
fn delete(
self,
_ctx: &mut DjogiContext,
) -> impl std::future::Future<Output = Result<(), DjogiError>> + Send {
async { unreachable!() }
}
fn refresh_from_db<'ctx>(
&'ctx self,
_ctx: &'ctx mut DjogiContext,
) -> impl std::future::Future<Output = Result<Self, DjogiError>> + Send + 'ctx {
async { unreachable!() }
}
}
impl FromPgRow for MiniClosure {
const COLUMNS: &'static [&'static str] = &[];
const COLUMN_LIST: &'static str = "";
fn from_pg_row(_row: &tokio_postgres::Row) -> Result<Self, DjogiError> {
unreachable!()
}
}
impl ClosureModel for MiniClosure {
type Source = MiniTree;
fn source_column() -> &'static str {
"mini_tree_id"
}
fn ancestor_column() -> &'static str {
"ancestor_id"
}
fn depth_column() -> &'static str {
"depth"
}
fn path_count_column() -> &'static str {
"path_count"
}
}
struct MiniPedigreeClosure;
impl crate::model::__sealed::Sealed for MiniPedigreeClosure {}
impl crate::model::Model for MiniPedigreeClosure {
type Pk = HeerId;
type Fields = ();
fn table_name() -> &'static str {
"mini_pedigree_closures"
}
fn pk_value(&self) -> &Self::Pk {
unreachable!()
}
fn descriptor() -> &'static ModelDescriptor {
unreachable!()
}
fn get(
_ctx: &mut DjogiContext,
_id: Self::Pk,
) -> impl std::future::Future<Output = Result<Self, DjogiError>> + Send {
async { unreachable!() }
}
fn create(
_ctx: &mut DjogiContext,
_v: Self,
) -> impl std::future::Future<Output = Result<Self, DjogiError>> + Send {
async { unreachable!() }
}
fn save<'ctx>(
&'ctx mut self,
_ctx: &'ctx mut DjogiContext,
) -> impl std::future::Future<Output = Result<(), DjogiError>> + Send + 'ctx {
async { unreachable!() }
}
fn delete(
self,
_ctx: &mut DjogiContext,
) -> impl std::future::Future<Output = Result<(), DjogiError>> + Send {
async { unreachable!() }
}
fn refresh_from_db<'ctx>(
&'ctx self,
_ctx: &'ctx mut DjogiContext,
) -> impl std::future::Future<Output = Result<Self, DjogiError>> + Send + 'ctx {
async { unreachable!() }
}
}
impl FromPgRow for MiniPedigreeClosure {
const COLUMNS: &'static [&'static str] = &[];
const COLUMN_LIST: &'static str = "";
fn from_pg_row(_row: &tokio_postgres::Row) -> Result<Self, DjogiError> {
unreachable!()
}
}
impl ClosureModel for MiniPedigreeClosure {
type Source = MiniPedigree;
fn source_column() -> &'static str {
"mini_pedigree_id"
}
fn ancestor_column() -> &'static str {
"ancestor_id"
}
fn depth_column() -> &'static str {
"depth"
}
fn path_count_column() -> &'static str {
"path_count"
}
}
const fn self_fk_field(name: &'static str) -> FieldDescriptor {
FieldDescriptor {
relation_kind: Some(RelationKind::ForeignKey),
is_self_fk: true,
..field_descriptor(name, FieldSqlType::BigInt, false)
}
}
static MINI_TREE_FIELDS: &[FieldDescriptor] = &[self_fk_field("parent_id")];
static MINI_TREE_DESC_ONE_EDGE: ModelDescriptor =
model_descriptor("MiniTree", "mini_trees", PkType::HeerId, MINI_TREE_FIELDS);
static MINI_PEDIGREE_FIELDS: &[FieldDescriptor] =
&[self_fk_field("mother_id"), self_fk_field("father_id")];
static MINI_PEDIGREE_DESC_TWO_EDGES: ModelDescriptor = model_descriptor(
"MiniPedigree",
"mini_pedigrees",
PkType::HeerId,
MINI_PEDIGREE_FIELDS,
);
fn opts_default() -> MaterializeClosureOptions<HeerId> {
MaterializeClosureOptions::<HeerId>::default()
}
#[test]
fn anchor_projects_source_ancestor_depth_path() {
let acc = build_materialize_closure_sql::<MiniTree, MiniClosure>(opts_default());
let sql = acc.sql();
assert!(
sql.contains("__djogi_closure (source_id, ancestor_id, depth, path)"),
"CTE column list must be (source_id, ancestor_id, depth, path): {sql}"
);
assert!(
sql.contains("SELECT s.id, s.id, 0, ARRAY[]::text[] FROM mini_trees s"),
"anchor must project (s.id, s.id, 0, empty path): {sql}"
);
}
#[test]
fn one_edge_emits_single_recursive_term() {
let acc = build_materialize_closure_sql::<MiniTree, MiniClosure>(opts_default());
let sql = acc.sql();
let union_count = sql.matches("UNION ALL").count();
assert_eq!(
union_count, 1,
"single-edge closure: exactly 1 UNION ALL (anchor + recursive): {sql}"
);
assert!(
sql.contains("a ON a.id = closure.ancestor_id"),
"must resolve ancestor row via closure.ancestor_id: {sql}"
);
assert!(
sql.contains("(a.parent_id, 'parent_id'::text)"),
"single-edge LATERAL VALUES tuple must contain (a.parent_id, 'parent_id'::text): {sql}"
);
assert!(
sql.contains("p ON p.id = step.pid"),
"must step to parent via the LATERAL `step.pid`: {sql}"
);
}
#[test]
fn two_edges_collapse_into_single_recursive_term() {
let opts = MaterializeClosureOptions::<HeerId>::default();
let acc = build_materialize_closure_sql::<MiniPedigree, MiniPedigreeClosure>(opts);
let sql = acc.sql();
let union_count = sql.matches("UNION ALL").count();
assert_eq!(
union_count, 1,
"two-edge closure: exactly 1 UNION ALL even with N edges (anchor + single recursive term): {sql}"
);
assert!(
sql.contains("(a.mother_id, 'mother_id'::text), (a.father_id, 'father_id'::text)"),
"LATERAL VALUES must enumerate both edges as (a.<col>, '<col>'::text) tuples: {sql}"
);
assert!(
sql.contains("p ON p.id = step.pid"),
"must step to parent via LATERAL `step.pid` — single JOIN, not per-edge: {sql}"
);
}
#[test]
fn cycle_clause_uses_two_columns() {
let acc = build_materialize_closure_sql::<MiniTree, MiniClosure>(opts_default());
let sql = acc.sql();
assert!(
sql.contains("CYCLE source_id, ancestor_id SET is_cycle USING cycle_path"),
"CYCLE clause must use (source_id, ancestor_id) and cycle_path array: {sql}"
);
}
#[test]
fn outer_select_groups_by_source_ancestor_depth() {
let acc = build_materialize_closure_sql::<MiniTree, MiniClosure>(opts_default());
let sql = acc.sql();
assert!(
sql.contains("SELECT source_id, ancestor_id, depth, COUNT(*) AS path_count"),
"outer SELECT must aggregate to (source, ancestor, depth, COUNT(*)): {sql}"
);
assert!(
sql.contains("GROUP BY source_id, ancestor_id, depth"),
"outer SELECT must GROUP BY (source, ancestor, depth): {sql}"
);
}
#[test]
fn on_conflict_clause_replaces_path_count() {
let acc = build_materialize_closure_sql::<MiniTree, MiniClosure>(opts_default());
let sql = acc.sql();
assert!(
sql.contains(
"ON CONFLICT (mini_tree_id, ancestor_id, depth) DO UPDATE SET path_count = EXCLUDED.path_count"
),
"ON CONFLICT clause must replace path_count from EXCLUDED: {sql}"
);
assert!(
!sql.contains("mini_tree_closures.path_count + EXCLUDED.path_count"),
"additive merge must not appear (would double on rerun): {sql}"
);
}
#[test]
fn returning_source_col_drives_outer_count() {
let acc = build_materialize_closure_sql::<MiniTree, MiniClosure>(opts_default());
let sql = acc.sql();
assert!(
sql.starts_with("WITH inserted AS (INSERT INTO mini_tree_closures"),
"outer wrap must start with WITH inserted AS (INSERT...): {sql}"
);
assert!(
sql.contains("RETURNING mini_tree_id"),
"INSERT must RETURN the source_col so the outer COUNT can run on it: {sql}"
);
assert!(
sql.ends_with(") SELECT COUNT(*), COUNT(DISTINCT mini_tree_id) FROM inserted"),
"outer SELECT must compute (rows_written, sources_visited): {sql}"
);
}
#[test]
fn no_max_depth_no_predicate() {
let acc = build_materialize_closure_sql::<MiniTree, MiniClosure>(opts_default());
let sql = acc.sql();
assert!(
!sql.contains("closure.depth <"),
"no max_depth must not emit `closure.depth <`: {sql}"
);
}
#[test]
fn with_max_depth_emits_depth_predicate_and_binds() {
let opts = MaterializeClosureOptions::<HeerId>::default().with_max_depth(5);
let acc = build_materialize_closure_sql::<MiniTree, MiniClosure>(opts);
let sql = acc.sql();
assert!(
sql.contains("closure.depth < $1"),
"max_depth must bind as $1 (no roots binds): {sql}"
);
assert_eq!(
acc.bind_count(),
1,
"single-edge + max_depth = 1 bind, got {}",
acc.bind_count()
);
}
#[test]
fn two_edges_max_depth_binds_once() {
let opts = MaterializeClosureOptions::<HeerId>::default().with_max_depth(3);
let acc = build_materialize_closure_sql::<MiniPedigree, MiniPedigreeClosure>(opts);
let sql = acc.sql();
assert!(
sql.contains("closure.depth < $1"),
"single recursive term => single depth bind at $1: {sql}"
);
assert!(
!sql.contains("closure.depth < $2"),
"second depth bind must not appear — only one recursive term: {sql}"
);
assert_eq!(
acc.bind_count(),
1,
"consolidated recursive term => 1 max_depth bind regardless of edge count, got {}",
acc.bind_count()
);
}
#[test]
fn no_roots_emits_where_true() {
let acc = build_materialize_closure_sql::<MiniTree, MiniClosure>(opts_default());
let sql = acc.sql();
assert!(
sql.contains("FROM mini_trees s WHERE TRUE"),
"no roots must emit WHERE TRUE: {sql}"
);
}
#[test]
fn with_roots_emits_in_clause_and_binds_each_id() {
let ids = vec![
HeerId::from_i64(1).unwrap(),
HeerId::from_i64(2).unwrap(),
HeerId::from_i64(3).unwrap(),
];
let opts = MaterializeClosureOptions::<HeerId>::default().with_roots(ids);
let acc = build_materialize_closure_sql::<MiniTree, MiniClosure>(opts);
let sql = acc.sql();
assert!(
sql.contains("FROM mini_trees s WHERE s.id IN ($1, $2, $3)"),
"with_roots must emit s.id IN (<binds>) with one bind per id: {sql}"
);
assert_eq!(
acc.bind_count(),
3,
"three roots = 3 binds, got {}",
acc.bind_count()
);
assert!(
!sql.contains("closure.depth <"),
"no max_depth must not emit a depth probe: {sql}"
);
}
#[test]
fn with_roots_and_max_depth_binds_after_roots() {
let ids = vec![
HeerId::from_i64(1).unwrap(),
HeerId::from_i64(2).unwrap(),
HeerId::from_i64(3).unwrap(),
];
let opts = MaterializeClosureOptions::<HeerId>::default()
.with_roots(ids)
.with_max_depth(7);
let acc = build_materialize_closure_sql::<MiniTree, MiniClosure>(opts);
let sql = acc.sql();
assert!(
sql.contains("s.id IN ($1, $2, $3)"),
"roots take $1..$3: {sql}"
);
assert!(
sql.contains("closure.depth < $4"),
"max_depth bind starts at $4 (after the 3 roots): {sql}"
);
assert_eq!(
acc.bind_count(),
4,
"3 roots + 1 depth cap = 4 binds, got {}",
acc.bind_count()
);
}
#[test]
fn empty_edges_descriptor_reports_zero_count() {
static NO_FIELDS: &[FieldDescriptor] = &[];
static MINI_DESC_NO_EDGES: ModelDescriptor =
model_descriptor("MiniNoEdges", "mini_no_edges", PkType::HeerId, NO_FIELDS);
assert_eq!(MINI_DESC_NO_EDGES.self_fk_count(), 0);
}
}