systemprompt_database/lifecycle/installation/fk_deferral/
mod.rs

1//! Split `FOREIGN KEY` constraints out of a declarative `CREATE TABLE`.
2//!
3//! The installer creates every table in the structural phase, before any
4//! migration or `CREATE INDEX` has run. A foreign key needs a unique index on
5//! the referenced columns at the moment it is created, and on an existing
6//! database that index may only arrive through a migration — so a key
7//! declared inline installs on a fresh database and fails on an upgraded one
8//! with "no unique constraint matching given keys". The table is therefore
9//! created without its foreign keys, and each key is re-emitted as an
10//! `ALTER TABLE … ADD CONSTRAINT` the installer runs after every extension's
11//! migrations and dependent DDL, exactly as `pg_dump` orders a schema.
12//!
13//! A statement that declares no foreign key is returned verbatim, so error
14//! messages for such tables still quote the author's text. A statement that
15//! does is re-emitted through `pg_query`'s deparser, which preserves
16//! `IF NOT EXISTS`, column defaults, `CHECK`, `GENERATED`, identity, `UNIQUE`
17//! and `PRIMARY KEY` clauses; only comments and formatting are lost.
18//!
19//! Copyright (c) systemprompt.io — Business Source License 1.1.
20//! See <https://systemprompt.io> for licensing details.
21
22mod emit;
23
24use pg_query::NodeEnum;
25use pg_query::protobuf::node::Node;
26use pg_query::protobuf::{ColumnDef, ConstrType, Constraint, CreateStmt};
27use thiserror::Error;
28
29use emit::{deferred_key, string_node};
30
31#[derive(Debug, Error)]
32pub enum FkDeferralError {
33    #[error("SQL parse failed: {0}")]
34    Parse(#[source] pg_query::Error),
35    #[error("not a CREATE TABLE statement")]
36    NotCreateTable,
37    #[error("could not re-emit CREATE TABLE {table} without its foreign keys: {source}")]
38    DeparseTable {
39        table: String,
40        #[source]
41        source: pg_query::Error,
42    },
43    #[error("could not emit the deferred foreign key on {table}: {source}")]
44    DeparseKey {
45        table: String,
46        #[source]
47        source: pg_query::Error,
48    },
49    #[error("foreign key on {table} names no referenced table")]
50    NoReferencedTable { table: String },
51}
52
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct DeferredForeignKey {
55    pub table: String,
56    pub source_table: String,
57    pub columns: Vec<String>,
58    pub referenced_table: String,
59    pub referenced_columns: Vec<String>,
60    pub constraint_name: String,
61    pub sql: String,
62}
63
64#[derive(Debug, Clone, PartialEq, Eq)]
65pub struct SplitCreateTable {
66    pub create_table_sql: String,
67    pub foreign_keys: Vec<DeferredForeignKey>,
68}
69
70pub(super) fn split_foreign_keys(
71    original_sql: &str,
72    create: &CreateStmt,
73) -> Result<SplitCreateTable, FkDeferralError> {
74    let Some(relation) = create.relation.as_ref() else {
75        return Ok(verbatim(original_sql));
76    };
77
78    let mut stripped = create.clone();
79    let mut constraints: Vec<Constraint> = Vec::new();
80    stripped.table_elts = create
81        .table_elts
82        .iter()
83        .filter_map(|elt| match elt.node.as_ref() {
84            Some(Node::Constraint(c)) if is_foreign(c) => {
85                constraints.push((**c).clone());
86                None
87            },
88            Some(Node::ColumnDef(cd)) => {
89                let (column, mut found) = strip_column_references(cd);
90                constraints.append(&mut found);
91                Some(pg_query::protobuf::Node {
92                    node: Some(Node::ColumnDef(Box::new(column))),
93                })
94            },
95            _ => Some(elt.clone()),
96        })
97        .collect();
98
99    if constraints.is_empty() {
100        return Ok(verbatim(original_sql));
101    }
102
103    let create_table_sql = NodeEnum::CreateStmt(stripped).deparse().map_err(|source| {
104        FkDeferralError::DeparseTable {
105            table: relation.relname.clone(),
106            source,
107        }
108    })?;
109
110    let foreign_keys = constraints
111        .into_iter()
112        .map(|c| deferred_key(relation, c))
113        .collect::<Result<Vec<_>, _>>()?;
114
115    Ok(SplitCreateTable {
116        create_table_sql,
117        foreign_keys,
118    })
119}
120
121pub fn split_create_table_foreign_keys(sql: &str) -> Result<SplitCreateTable, FkDeferralError> {
122    let parsed = pg_query::parse(sql).map_err(FkDeferralError::Parse)?;
123    let create = parsed
124        .protobuf
125        .stmts
126        .iter()
127        .find_map(|raw| match raw.stmt.as_ref()?.node.as_ref()? {
128            Node::CreateStmt(create) => Some(create),
129            _ => None,
130        })
131        .ok_or(FkDeferralError::NotCreateTable)?;
132    split_foreign_keys(sql, create)
133}
134
135fn verbatim(original_sql: &str) -> SplitCreateTable {
136    SplitCreateTable {
137        create_table_sql: original_sql.to_owned(),
138        foreign_keys: Vec::new(),
139    }
140}
141
142fn is_foreign(c: &Constraint) -> bool {
143    ConstrType::try_from(c.contype) == Ok(ConstrType::ConstrForeign)
144}
145
146fn strip_column_references(cd: &ColumnDef) -> (ColumnDef, Vec<Constraint>) {
147    let mut column = (*cd).clone();
148    let mut kept = Vec::with_capacity(cd.constraints.len());
149    let mut found = Vec::new();
150    let mut last_primary_is_foreign = false;
151
152    for node in &cd.constraints {
153        let Some(Node::Constraint(c)) = node.node.as_ref() else {
154            kept.push(node.clone());
155            continue;
156        };
157        if is_foreign(c) {
158            let mut key = (**c).clone();
159            key.fk_attrs = vec![string_node(&cd.colname)];
160            found.push(key);
161            last_primary_is_foreign = true;
162            continue;
163        }
164        // Why: `REFERENCES t DEFERRABLE INITIALLY DEFERRED` parses as three
165        // sibling constraints; Postgres attaches the attributes to the last
166        // key-like constraint (`lastprimarycon`: PRIMARY KEY, UNIQUE, FOREIGN
167        // KEY or EXCLUDE) on the column, so an attribute is folded into the
168        // deferred key only when that is the foreign key.
169        if last_primary_is_foreign
170            && let Some(previous) = found.last_mut()
171            && fold_attribute(previous, c)
172        {
173            continue;
174        }
175        if is_key_like(c) {
176            last_primary_is_foreign = false;
177        }
178        kept.push(node.clone());
179    }
180
181    column.constraints = kept;
182    (column, found)
183}
184
185fn is_key_like(c: &Constraint) -> bool {
186    matches!(
187        ConstrType::try_from(c.contype),
188        Ok(ConstrType::ConstrPrimary | ConstrType::ConstrUnique | ConstrType::ConstrExclusion)
189    )
190}
191
192fn fold_attribute(key: &mut Constraint, attribute: &Constraint) -> bool {
193    match ConstrType::try_from(attribute.contype) {
194        Ok(ConstrType::ConstrAttrDeferrable) => key.deferrable = true,
195        Ok(ConstrType::ConstrAttrNotDeferrable) => key.deferrable = false,
196        Ok(ConstrType::ConstrAttrDeferred) => {
197            key.deferrable = true;
198            key.initdeferred = true;
199        },
200        Ok(ConstrType::ConstrAttrImmediate) => key.initdeferred = false,
201        _ => return false,
202    }
203    true
204}
systemprompt_database/lifecycle/installation/fk_deferral/mod.rs

systemprompt_database/lifecycle/installation/fk_deferral/
mod.rs
