Skip to main content

uqa_sql/schema/
referenced_partitions.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! The constraints a foreign key derives on the partitions of the partitioned table it references, named and identified as `PostgreSQL`'s `addFkRecurseReferenced` and `CloneFkReferenced` create them and removed as `DetachPartitionFinalize` and `AttachPartitionForeignKey` remove them.
8
9use std::collections::{BTreeMap, BTreeSet};
10
11use crate::ast::{ColumnDef, ReferencedPartitionConstraint, TableConstraintSet};
12use crate::schema::constraint_metadata::{
13    CatalogIdentityAllocator, ConstraintMetadataError, ConstraintMetadataResult,
14};
15use crate::SQLError;
16
17#[cfg(test)]
18mod tests;
19
20/// A referenced partition and the partitioned partition it belongs to, or `None` when it belongs to the foreign key's referenced table.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct ReferencedPartition {
23    pub partition: [u8; 16],
24    pub parent: Option<[u8; 16]>,
25}
26
27/// The state of the foreign key that its derived constraints follow.
28#[derive(Debug, Clone, Copy)]
29pub struct ReferencingConstraint<'a> {
30    pub name: &'a str,
31    pub validated: bool,
32    pub enforced: bool,
33}
34
35/// What derivation reads from the catalog.
36pub trait ReferencedPartitionSource {
37    /// The partitions below `table`, each before its own partitions, with siblings in partition order; empty when `table` is not partitioned.
38    fn referenced_partitions(&self, table: &str) -> Result<Vec<ReferencedPartition>, SQLError>;
39    /// The object identities of the foreign keys `table` declares, which its partitions' copies share.
40    fn declared_foreign_key_ids(&self, table: &str) -> Result<BTreeSet<[u8; 16]>, SQLError>;
41}
42
43/// Make `constraints` hold one derived constraint for each of `partitions`, the referenced table's partition tree in order, keeping the constraints of partitions that remain under the same parent. The partitions that joined below one constraint since the constraints were last derived take their names from that constraint and its validity, as `PostgreSQL` recurses into a new partition subtree with the name of the constraint it joins: the foreign key's own for a partition of the referenced table, else the derived constraint of the partition's parent. A valid foreign key validates every derived constraint, and one that is not enforced validates none. Reports whether the constraints changed.
44pub fn reconcile_referenced_partition_constraints(
45    constraints: &mut Vec<ReferencedPartitionConstraint>,
46    foreign_key: ReferencingConstraint<'_>,
47    partitions: &[ReferencedPartition],
48    used: &mut BTreeSet<String>,
49    allocate: &mut CatalogIdentityAllocator<'_>,
50) -> ConstraintMetadataResult<bool> {
51    let previous = std::mem::take(constraints);
52    let mut kept = previous
53        .iter()
54        .filter(|constraint| {
55            partitions.iter().any(|partition| {
56                partition.partition == constraint.partition && partition.parent == constraint.parent
57            })
58        })
59        .map(|constraint| (constraint.partition, constraint.clone()))
60        .collect::<BTreeMap<_, _>>();
61    let mut joined = BTreeMap::<[u8; 16], (String, bool)>::new();
62    let mut reconciled = Vec::with_capacity(partitions.len());
63    for partition in partitions {
64        if let Some(constraint) = kept.remove(&partition.partition) {
65            reconciled.push(constraint);
66            continue;
67        }
68        let (base, validated) = match partition.parent {
69            None => (foreign_key.name.to_string(), foreign_key.validated),
70            Some(parent) => joined
71                .get(&parent)
72                .cloned()
73                .or_else(|| {
74                    reconciled
75                        .iter()
76                        .find(|constraint: &&ReferencedPartitionConstraint| {
77                            constraint.partition == parent
78                        })
79                        .map(|constraint| (constraint.name.clone(), constraint.validated))
80                })
81                .ok_or_else(|| {
82                    ConstraintMetadataError::Invalid(
83                        "a referenced partition precedes its parent in partition order".into(),
84                    )
85                })?,
86        };
87        let name = super::constraint_metadata::choose_suffixed_constraint_name(&base, used)?;
88        let mut catalog_identity = None;
89        super::constraint_metadata::identity::foreign_keys::materialize(
90            &mut catalog_identity,
91            allocate,
92        )?;
93        let catalog_identity = catalog_identity.ok_or_else(|| {
94            ConstraintMetadataError::Invalid("a derived constraint has no catalog identity".into())
95        })?;
96        joined.insert(partition.partition, (base, validated));
97        reconciled.push(ReferencedPartitionConstraint {
98            partition: partition.partition,
99            parent: partition.parent,
100            name,
101            catalog_identity,
102            validated,
103        });
104    }
105    for constraint in &mut reconciled {
106        if foreign_key.validated {
107            constraint.validated = true;
108        }
109        if !foreign_key.enforced {
110            constraint.validated = false;
111        }
112    }
113    let changed = reconciled != previous;
114    *constraints = reconciled;
115    Ok(changed)
116}
117
118fn source_error(error: SQLError) -> ConstraintMetadataError {
119    ConstraintMetadataError::Execution(Box::new(error))
120}
121
122fn unnamed() -> ConstraintMetadataError {
123    ConstraintMetadataError::Invalid("a foreign key derives constraints before it is named".into())
124}
125
126/// Reconcile the derived constraints of every foreign key a table declares with the current partition trees of the tables they reference. A copy of a foreign key that the table's partitioned parent declares derives none, as only the constraint without a parent derives constraints in `PostgreSQL`; a copy that became the table's own when the table was detached derives them under its own name. Reports whether any foreign key changed.
127pub fn reconcile_table_referenced_partition_constraints(
128    source: &dyn ReferencedPartitionSource,
129    columns: &mut [ColumnDef],
130    constraints: &mut TableConstraintSet,
131    used: &mut BTreeSet<String>,
132    allocate: &mut CatalogIdentityAllocator<'_>,
133) -> ConstraintMetadataResult<bool> {
134    let inherited = match constraints
135        .hierarchy
136        .parents
137        .first()
138        .filter(|_| constraints.hierarchy.is_partition())
139    {
140        Some(parent) => source
141            .declared_foreign_key_ids(parent)
142            .map_err(source_error)?,
143        None => BTreeSet::new(),
144    };
145    let mut changed = false;
146    for reference in columns
147        .iter_mut()
148        .filter_map(|column| column.references.as_mut())
149    {
150        if reference
151            .object_id
152            .is_some_and(|object_id| inherited.contains(&object_id))
153        {
154            changed |= !reference.referenced_partitions.is_empty();
155            reference.referenced_partitions.clear();
156            continue;
157        }
158        let partitions = source
159            .referenced_partitions(&reference.table)
160            .map_err(source_error)?;
161        changed |= reconcile_referenced_partition_constraints(
162            &mut reference.referenced_partitions,
163            ReferencingConstraint {
164                name: reference.name.as_deref().ok_or_else(unnamed)?,
165                validated: reference.validated,
166                enforced: reference.enforced,
167            },
168            &partitions,
169            used,
170            allocate,
171        )?;
172    }
173    for foreign_key in &mut constraints.foreign_keys {
174        if foreign_key
175            .object_id
176            .is_some_and(|object_id| inherited.contains(&object_id))
177        {
178            changed |= !foreign_key.referenced_partitions.is_empty();
179            foreign_key.referenced_partitions.clear();
180            continue;
181        }
182        let partitions = source
183            .referenced_partitions(&foreign_key.ref_table)
184            .map_err(source_error)?;
185        changed |= reconcile_referenced_partition_constraints(
186            &mut foreign_key.referenced_partitions,
187            ReferencingConstraint {
188                name: foreign_key.name.as_deref().ok_or_else(unnamed)?,
189                validated: foreign_key.validated,
190                enforced: foreign_key.enforced,
191            },
192            &partitions,
193            used,
194            allocate,
195        )?;
196    }
197    Ok(changed)
198}
199
200/// The constraint whose triggers a change to a referenced row fires, and the relation that holds it.
201#[derive(Debug, Clone, Copy)]
202pub struct FiringConstraint<'a> {
203    /// The partition that holds the row, or the relation an `UPDATE` names when the row moves to another partition.
204    pub relation: &'a str,
205    /// The constraint the foreign key derives on `relation`, or `None` for the foreign key's own.
206    pub derived: Option<&'a ReferencedPartitionConstraint>,
207}
208
209impl FiringConstraint<'_> {
210    /// The name of the constraint, given the foreign key's own.
211    pub fn name<'n>(&'n self, foreign_key: &'n crate::ast::ForeignKey) -> &'n str {
212        self.derived.map_or_else(
213            || foreign_key.name.as_deref().unwrap_or("<unnamed>"),
214            |derived| derived.name.as_str(),
215        )
216    }
217}
218
219/// The constraint of `foreign_key` whose triggers a change to a referenced row of `relation`, the partition that holds it, fires: the constraint it derives on `relation`, or the foreign key's own when the foreign key references `relation`. A row that moves to another partition fires the constraint at `moved_through`, the relation the `UPDATE` names, when the foreign key references it or a table it is a partition of, as `PostgreSQL` fires the update triggers of the update's root for a moved row.
220pub fn firing_constraint<'a>(
221    catalog: &dyn crate::semantics::partition::PartitionCatalog,
222    relation: &'a str,
223    moved_through: Option<&'a str>,
224    foreign_key: &'a crate::ast::ForeignKey,
225) -> Result<FiringConstraint<'a>, SQLError> {
226    let object_id = |table: &str| {
227        catalog
228            .try_table_object_id(table)
229            .map_err(SQLError::Internal)
230    };
231    let derived_on = |table: &str| -> Result<Option<&'a ReferencedPartitionConstraint>, SQLError> {
232        Ok(object_id(table)?.and_then(|object_id| {
233            foreign_key
234                .referenced_partitions
235                .iter()
236                .find(|constraint| constraint.partition == object_id)
237        }))
238    };
239    if let Some(target) = moved_through {
240        let referenced = object_id(&foreign_key.ref_table)?;
241        if referenced.is_some() && object_id(target)? == referenced {
242            return Ok(FiringConstraint {
243                relation: target,
244                derived: None,
245            });
246        }
247        if let Some(derived) = derived_on(target)? {
248            return Ok(FiringConstraint {
249                relation: target,
250                derived: Some(derived),
251            });
252        }
253    }
254    Ok(FiringConstraint {
255        relation,
256        derived: derived_on(relation)?,
257    })
258}