#![deny(rust_2018_idioms)]
#![warn(missing_debug_implementations)]
use std::collections::BTreeSet;
#[derive(Debug, Clone)]
pub struct SchemaBinding {
pub constraint: &'static str,
pub allowed: BTreeSet<&'static str>,
}
pub trait SchemaConstrained: Sized + 'static {
const ALL: &'static [Self];
const CHECK_CONSTRAINTS: &'static [&'static str];
fn db_value(&self) -> &'static str;
fn bindings() -> Vec<SchemaBinding> {
let allowed: BTreeSet<&'static str> = Self::ALL.iter().map(Self::db_value).collect();
Self::CHECK_CONSTRAINTS
.iter()
.map(|&constraint| SchemaBinding {
constraint,
allowed: allowed.clone(),
})
.collect()
}
}
#[macro_export]
macro_rules! impl_schema_constrained {
(
$ty:ident via $method:ident {
all: [ $( $variant:ident ),+ $(,)? ],
$( non_stored: [ $( $ns:ident ),+ $(,)? ], )?
constraints: [ $( $constraint:literal ),+ $(,)? ] $(,)?
}
) => {
impl $crate::SchemaConstrained for $ty {
const ALL: &'static [Self] = &[ $( $ty::$variant ),+ ];
const CHECK_CONSTRAINTS: &'static [&'static str] = &[ $( $constraint ),+ ];
fn db_value(&self) -> &'static str {
self.$method()
}
}
const _: fn($ty) = |x| match x {
$( $ty::$variant => () ),+
$( , $( $ty::$ns => () ),+ )?
};
};
}
pub mod migration_scan {
use std::collections::BTreeSet;
pub fn value_set_constraints_in(sql: &str) -> BTreeSet<String> {
let flat: String = {
let mut out = String::with_capacity(sql.len());
let mut in_space = false;
for ch in sql.chars() {
if ch.is_whitespace() {
if !in_space {
out.push(' ');
}
in_space = true;
} else {
out.push(ch);
in_space = false;
}
}
out
};
let mut found = BTreeSet::new();
for (at, _) in flat.match_indices("CONSTRAINT ") {
let rest = &flat[at + "CONSTRAINT ".len()..];
let Some((name, tail)) = rest.split_once(' ') else {
continue;
};
let Some(body) = tail.strip_prefix("CHECK ") else {
continue;
};
let Some(body) = balanced(body) else {
continue;
};
if is_pure_value_set(body) {
found.insert(name.to_string());
}
}
found
}
fn balanced(s: &str) -> Option<&str> {
let s = s.strip_prefix('(')?;
let mut depth = 1usize;
for (i, ch) in s.char_indices() {
match ch {
'(' => depth += 1,
')' => {
depth -= 1;
if depth == 0 {
return Some(&s[..i]);
}
}
_ => {}
}
}
None
}
fn is_pure_value_set(expr: &str) -> bool {
let mut expr = expr.trim();
while let Some(inner) = balanced(expr) {
if inner.len() + 2 == expr.len() {
expr = inner.trim();
} else {
break;
}
}
let mut after = if expr.starts_with('(') {
match balanced(expr) {
Some(inner)
if !inner.is_empty()
&& inner.chars().all(|c| c.is_alphanumeric() || c == '_') =>
{
expr[inner.len() + 2..].trim_start()
}
_ => return false,
}
} else {
let column_len = expr
.find(|c: char| !(c.is_alphanumeric() || c == '_'))
.unwrap_or(expr.len());
if column_len == 0 {
return false;
}
expr[column_len..].trim_start()
};
loop {
if let Some(r) = after.strip_prefix(')') {
after = r.trim_start();
continue;
}
if let Some(r) = after.strip_prefix("::") {
after = r
.trim_start_matches(|ch: char| ch.is_alphanumeric() || ch == '_')
.trim_start();
continue;
}
break;
}
let list = if let Some(r) = after.strip_prefix("IN ") {
r.trim_start()
} else if let Some(r) = after.strip_prefix("= ANY ") {
r.trim_start()
} else {
return false;
};
match balanced(list) {
Some(inner) => list.len() == inner.len() + 2,
None => false,
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
mod tests {
#[test]
fn the_parser_sees_the_shapes_a_migration_is_actually_written_in() {
let must_find = [
(
"hand-written ALTER, name and CHECK on separate lines, IN spelling",
"ALTER TABLE scchat.t\n ADD CONSTRAINT ck_target\n CHECK (kind IN ('a', 'b'));",
),
(
"hand-written ALTER, separate lines, pg_dump spelling",
"ALTER TABLE scchat.t\n ADD CONSTRAINT ck_target\n CHECK ((kind = ANY (ARRAY['a'::text])));",
),
(
"inline in CREATE TABLE, IN spelling",
"CREATE TABLE scchat.t (\n kind text,\n CONSTRAINT ck_target CHECK (kind IN ('a', 'b'))\n);",
),
(
"cast between column and operator",
"ADD CONSTRAINT ck_target CHECK (((kind)::text = ANY (ARRAY['a'::text])));",
),
(
"single-line pg_dump form (the only one the first matcher caught)",
"ADD CONSTRAINT ck_target CHECK ((kind = ANY (ARRAY['a'::text, 'b'::text])));",
),
];
for (label, sql) in must_find {
assert!(
super::value_set_constraints_in(sql).contains("ck_target"),
"parser missed a real value-set CHECK — {label}\n input: {sql}"
);
}
let must_ignore = [
(
"multi-clause pairing check (this repo has one)",
"ADD CONSTRAINT ck_pairing CHECK ((plan = 'enterprise') = (monthly_message_limit IS NULL));",
),
(
"range check",
"ADD CONSTRAINT ck_range CHECK ((retention_days > 0));",
),
(
"regex check",
"ADD CONSTRAINT ck_regex CHECK ((ppnum ~ '^[0-9]{4}$'));",
),
(
"membership nested inside a compound predicate is not a value-set",
"ADD CONSTRAINT ck_compound CHECK ((a IS NULL) OR (kind IN ('a', 'b')));",
),
(
"membership as the FIRST clause of a compound predicate — PAS has\
exactly this shape in ck_signup_sessions_reservation_pair, and a\
parser that only checks the start of the expression accepts it",
"ADD CONSTRAINT ck_pair CHECK ((((state = ANY (ARRAY['a'::text])) AND\
(r IS NOT NULL)) OR ((state = 'b'::text) AND (r IS NULL))));",
),
];
for (label, sql) in must_ignore {
assert!(
super::value_set_constraints_in(sql).is_empty(),
"parser claimed a non-value-set CHECK — {label}\n input: {sql}"
);
}
}
}
}