use rupg_catalog::{
ConKind, Constraint, FkSide, NewForeignKey, ObjRef, PG_CONSTRAINT, PG_TRIGGER, RelKind,
};
use rupg_common::{Error, Result, SqlState};
use super::Definer;
impl Definer<'_, '_> {
pub(super) fn fk_recurse_referenced(
&mut self,
constraint: u32,
parents: Option<(u32, u32)>,
) -> Result<()> {
let con = self.foreign_constraint(constraint)?;
let triggers = if con.enforced {
self.catalog.add_fk_action_triggers(constraint, parents)?
} else {
(0, 0)
};
let Some(foreign) = con.foreign.clone() else { return Ok(()) };
if !self.is_partitioned(foreign.table) {
return Ok(());
}
for partition in self.partition_oids(foreign.table) {
let index = self.partition_index(partition, con.index)?;
let keys = self.map_keys(foreign.table, partition, &foreign.keys)?;
let new = NewForeignKey {
table: con.relation,
name: Some(con.name.clone()),
keys: con.keys.clone(),
index,
deferrable: con.deferrable,
deferred: con.deferred,
foreign: rupg_catalog::ForeignKey { table: partition, keys, ..foreign.clone() },
};
let clone = self.catalog.add_fk_constraint(new, constraint, FkSide::Referenced)?;
self.fk_recurse_referenced(clone, Some(triggers))?;
}
Ok(())
}
pub(super) fn fk_recurse_referencing(
&mut self,
constraint: u32,
parents: Option<(u32, u32)>,
) -> Result<()> {
let con = self.foreign_constraint(constraint)?;
let triggers = if con.enforced {
self.catalog.add_fk_check_triggers(constraint, parents)?
} else {
(0, 0)
};
let Some(foreign) = con.foreign.clone() else { return Ok(()) };
if !self.is_partitioned(con.relation) {
return Ok(());
}
for partition in self.partition_oids(con.relation) {
let keys = self.map_keys(con.relation, partition, &con.keys)?;
if self.try_attach_foreign_key(partition, &con, &keys, triggers)? {
continue;
}
let new = NewForeignKey {
table: partition,
name: Some(con.name.clone()),
keys,
index: con.index,
deferrable: con.deferrable,
deferred: con.deferred,
foreign: foreign.clone(),
};
let clone = self.catalog.add_fk_constraint(new, constraint, FkSide::Referencing)?;
self.fk_recurse_referencing(clone, Some(triggers))?;
}
Ok(())
}
pub(super) fn clone_foreign_keys(&mut self, parent: u32, partition: u32) -> Result<()> {
let own: Vec<Constraint> = self
.catalog
.constraints_of(parent)
.filter(|c| c.kind == ConKind::Foreign)
.cloned()
.collect();
let oids: Vec<u32> = own.iter().map(|c| c.oid).collect();
for con in own.iter().filter(|c| !oids.contains(&c.parent)) {
let Some(foreign) = con.foreign.clone() else { continue };
let keys = self.map_keys(parent, partition, &con.keys)?;
let triggers = self.fk_triggers(con.oid, parent, true);
if self.try_attach_foreign_key(partition, con, &keys, triggers)? {
continue;
}
let new = NewForeignKey {
table: partition,
name: Some(con.name.clone()),
keys,
index: con.index,
deferrable: con.deferrable,
deferred: con.deferred,
foreign,
};
let clone = self.catalog.add_fk_constraint(new, con.oid, FkSide::Referencing)?;
self.fk_recurse_referencing(clone, Some(triggers))?;
}
let referring: Vec<Constraint> = self
.catalog
.constraints()
.filter(|c| c.foreign.as_ref().is_some_and(|f| f.table == parent))
.cloned()
.collect();
let oids: Vec<u32> = referring.iter().map(|c| c.oid).collect();
for con in &referring {
if oids.contains(&con.parent) || con.relation == parent || con.relation == partition {
continue;
}
let Some(foreign) = con.foreign.clone() else { continue };
let index = self.partition_index(partition, con.index)?;
let keys = self.map_keys(parent, partition, &foreign.keys)?;
let triggers = self.fk_triggers(con.oid, parent, false);
let new = NewForeignKey {
table: con.relation,
name: Some(con.name.clone()),
keys: con.keys.clone(),
index,
deferrable: con.deferrable,
deferred: con.deferred,
foreign: rupg_catalog::ForeignKey { table: partition, keys, ..foreign },
};
let clone = self.catalog.add_fk_constraint(new, con.oid, FkSide::Referenced)?;
self.fk_recurse_referenced(clone, Some(triggers))?;
}
Ok(())
}
fn try_attach_foreign_key(
&mut self,
partition: u32,
parent: &Constraint,
keys: &[i16],
triggers: (u32, u32),
) -> Result<bool> {
let Some(wanted) = parent.foreign.as_ref() else { return Ok(false) };
let same_keys = |c: &&Constraint| {
c.kind == ConKind::Foreign
&& c.keys == keys
&& c.foreign.as_ref().is_some_and(|f| {
f.table == wanted.table && f.keys == wanted.keys && f.pf_eq == wanted.pf_eq
})
};
if let Some(c) = self
.catalog
.constraints_of(partition)
.find(|c| same_keys(c) && c.enforced != parent.enforced)
{
let relation = self.catalog.relation(partition).map(|r| r.name.clone());
return Err(Error::new(
SqlState::INVALID_OBJECT_DEFINITION,
format!(
"constraint \"{}\" enforceability conflicts with constraint \"{}\" on relation \"{}\"",
parent.name,
c.name,
relation.unwrap_or_default()
),
));
}
let found = self.catalog.constraints_of(partition).filter(same_keys).find(|c| {
c.parent == 0
&& c.deferrable == parent.deferrable
&& c.deferred == parent.deferred
&& c.foreign.as_ref().is_some_and(|f| {
f.update == wanted.update
&& f.delete == wanted.delete
&& f.match_type == wanted.match_type
})
});
let Some((oid, valid)) = found.map(|c| (c.oid, c.validated)) else { return Ok(false) };
if self.is_partitioned(wanted.table) {
let clones: Vec<u32> = self
.catalog
.constraints_of(partition)
.filter(|c| c.parent == oid)
.map(|c| c.oid)
.collect();
for clone in clones {
self.drop_fk_clone(clone)?;
}
}
let actions: Vec<u32> = self
.catalog
.triggers_of(wanted.table)
.into_iter()
.filter(|t| t.constraint == oid && t.constraint_relation == partition)
.map(|t| t.oid)
.collect();
for trigger in actions {
let object = ObjRef::new(PG_TRIGGER, trigger);
self.catalog.drop_depends_of(object);
self.perform(&[object], false)?;
}
self.catalog.set_constraint_parent(oid, parent.oid)?;
let (insert, update) = self.fk_triggers(oid, partition, true);
self.catalog.set_trigger_parent(insert, triggers.0);
self.catalog.set_trigger_parent(update, triggers.1);
if parent.validated && !valid {
self.validate_foreign_key(oid)?;
}
Ok(true)
}
pub(super) fn detach_foreign_keys(&mut self, partition: u32) -> Result<()> {
let own: Vec<Constraint> = self
.catalog
.constraints_of(partition)
.filter(|c| c.kind == ConKind::Foreign)
.cloned()
.collect();
let oids: Vec<u32> = own.iter().map(|c| c.oid).collect();
for con in own.iter().filter(|c| c.parent != 0 && !oids.contains(&c.parent)) {
self.catalog.clear_constraint_parent(con.oid)?;
let (insert, update) = self.fk_triggers(con.oid, partition, true);
self.catalog.clear_trigger_parent(insert);
self.catalog.clear_trigger_parent(update);
self.fk_recurse_referenced(con.oid, None)?;
}
for clone in self.parented_refs(partition) {
self.catalog.clear_constraint_parent(clone)?;
self.drop_fk_clone(clone)?;
}
Ok(())
}
fn drop_fk_clone(&mut self, clone: u32) -> Result<()> {
let triggers: Vec<u32> =
self.catalog.triggers().filter(|t| t.constraint == clone).map(|t| t.oid).collect();
for trigger in triggers {
self.catalog.clear_trigger_parent(trigger);
}
let object = ObjRef::new(PG_CONSTRAINT, clone);
self.catalog.drop_depends_of(object);
self.perform(&[object], false)
}
pub(super) fn detach_checks(&mut self, partition: u32) {
if !self.parented_refs(partition).is_empty() {
self.detached.push(partition);
}
}
fn parented_refs(&self, partition: u32) -> Vec<u32> {
self.catalog
.constraints()
.filter(|c| c.parent != 0 && c.foreign.as_ref().is_some_and(|f| f.table == partition))
.map(|c| c.oid)
.collect()
}
fn foreign_constraint(&self, oid: u32) -> Result<Constraint> {
self.catalog
.constraint(oid)
.filter(|c| c.kind == ConKind::Foreign)
.cloned()
.ok_or_else(|| Error::internal(format!("cache lookup failed for constraint {oid}")))
}
fn fk_triggers(&self, constraint: u32, table: u32, check: bool) -> (u32, u32) {
let (mut first, mut second) = (0, 0);
for trigger in self.catalog.triggers_of(table) {
if trigger.constraint != constraint {
continue;
}
let Some(function) = rupg_catalog::foreign_key_function(trigger.function) else {
continue;
};
if (function < 2) != check {
continue;
}
if function % 2 == 1 {
second = trigger.oid;
} else {
first = trigger.oid;
}
}
(first, second)
}
fn is_partitioned(&self, table: u32) -> bool {
self.catalog.relation(table).is_some_and(|r| r.kind == RelKind::Partitioned)
}
fn partition_oids(&self, table: u32) -> Vec<u32> {
self.catalog.partitions(table).iter().map(|r| r.oid).collect()
}
fn partition_index(&self, partition: u32, index: u32) -> Result<u32> {
self.catalog
.relations()
.find(|r| r.index.as_ref().is_some_and(|i| i.table == partition && i.parent == index))
.map(|r| r.oid)
.ok_or_else(|| {
let name = self.catalog.relation(partition).map_or("", |r| r.name.as_str());
Error::internal(format!("index for {index} not found in partition {name}"))
})
}
fn map_keys(&self, from: u32, to: u32, keys: &[i16]) -> Result<Vec<i16>> {
let (Some(from_rel), Some(to_rel)) =
(self.catalog.relation(from), self.catalog.relation(to))
else {
return Err(Error::internal(format!("cache lookup failed for relation {to}")));
};
keys.iter()
.map(|&key| {
from_rel.column(key).and_then(|c| to_rel.column_number(&c.name)).ok_or_else(|| {
Error::internal(format!("could not map column {key} of relation {from}"))
})
})
.collect()
}
}