1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct ReferencedPartition {
23 pub partition: [u8; 16],
24 pub parent: Option<[u8; 16]>,
25}
26
27#[derive(Debug, Clone, Copy)]
29pub struct ReferencingConstraint<'a> {
30 pub name: &'a str,
31 pub validated: bool,
32 pub enforced: bool,
33}
34
35pub trait ReferencedPartitionSource {
37 fn referenced_partitions(&self, table: &str) -> Result<Vec<ReferencedPartition>, SQLError>;
39 fn declared_foreign_key_ids(&self, table: &str) -> Result<BTreeSet<[u8; 16]>, SQLError>;
41}
42
43pub 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
126pub 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#[derive(Debug, Clone, Copy)]
202pub struct FiringConstraint<'a> {
203 pub relation: &'a str,
205 pub derived: Option<&'a ReferencedPartitionConstraint>,
207}
208
209impl FiringConstraint<'_> {
210 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
219pub 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}