use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};
use autumn_schema_core::{Column, ColumnType, Table};
use diesel::connection::SimpleConnection as _;
use diesel::{Connection as _, PgConnection, RunQueryDsl as _, sql_query};
use serde::Deserialize;
use crate::migrate;
use crate::schema::introspect;
pub mod sample;
use super::{quote_ident, quote_literal};
pub const SCRUB_CONFIG_FILE: &str = "scrub.toml";
const TOKEN_HEX_LEN: usize = 32;
const UNIQUE_TOKEN_HEX_LEN: usize = 64;
const MIN_TOKEN_WIDTH: usize = 8;
const MIN_UNIQUE_TOKEN_WIDTH: usize = 16;
const SCRUB_EMAIL_DOMAIN: &str = "@example.invalid";
#[derive(Debug, Clone, Default)]
#[allow(clippy::struct_excessive_bools)]
pub struct ScrubArgs {
pub profile: Option<String>,
pub artifact: Option<PathBuf>,
pub output: Option<PathBuf>,
pub config: Option<PathBuf>,
pub check: bool,
pub dry_run: bool,
pub force: bool,
pub allow_source_overwrite: bool,
pub sample: Vec<String>,
pub seed: u64,
}
#[derive(Debug)]
pub enum ScrubError {
Unclassified {
columns: Vec<String>,
},
StaleConfig {
entries: Vec<String>,
},
Contradiction {
columns: Vec<String>,
},
SafeOverridesEncrypted {
columns: Vec<String>,
},
PiiOnKeyColumn {
columns: Vec<String>,
},
NoAutoStrategy {
column: String,
detail: String,
},
NullOnNotNull {
column: String,
},
NonUniqueStrategy {
column: String,
strategy: &'static str,
},
StrategyTypeMismatch {
column: String,
strategy: &'static str,
detail: String,
},
ColumnTooNarrow {
column: String,
limit: usize,
overhead: usize,
floor: usize,
},
PlaintextIntoEncrypted {
columns: Vec<String>,
},
CheckConstrainedColumn {
column: String,
},
EncryptionKeyUnavailable {
profile: String,
detail: String,
},
EncryptedMetadataUnavailable,
InaccessibleTables {
tables: Vec<String>,
},
UnsupportedSchemas {
schemas: Vec<String>,
},
LegacyInheritance {
tables: Vec<String>,
},
ReplicaSessionRole {
role: String,
},
UnprintableTarget {
targets: Vec<String>,
},
UnprintableEncryptedRewrite {
columns: Vec<String>,
},
UnprintableProductionTarget {
profile: String,
},
UnrefreshableViews {
views: Vec<String>,
},
UnprintableAmbiguousTarget {
targets: Vec<String>,
},
UntraceableViewFunction {
views: Vec<String>,
},
UnprintablePortlessTarget {
targets: Vec<String>,
},
UnprintableHostaddrTarget {
targets: Vec<String>,
},
ViewReadsRelationsByName {
views: Vec<String>,
},
UnprintableMultiHostTarget {
targets: Vec<String>,
},
EmptyingTriggerLeak {
tables: Vec<String>,
},
RowLevelSecurity {
tables: Vec<String>,
},
FrameworkTableDeclared {
tables: Vec<String>,
},
PurgeSchemaBookkeeping {
tables: Vec<String>,
},
Config {
path: String,
detail: String,
},
SourceScan {
detail: String,
},
UnresolvableAnonymize {
detail: String,
},
Introspect {
label: String,
detail: String,
},
Sql(String),
ProductionRefused {
profile: String,
},
OverwritesConfiguredTarget {
profile: String,
database: String,
},
PurgeNotFrameworkTable {
tables: Vec<String>,
},
PurgeLedgerTablesUnpaired {
listed: String,
missing: String,
},
Backup(Box<super::backup::BackupError>),
Sample(Box<sample::SampleError>),
}
impl std::fmt::Display for ScrubError {
#[allow(clippy::too_many_lines)]
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Unclassified { columns } => write!(
f,
"{} column(s) are neither PII-classified nor declared safe, so the scrub \
cannot prove they carry no real data:\n{}\n \
Declare each one in {SCRUB_CONFIG_FILE} — under [tables.<table>.pii] to \
replace it, or in `safe` to keep it verbatim.",
columns.len(),
bullet_list(columns),
),
Self::StaleConfig { entries } => write!(
f,
"{SCRUB_CONFIG_FILE} names {} table(s)/column(s) the database does not have:\n{}\n \
The declaration has drifted from the schema — remove or rename the stale \
entries (a renamed column must be re-declared under its new name).",
entries.len(),
bullet_list(entries),
),
Self::Contradiction { columns } => write!(
f,
"{} column(s) are declared both `safe` and PII in {SCRUB_CONFIG_FILE}:\n{}\n \
Pick one.",
columns.len(),
bullet_list(columns),
),
Self::SafeOverridesEncrypted { columns } => write!(
f,
"{} column(s) carry #[encrypted] in the model but are declared `safe` in \
{SCRUB_CONFIG_FILE}:\n{}\n \
An at-rest-encrypted column is PII by construction and cannot be declared \
safe. Remove the `safe` entry (or drop #[encrypted] from the model if the \
column really is not sensitive).",
columns.len(),
bullet_list(columns),
),
Self::PiiOnKeyColumn { columns } => write!(
f,
"{} column(s) are primary or foreign keys but declared PII:\n{}\n \
Rewriting a key column would break referential integrity. Scrub the \
referenced table's own PII columns instead, and declare these `safe`.",
columns.len(),
bullet_list(columns),
),
Self::NoAutoStrategy { column, detail } => write!(
f,
"No replacement can be derived for {column:?} from its type ({detail}).\n \
Declare an explicit strategy for it under [tables.<table>.pii] in \
{SCRUB_CONFIG_FILE} (for example `= \"redact\"`), or declare the column \
`safe`."
),
Self::NullOnNotNull { column } => write!(
f,
"{column:?} is NOT NULL, so the `null` strategy would violate the column \
constraint.\n Use `redact` (or another value-producing strategy) instead."
),
Self::NonUniqueStrategy { column, strategy } => write!(
f,
"{column:?} is UNIQUE, but the `{strategy}` strategy writes the same value \
into every row and would violate the unique constraint.\n \
Use a per-row-unique strategy (`redact`, `email`, `name`, `uuid`, `bytes`)."
),
Self::StrategyTypeMismatch {
column,
strategy,
detail,
} => write!(
f,
"The `{strategy}` strategy cannot produce a value for {column:?} ({detail}).\n \
Pick a strategy that matches the column type."
),
Self::ColumnTooNarrow {
column,
limit,
overhead,
floor,
} => write!(
f,
"{column:?} holds at most {limit} characters, but the chosen strategy needs \
{overhead} for its fixed text plus at least {floor} more for a per-row token.\n \
Strategies by fixed overhead: `phone` (5), `name` (9), `redact` (11), \
`email` (25). Pick one that fits, use `null` if the column is nullable, \
widen the column, or declare it `safe`."
),
Self::PlaintextIntoEncrypted { columns } => write!(
f,
"{} column(s) carry #[encrypted] in the model but {SCRUB_CONFIG_FILE} declares a \
plaintext strategy for them:\n{}\n \
Writing a plain string into an at-rest-encrypted column makes every later read \
of that row fail as malformed ciphertext. An #[encrypted] column is \
re-encrypted automatically — remove the declaration, or use `null` if the \
column is nullable.",
columns.len(),
bullet_list(columns),
),
Self::CheckConstrainedColumn { column } => write!(
f,
"{column:?} is covered by a CHECK constraint, so no fabricated value can be \
proven to satisfy it (a closed-set column reaches the database as TEXT plus a \
CHECK, which is why the type alone does not reveal this).\n \
Declare the column `safe` if the constraint means it holds no free-form PII, \
or drop the constraint on the copy before scrubbing."
),
Self::EncryptionKeyUnavailable { profile, detail } => write!(
f,
"This database has #[encrypted] columns, but no encryption key could be resolved \
for the {profile:?} profile: {detail}\n \
A scrub must write a VALID ciphertext envelope into an encrypted column — a \
plaintext replacement would make every later read of that row fail. Provide the \
target's `active_record_encryption` credentials (`autumn credentials edit`), or \
declare those columns `null` in {SCRUB_CONFIG_FILE} if they are nullable."
),
Self::EncryptedMetadataUnavailable => write!(
f,
"No model source was found, so the scrub cannot tell which columns are \
#[encrypted].\n \
That matters more than it sounds: an unrecognised encrypted column declared \
`safe` keeps its production ciphertext, and one given a plaintext strategy \
becomes permanently unreadable. Run the scrub from the project root (where \
`src/models` lives), or name them — with their mode — in {SCRUB_CONFIG_FILE}:\n \
[tables.users.encrypted]\n api_token = \"randomized\"\n \
email = \"deterministic\"\n \
An app with no encrypted columns at all still needs one empty \
`[tables.<any>.encrypted]` section to say so deliberately."
),
Self::InaccessibleTables { tables } => write!(
f,
"{} table(s) exist in `public` but could not be read by the connecting \
role:\n{}\n \
They never reached the classifier, so the scrub cannot claim they carry no \
PII — and it will not rewrite them either. Connect as a role that can see \
every table (the owner, or one with the needed privileges).",
tables.len(),
bullet_list(tables),
),
Self::UnsupportedSchemas { schemas } => write!(
f,
"This database has base tables in {} schema(s) outside `public`:\n{}\n \
The classification universe is `public`-only, so a scrub cannot prove those \
tables carry no PII and refuses rather than reporting a completeness it did \
not check. Scrub those schemas separately, or drop them from the copy.",
schemas.len(),
bullet_list(schemas),
),
Self::LegacyInheritance { tables } => write!(
f,
"This database uses legacy table inheritance on {} table(s):\n{}\n \
Unlike a declarative partition, an inheritance child is an ordinary table \
that the sample plans separately — while `DELETE FROM parent` and \
`SELECT ... FROM parent` reach its rows too, because they are not written \
`ONLY parent`. Parent and child would then select rows independently and \
delete each other's, and the run would report a success it cannot stand \
behind. Sample a copy without the inheritance, or drop the child tables \
from it.",
tables.len(),
bullet_list(tables),
),
Self::UnprintableTarget { targets } => write!(
f,
"`--dry-run` cannot print a runnable script for {} target(s):\n{}\n \
Each block it prints is destructive, and the `\\connect` line above it \
is what points psql at the right database. These targets are configured \
with a keyword-form connection string (`host=... password=...`), which \
cannot be printed with its password removed for certain — libpq allows \
whitespace around the `=` and quoted values with escapes — or with a \
connection-string parameter this command does not recognise, which it \
refuses to print rather than guess about (`password_free_conninfo` \
allowlists known-non-secret keywords and strips known credentials; \
anything else fails closed). A script without the boundary does not fail: \
it runs this target's plan against whichever database the pasting session \
is already on, which in a multi-target stream is the previous target. \
Configure the target as a URI (`postgres://user@host/db`), or run without \
`--dry-run`.",
targets.len(),
bullet_list(targets),
),
Self::UntraceableViewFunction { views } => write!(
f,
"{} materialized view(s) read through a function this run cannot \
follow:\n{}\n \
Refresh order is taken from the dependency graph `PostgreSQL` records, \
and it records nothing about what a function whose body it cannot \
parse reads — a string literal, `plpgsql`, or C. Measured on `a_report -> bridge_fn() -> z_source`: \
`pg_depend` holds `a_report -> pg_proc(bridge_fn)` and the function \
holds NO relation dependency at all, so the two views sorted by name, \
`a_report` refreshed FIRST from a `z_source` still holding pre-scrub \
rows, and the run reported success with `users` at 2 rows, 0 original \
addresses in the base tables and in `z_source`, and all 200 still in \
`a_report`. A `BEGIN ATOMIC` body IS tracked and is followed normally. \
Rewrite the function with `BEGIN ATOMIC`, inline its query into the \
view, or drop the view and rebuild it after the run.",
views.len(),
bullet_list(views),
),
Self::ViewReadsRelationsByName { views } => write!(
f,
"{} materialized view(s) read relations named in a string the server \
executes:\n{}\n \
`query_to_xml` takes its query as text, and `schema_to_xml` and \
`database_to_xml` take no relation argument at all, so PostgreSQL \
records no dependency on anything they read — measured, a view \
defined `SELECT query_to_xml('SELECT email FROM z_source', …)` records \
only itself. Refresh order comes from those dependencies, so the views \
sorted by name instead: measured, the dependent refreshed FIRST from a \
stale source and kept all 200 pre-scrub addresses under a reported \
success. Nothing in the catalog can recover the edge, so it is refused \
rather than ordered around a gap. `table_to_xml` is fine and not \
refused: its `regclass` argument IS recorded. Name the relation in the \
view's own query, or take the view out of the database this run \
scrubs.",
views.len(),
bullet_list(views),
),
Self::UnprintableMultiHostTarget { targets } => write!(
f,
"`--dry-run` cannot print a runnable script for {} target(s):\n{}\n \
Their connection strings name more than one endpoint, and libpq chooses \
a member per connection — measured on 16.13, 20 connections with \
`load_balance_hosts=random` split 7/13 across two members. The run sizes \
the sample on one connection and the pasted script runs on another, so \
the root `LIMIT` need not describe the database it lands on: measured \
against a two-member URI whose first member held 200 rows and whose \
second held none, ten dry runs printed `LIMIT 2` six times and `LIMIT 0` \
four times, and `LIMIT 0` selects no root rows at all, so the delete pass \
empties the table instead of sampling it. The endpoint proof cannot \
separate the members either — it has to accept any of them, which is the \
discriminator it exists for. Name one endpoint \
(`postgres://user@host:5432/db`), or run without `--dry-run`, where the \
command sizes and writes on the one connection it holds.",
targets.len(),
bullet_list(targets),
),
Self::UnprintableHostaddrTarget { targets } => write!(
f,
"`--dry-run` cannot print a runnable script for {} target(s):\n{}\n \
Something other than `host` and `port` chooses their endpoint — \
`hostaddr` in the connection string, `PGHOSTADDR` in the environment, \
or a service file named by `service=` or `PGSERVICE`. All three were \
measured reaching 127.0.0.1 through `host=not-a-real-host.example`, a \
name that does not resolve at all, with psql still reporting \
`HOST=not-a-real-host.example`. psql exposes no `:HOSTADDR` to pin \
instead: `\\echo [:HOSTADDR]` prints the name back unexpanded. So two \
targets sharing a host, port and database name but reaching different \
servers are indistinguishable to the reconnect proof, which is the one \
thing it exists to tell apart — and a pasting session brings its own \
environment, so what this run resolved need not be what it resolves. \
Name the endpoint in `host` (`postgres://user@10.0.0.7:5432/db`) with \
no `hostaddr`, `service`, `PGHOSTADDR` or `PGSERVICE` in play, or run \
without `--dry-run`, where the command holds its own connection and \
never has to prove which one it is.",
targets.len(),
bullet_list(targets),
),
Self::UnprintablePortlessTarget { targets } => write!(
f,
"`--dry-run` cannot print a runnable script for {} target(s):\n{}\n \
Their connection strings name a host but no port, so psql resolves one \
when the script is pasted — from `PGPORT`, or 5432 — and that need not \
be what libpq resolved while this run planned. The block's proof that \
`\\connect` reached the intended endpoint would then have to accept ANY \
port for that host, which is the discriminator it exists for: measured, \
the same URI resolves to 5433 under `PGPORT=5433` and to 5432 without \
it, two different servers. Embedding the port this run resolved is not \
an option either — libpq's own resolution can come from a service file \
this command does not read. State the port in the connection string \
(`postgres://user@host:5432/db`), or run without `--dry-run`, where the \
command holds its own connection and never has to prove which one it \
is.",
targets.len(),
bullet_list(targets),
),
Self::UnprintableEncryptedRewrite { columns } => write!(
f,
"`--dry-run` cannot print a runnable script: {} column(s) are \
#[encrypted]:\n{}\n \
The scrub replaces each of these with an AEAD envelope sealed per row \
under the target's key, so there is no SQL text for it that does not \
embed the key. Printing the rest and omitting these is worse than \
printing nothing: the script is meant to be pasted into psql, and \
nothing in a pasted stream can stop it partway — an aborted \
transaction ends at the next `COMMIT`, after which the following \
target's `\\connect` and its DELETEs run for real. So the whole \
script is withheld, not just the statements that cannot be written. \
The plan above is complete and accurate — run `autumn db scrub` \
without --dry-run to apply it.",
columns.len(),
bullet_list(columns),
),
Self::UnprintableProductionTarget { profile } => write!(
f,
"`--dry-run` cannot print a runnable script for the {profile:?} profile.\n \
A scrub REWRITES data in place, and this command refuses to run against \
{profile:?} without `--force`. The script it prints is executable and its \
`\\connect` line names that same database — with the password removed on \
purpose, which is what `.pgpass` is for — so printing it would hand over \
the run the profile guard exists to stop. The plan above is complete and \
accurate; add `--force` to print the script too, or point `--profile` at a \
staging target.",
),
Self::UnrefreshableViews { views } => write!(
f,
"{} materialized view(s) cannot be refreshed in dependency order:\n{}\n \
The run refreshes views so each is rebuilt from scrubbed data, and it \
orders them by walking `pg_depend`. That walk stops at a fixed depth, so \
a chain longer than it leaves these views unreached — and a view that is \
not refreshed keeps the rows it selected BEFORE the scrub, including the \
values the run just removed from the tables it reads. Measured on a \
36-deep chain: 33 views refreshed, `users` left with 0 original \
addresses, and the deepest view still holding all 200. Refusing is the \
only honest answer until the walk covers the whole graph — shorten the \
chain, or drop the views this run cannot reach and rebuild them after \
it.",
views.len(),
bullet_list(views),
),
Self::UnprintableAmbiguousTarget { targets } => write!(
f,
"`--dry-run` cannot print a runnable script for {} target(s):\n{}\n \
These are reached over a Unix socket, where the server reports no \
address and no port — and nothing else it reports identifies the \
instance either. A physical copy carries its origin's \
`system_identifier`; the configured port is shared by two clusters on \
different socket directories; and `data_directory` is server-LOCAL, so \
two containers each answering `/var/lib/postgresql/data` match on it \
while being different databases. Every value the guard can ask for is \
either cloned or container-local, so a printed block cannot tell this \
target from a copy of it — and `\\connect` keeps the PREVIOUS connection \
when it fails. Measured on two clusters sharing port 5433: the clone's \
script pasted at its origin passed the guard, committed, and took the \
ORIGIN from 200 users to 25. Connect over TCP, where the address and \
port identify the endpoint and two clones cannot hold the same pair, or \
run without --dry-run — the command opens its own connection and cannot \
be on the wrong database.",
targets.len(),
bullet_list(targets),
),
Self::ReplicaSessionRole { role } => write!(
f,
"This connection runs with `session_replication_role = {role}`, not `origin`.\n \
That inverts which hooks fire: a trigger or rule marked `ENABLE REPLICA` runs \
and an ordinary one does not, so every check here that asks whether a `DELETE` \
can execute code is answering for the wrong session. A replica-only archive \
trigger on a table this run empties would copy the rows it removes into a \
table nothing verifies, past a refusal that never saw it. Connect without \
`options=-c session_replication_role=...`, or reset it before scrubbing.",
),
Self::EmptyingTriggerLeak { tables } => write!(
f,
"{} table(s) this run empties carry a user-defined trigger or rewrite rule \
that fires on `DELETE`:\n{}\n \
Emptying them is the LAST thing the run writes \u{2014} it has to be, because a \
trigger on a scrubbed table can otherwise re-fill them with the very PII \
being removed. Anything attached to a `DELETE` on one of them therefore runs \
after every column rewrite, and can copy the rows it is removing into an \
ordinary table that has already been scrubbed. Nothing downstream can catch \
that: the run verifies these tables are empty, not where their triggers and \
rules wrote. On the copy, `ALTER TABLE ... DISABLE TRIGGER ...` or \
`DROP RULE ... ON ...` before scrubbing.",
tables.len(),
bullet_list(tables),
),
Self::RowLevelSecurity { tables } => write!(
f,
"{} table(s) this run reads or writes have row-level security enabled:\n{}\n \
A role that does not bypass RLS sees only the rows its policies expose, and \
every phase then reports a success it cannot stand behind: an `UPDATE` \
rewrites part of the PII, a `DELETE` reports a table emptied that is not, and \
the foreign key re-count misses the very orphan it exists to find. \
Connect as the table owner or a BYPASSRLS role.",
tables.len(),
bullet_list(tables),
),
Self::FrameworkTableDeclared { tables } => write!(
f,
"{SCRUB_CONFIG_FILE} declares {} framework-owned table(s) under [tables.*]:\n{}\n \
Framework-owned tables are deliberately outside the column classification \
(exactly as `autumn db pull` and `autumn schema pull` skip them), so a \
per-column declaration for them has no effect. Use \
`[framework] purge = [...]` to empty one instead.",
tables.len(),
bullet_list(tables),
),
Self::PurgeSchemaBookkeeping { tables } => write!(
f,
"[framework] purge in {SCRUB_CONFIG_FILE} names {} table(s) that hold schema \
bookkeeping, not app payloads:\n{}\n \
Emptying them would make the copy un-migratable or un-routable. Remove them \
from `purge`.",
tables.len(),
bullet_list(tables),
),
Self::Config { path, detail } => write!(f, "Cannot read {path}: {detail}"),
Self::SourceScan { detail } => write!(
f,
"Could not scan the app source for PII annotations: {detail}"
),
Self::UnresolvableAnonymize { detail } => write!(
f,
"A GDPR anonymize registration names a table this scanner cannot resolve: \
{detail}\n `autumn db scrub` reads `ModelRegistration::anonymize(\"...\")` \
with a string-literal table name. Use a literal there, or declare the \
table's columns explicitly in {SCRUB_CONFIG_FILE}."
),
Self::Introspect { label, detail } => {
write!(f, "Could not read the schema of {label}: {detail}")
}
Self::Sql(message) => write!(f, "{message}"),
Self::ProductionRefused { profile } => write!(
f,
"Refusing to scrub the {profile:?} profile database.\n \
A scrub REWRITES data in place — running it against production would \
destroy the real values. Point --profile at your staging/dev target, or \
re-run with --force if you really mean it."
),
Self::OverwritesConfiguredTarget { profile, database } => write!(
f,
"Refusing to scrub {database:?}: it is the database the {profile:?} profile's \
config file declares.\n \
The scrub would overwrite the source the artifact was taken from. Point the \
target at a separate staging database, or re-run with \
--allow-source-overwrite if that really is what you mean."
),
Self::PurgeNotFrameworkTable { tables } => write!(
f,
"[framework] purge in {SCRUB_CONFIG_FILE} names {} table(s) that are not \
framework-owned:\n{}\n \
`purge` empties a table outright and only accepts framework-owned names \
(`autumn_*` / `_autumn*`, plus the framework's unprefixed tables). Declare \
a user table's columns under [tables.<table>.pii] instead.",
tables.len(),
bullet_list(tables),
),
Self::PurgeLedgerTablesUnpaired { listed, missing } => write!(
f,
"[framework] purge in {SCRUB_CONFIG_FILE} names {listed:?} but not \
{missing:?}.\n \
The ledger's two tables are emptied together or not at all: a high-water \
mark is built to outlive the revisions it names, so a copy holding one \
without the other makes `ledger_verify` report every ledgered record as \
tampered — and, with the marks kept, makes the write path refuse every \
subsequent write. Add {missing:?} to `purge`, or remove {listed:?}."
),
Self::Backup(e) => write!(f, "{e}"),
Self::Sample(e) => write!(f, "{e}"),
}
}
}
impl From<sample::SampleError> for ScrubError {
fn from(e: sample::SampleError) -> Self {
Self::Sample(Box::new(e))
}
}
impl From<super::backup::BackupError> for ScrubError {
fn from(e: super::backup::BackupError) -> Self {
Self::Backup(Box::new(e))
}
}
fn bullet_list(items: &[String]) -> String {
items
.iter()
.map(|item| format!(" - {item}"))
.collect::<Vec<_>>()
.join("\n")
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Strategy {
Auto,
Email,
Name,
Phone,
Redact,
Null,
Uuid,
Bytes,
Json,
Zero,
Epoch,
Encrypted,
}
impl Strategy {
const fn as_str(self) -> &'static str {
match self {
Self::Auto => "auto",
Self::Email => "email",
Self::Name => "name",
Self::Phone => "phone",
Self::Redact => "redact",
Self::Null => "null",
Self::Uuid => "uuid",
Self::Bytes => "bytes",
Self::Json => "json",
Self::Zero => "zero",
Self::Epoch => "epoch",
Self::Encrypted => "encrypted",
}
}
const fn allowed_on_unique(self) -> bool {
matches!(
self,
Self::Email
| Self::Name
| Self::Redact
| Self::Null
| Self::Uuid
| Self::Bytes
| Self::Encrypted
)
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ScrubDefaults {
#[serde(default)]
pub safe_columns: Vec<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum EncryptionMode {
Randomized,
Deterministic,
}
impl EncryptionMode {
const fn is_deterministic(self) -> bool {
matches!(self, Self::Deterministic)
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct TableRule {
#[serde(default)]
pub safe: Vec<String>,
#[serde(default)]
pub pii: BTreeMap<String, Strategy>,
#[serde(default)]
pub encrypted: BTreeMap<String, EncryptionMode>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct FrameworkRule {
#[serde(default)]
pub purge: Vec<String>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ScrubConfig {
#[serde(default)]
pub defaults: ScrubDefaults,
#[serde(default)]
pub tables: BTreeMap<String, TableRule>,
#[serde(default)]
pub framework: FrameworkRule,
#[serde(default)]
pub sample: sample::SampleRules,
}
const FRAMEWORK_PAYLOAD_TABLES: &[&str] = &[
"_autumn_ledger_revisions",
"_autumn_ledger_high_water",
"_autumn_version_history",
"api_tokens",
"autumn_experiment_assignments",
"autumn_experiment_changes",
"autumn_experiment_overrides",
"autumn_feature_flags",
"autumn_job_tracking",
"autumn_jobs",
"autumn_repository_commit_hooks",
"autumn_runtime_config_changes",
"autumn_runtime_config_values",
"autumn_search_documents",
"autumn_sync_applied",
"autumn_sync_pending",
"autumn_sync_rows",
"feature_flag_changes",
];
const UNPREFIXED_FRAMEWORK_TABLES: &[&str] = &[
"api_tokens",
"feature_flag_changes",
"__diesel_schema_migrations",
];
fn is_framework_table(table: &str) -> bool {
table.starts_with("autumn_")
|| table.starts_with("_autumn")
|| UNPREFIXED_FRAMEWORK_TABLES.contains(&table)
}
const NEVER_PURGEABLE_TABLES: &[&str] = &[
"__diesel_schema_migrations",
"_autumn_shard_directory",
"_autumn_shard_map",
"autumn_migration_checksums",
];
fn check_purge_list(config: &ScrubConfig) -> Result<(), ScrubError> {
let mut bookkeeping: Vec<String> = config
.framework
.purge
.iter()
.filter(|t| NEVER_PURGEABLE_TABLES.contains(&t.as_str()))
.cloned()
.collect();
if !bookkeeping.is_empty() {
bookkeeping.sort();
return Err(ScrubError::PurgeSchemaBookkeeping {
tables: bookkeeping,
});
}
let mut offenders: Vec<String> = config
.framework
.purge
.iter()
.filter(|t| !is_framework_table(t))
.cloned()
.collect();
if !offenders.is_empty() {
offenders.sort();
return Err(ScrubError::PurgeNotFrameworkTable { tables: offenders });
}
check_ledger_purge_pairing(config)
}
fn check_ledger_purge_pairing(config: &ScrubConfig) -> Result<(), ScrubError> {
let has = |table: &str| config.framework.purge.iter().any(|t| t == table);
let (revisions, marks) = (has(LEDGER_REVISIONS_TABLE), has(LEDGER_HIGH_WATER_TABLE));
if revisions == marks {
return Ok(());
}
let (listed, missing) = if revisions {
(LEDGER_REVISIONS_TABLE, LEDGER_HIGH_WATER_TABLE)
} else {
(LEDGER_HIGH_WATER_TABLE, LEDGER_REVISIONS_TABLE)
};
Err(ScrubError::PurgeLedgerTablesUnpaired {
listed: listed.to_owned(),
missing: missing.to_owned(),
})
}
const LEDGER_REVISIONS_TABLE: &str = "_autumn_ledger_revisions";
const LEDGER_HIGH_WATER_TABLE: &str = "_autumn_ledger_high_water";
#[cfg(test)]
fn parse_config_str(src: &str) -> Result<ScrubConfig, ScrubError> {
parse_config_at(src, Path::new(SCRUB_CONFIG_FILE))
}
fn parse_config_at(src: &str, path: &Path) -> Result<ScrubConfig, ScrubError> {
toml::from_str(src).map_err(|e| ScrubError::Config {
path: path.display().to_string(),
detail: e.to_string(),
})
}
fn load_config(explicit: Option<&Path>) -> Result<ScrubConfig, ScrubError> {
load_config_at(explicit, Path::new(SCRUB_CONFIG_FILE))
}
fn load_config_at(explicit: Option<&Path>, default: &Path) -> Result<ScrubConfig, ScrubError> {
let path = explicit.map_or_else(|| default.to_path_buf(), Path::to_path_buf);
match std::fs::read_to_string(&path) {
Ok(src) => parse_config_at(&src, &path),
Err(e) if e.kind() == std::io::ErrorKind::NotFound && explicit.is_none() => {
Ok(ScrubConfig::default())
}
Err(e) => Err(ScrubError::Config {
path: path.display().to_string(),
detail: e.to_string(),
}),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ClassSource {
Config,
Encrypted,
GdprAnonymize,
}
impl ClassSource {
const fn as_str(self) -> &'static str {
match self {
Self::Config => "declared",
Self::Encrypted => "#[encrypted]",
Self::GdprAnonymize => "gdpr:anonymize",
}
}
}
pub struct ClassificationInputs<'a> {
pub tables: &'a [Table],
pub config: &'a ScrubConfig,
pub encrypted: &'a BTreeMap<String, BTreeMap<String, bool>>,
pub anonymize_tables: &'a BTreeSet<String>,
pub facts: &'a DatabaseFacts,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ColumnPlan {
pub column: String,
pub strategy: Strategy,
pub source: ClassSource,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EncryptedRewrite {
pub column: String,
pub deterministic: bool,
pub shape: Strategy,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TablePlan {
pub table: String,
pub columns: Vec<ColumnPlan>,
pub sql: Option<String>,
pub row_key: String,
pub encrypted: Vec<EncryptedRewrite>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ScrubPlan {
pub tables: Vec<TablePlan>,
}
impl ScrubPlan {
#[cfg(test)]
fn column(&self, table: &str, column: &str) -> Option<&ColumnPlan> {
self.tables
.iter()
.find(|t| t.table == table)?
.columns
.iter()
.find(|c| c.column == column)
}
#[must_use]
pub fn column_count(&self) -> usize {
self.tables.iter().map(|t| t.columns.len()).sum()
}
}
#[allow(clippy::too_many_lines)]
pub fn build_plan(inputs: &ClassificationInputs<'_>) -> Result<ScrubPlan, ScrubError> {
let by_name: BTreeMap<&str, &Table> =
inputs.tables.iter().map(|t| (t.name.as_str(), t)).collect();
check_config_freshness(inputs, &by_name)?;
check_contradictions(inputs)?;
check_safe_overrides_encrypted(inputs, &by_name)?;
let mut unclassified = Vec::new();
let mut key_pii = Vec::new();
let mut plaintext_into_encrypted: Vec<(String, &'static str)> = Vec::new();
let mut planned: Vec<(&Table, Vec<(ColumnPlan, &Column)>)> = Vec::new();
for table in inputs.tables {
if inputs.facts.partitions.contains_key(&table.name) {
continue;
}
let rule = inputs.config.tables.get(&table.name);
let anonymized = inputs.anonymize_tables.contains(&table.name);
let encrypted = inputs.encrypted.get(&table.name);
let mut columns = Vec::new();
for column in &table.columns {
let qualified = format!("{}.{}", table.name, column.name);
if inputs
.facts
.generated_columns
.contains(&(table.name.clone(), column.name.clone()))
{
continue;
}
let is_key = is_key_column(table, column, inputs.facts);
let declared_pii = rule.and_then(|r| r.pii.get(&column.name)).copied();
let is_encrypted = encrypted.is_some_and(|e| e.contains_key(&column.name));
let declared_safe_here = rule.is_some_and(|r| r.safe.contains(&column.name));
let declared_safe = declared_safe_here
|| (!anonymized && inputs.config.defaults.safe_columns.contains(&column.name));
let (spec, source) = if is_encrypted {
match declared_pii {
None | Some(Strategy::Encrypted) => {
(Strategy::Encrypted, ClassSource::Encrypted)
}
Some(Strategy::Null) => (Strategy::Null, ClassSource::Config),
Some(other) => {
plaintext_into_encrypted.push((qualified, other.as_str()));
continue;
}
}
} else if let Some(strategy) = declared_pii {
(strategy, ClassSource::Config)
} else if declared_safe {
continue;
} else if anonymized {
if is_key {
continue;
}
(Strategy::Auto, ClassSource::GdprAnonymize)
} else {
unclassified.push(qualified);
continue;
};
if is_key {
key_pii.push(qualified);
continue;
}
columns.push((
ColumnPlan {
column: column.name.clone(),
strategy: spec,
source,
},
column,
));
}
planned.push((table, columns));
}
if !plaintext_into_encrypted.is_empty() {
plaintext_into_encrypted.sort();
return Err(ScrubError::PlaintextIntoEncrypted {
columns: plaintext_into_encrypted
.into_iter()
.map(|(column, strategy)| format!("{column} (declared `{strategy}`)"))
.collect(),
});
}
if !key_pii.is_empty() {
key_pii.sort();
return Err(ScrubError::PiiOnKeyColumn { columns: key_pii });
}
if !unclassified.is_empty() {
unclassified.sort();
return Err(ScrubError::Unclassified {
columns: unclassified,
});
}
let mut out = ScrubPlan::default();
for (table, columns) in planned {
if columns.is_empty() {
continue;
}
let mut resolved = Vec::with_capacity(columns.len());
let mut assignments = Vec::with_capacity(columns.len());
let mut encrypted_rewrites = Vec::new();
for (mut plan, column) in columns {
let qualified = format!("{}.{}", table.name, column.name);
let pair = (table.name.clone(), column.name.clone());
let unique = is_unique_column(table, column, inputs.facts);
if inputs.facts.checked_columns.contains(&pair) {
return Err(ScrubError::CheckConstrainedColumn { column: qualified });
}
if plan.strategy == Strategy::Auto {
plan.strategy = auto_strategy(column).map_err(|e| qualify(e, &qualified))?;
}
if plan.strategy == Strategy::Null {
if !column.nullable {
return Err(ScrubError::NullOnNotNull { column: qualified });
}
if inputs.facts.nulls_not_distinct_columns.contains(&pair) {
return Err(ScrubError::NonUniqueStrategy {
column: qualified,
strategy: plan.strategy.as_str(),
});
}
}
if unique && !plan.strategy.allowed_on_unique() {
return Err(ScrubError::NonUniqueStrategy {
column: qualified,
strategy: plan.strategy.as_str(),
});
}
if plan.strategy == Strategy::Encrypted {
encrypted_rewrites.push(EncryptedRewrite {
column: column.name.clone(),
deterministic: inputs
.encrypted
.get(&table.name)
.and_then(|c| c.get(&column.name))
.copied()
.unwrap_or(false),
shape: email_shaped(&column.name),
});
resolved.push(plan);
continue;
}
let token = token_expr(table, &column.name, unique);
let value = replacement_expr(plan.strategy, column, &token, unique)
.map_err(|e| qualify(e, &qualified))?;
assignments.push(assignment(column, &value, plan.strategy));
resolved.push(plan);
}
out.tables.push(TablePlan {
table: table.name.clone(),
columns: resolved,
sql: (!assignments.is_empty()).then(|| {
format!(
"UPDATE {} SET {}",
qualified_ident(&table.name),
assignments.join(", ")
)
}),
row_key: row_key_expr(table),
encrypted: encrypted_rewrites,
});
}
Ok(out)
}
fn qualify(error: ScrubError, qualified: &str) -> ScrubError {
let column = qualified.to_owned();
match error {
ScrubError::NoAutoStrategy { detail, .. } => ScrubError::NoAutoStrategy { column, detail },
ScrubError::StrategyTypeMismatch {
strategy, detail, ..
} => ScrubError::StrategyTypeMismatch {
column,
strategy,
detail,
},
ScrubError::ColumnTooNarrow {
limit,
overhead,
floor,
..
} => ScrubError::ColumnTooNarrow {
column,
limit,
overhead,
floor,
},
other => other,
}
}
fn check_config_freshness(
inputs: &ClassificationInputs<'_>,
by_name: &BTreeMap<&str, &Table>,
) -> Result<(), ScrubError> {
let mut stale = Vec::new();
let mut framework = Vec::new();
for (name, rule) in &inputs.config.tables {
let Some(table) = by_name.get(name.as_str()) else {
if is_framework_table(name) {
framework.push(name.clone());
} else {
stale.push(name.clone());
}
continue;
};
let columns: BTreeSet<&str> = table.columns.iter().map(|c| c.name.as_str()).collect();
for column in rule
.safe
.iter()
.chain(rule.pii.keys())
.chain(rule.encrypted.keys())
{
if !columns.contains(column.as_str()) {
stale.push(format!("{name}.{column}"));
}
}
}
if !framework.is_empty() {
framework.sort();
return Err(ScrubError::FrameworkTableDeclared { tables: framework });
}
if stale.is_empty() {
return Ok(());
}
stale.sort();
stale.dedup();
Err(ScrubError::StaleConfig { entries: stale })
}
fn check_contradictions(inputs: &ClassificationInputs<'_>) -> Result<(), ScrubError> {
let mut columns = Vec::new();
for (name, rule) in &inputs.config.tables {
for column in &rule.safe {
if rule.pii.contains_key(column) {
columns.push(format!("{name}.{column}"));
}
}
}
if columns.is_empty() {
return Ok(());
}
columns.sort();
Err(ScrubError::Contradiction { columns })
}
fn check_safe_overrides_encrypted(
inputs: &ClassificationInputs<'_>,
by_name: &BTreeMap<&str, &Table>,
) -> Result<(), ScrubError> {
let mut columns = Vec::new();
for (table, encrypted) in inputs.encrypted {
if !by_name.contains_key(table.as_str()) {
continue;
}
let rule = inputs.config.tables.get(table);
for column in encrypted.keys() {
if rule.is_some_and(|r| r.pii.contains_key(column)) {
continue;
}
if by_name
.get(table.as_str())
.and_then(|t| {
t.columns
.iter()
.find(|c| &c.name == column)
.map(|c| (*t, c))
})
.is_some_and(|(t, c)| is_key_column(t, c, inputs.facts))
{
continue;
}
if rule.is_some_and(|r| r.safe.contains(column))
|| inputs.config.defaults.safe_columns.contains(column)
{
columns.push(format!("{table}.{column}"));
}
}
}
if columns.is_empty() {
return Ok(());
}
columns.sort();
Err(ScrubError::SafeOverridesEncrypted { columns })
}
fn unreachable_tables(facts: &DatabaseFacts, introspected: &[Table]) -> Vec<String> {
let seen: BTreeSet<&str> = introspected.iter().map(|t| t.name.as_str()).collect();
let mut missing: Vec<String> = facts
.public_base_tables
.iter()
.filter(|t| !seen.contains(t.as_str()) && !is_framework_table(t))
.cloned()
.collect();
let seen_columns: BTreeSet<(&str, &str)> = introspected
.iter()
.flat_map(|t| {
t.columns
.iter()
.map(move |c| (t.name.as_str(), c.name.as_str()))
})
.collect();
missing.extend(
facts
.public_columns
.iter()
.filter(|(table, column)| {
seen.contains(table.as_str())
&& !is_framework_table(table)
&& !seen_columns.contains(&(table.as_str(), column.as_str()))
})
.map(|(table, column)| format!("{table}.{column}")),
);
missing.sort();
missing.dedup();
missing
}
fn is_key_column(table: &Table, column: &Column, facts: &DatabaseFacts) -> bool {
let pair = (table.name.clone(), column.name.clone());
column.primary_key
|| table.primary_key.contains(&column.name)
|| column.references.is_some()
|| facts.foreign_key_columns.contains(&pair)
|| facts.generated_columns.contains(&pair)
}
fn is_unique_column(table: &Table, column: &Column, facts: &DatabaseFacts) -> bool {
column.unique
|| facts
.unique_columns
.contains(&(table.name.clone(), column.name.clone()))
|| table.indexes.iter().any(|index| {
if !index.unique {
return false;
}
let keys = if index.key_columns.is_empty() {
&index.columns
} else {
&index.key_columns
};
keys.iter().any(|k| k == &column.name)
})
}
fn row_key_expr(table: &Table) -> String {
let mut keys: Vec<&str> = table.primary_key.iter().map(String::as_str).collect();
if keys.is_empty() {
keys = table
.columns
.iter()
.filter(|c| c.primary_key)
.map(|c| c.name.as_str())
.collect();
}
if keys.is_empty() {
return "ctid::text".to_owned();
}
if keys.len() == 1 {
return format!("coalesce({}::text, '')", quote_ident(keys[0]));
}
format!(
"ROW({})::text",
keys.iter()
.map(|key| quote_ident(key))
.collect::<Vec<_>>()
.join(", ")
)
}
fn token_expr(table: &Table, column: &str, unique: bool) -> String {
token_expr_from(&row_key_expr(table), column, unique)
}
fn token_expr_from(row_key: &str, column: &str, unique: bool) -> String {
let seed = format!("{} || '|' || {row_key}", quote_literal(column));
if unique {
format!("md5({seed}) || md5(({seed}) || '#2')")
} else {
format!("md5({seed})")
}
}
fn char_max_len(ty: &ColumnType) -> Option<usize> {
let ColumnType::Opaque { pg_type } = ty else {
return None;
};
let rest = pg_type
.strip_prefix("varchar(")
.or_else(|| pg_type.strip_prefix("char("))?;
rest.strip_suffix(')')?.parse().ok()
}
fn is_texty(ty: &ColumnType) -> bool {
match ty {
ColumnType::Text => true,
ColumnType::Opaque { pg_type } => {
char_max_len(ty).is_some() || matches!(pg_type.as_str(), "citext" | "name")
}
_ => false,
}
}
fn auto_strategy(column: &Column) -> Result<Strategy, ScrubError> {
Ok(match &column.ty {
ColumnType::Text => email_shaped(&column.name),
ColumnType::Uuid => Strategy::Uuid,
ColumnType::Bytes => Strategy::Bytes,
ColumnType::Json | ColumnType::Attachment => Strategy::Json,
ColumnType::Int32
| ColumnType::Int64
| ColumnType::Float32
| ColumnType::Float64
| ColumnType::Bool
| ColumnType::Decimal { .. } => Strategy::Zero,
ColumnType::Timestamp | ColumnType::TimestampTz => Strategy::Epoch,
ty @ ColumnType::Opaque { .. } if is_texty(ty) => email_shaped(&column.name),
ColumnType::Opaque { pg_type }
if matches!(pg_type.as_str(), "date" | "time" | "timetz") =>
{
Strategy::Epoch
}
ColumnType::Opaque { pg_type }
if matches!(pg_type.as_str(), "int2" | "money" | "oid")
|| pg_type.starts_with("numeric") =>
{
Strategy::Zero
}
other => {
return Err(ScrubError::NoAutoStrategy {
column: column.name.clone(),
detail: format!("{other:?}"),
});
}
})
}
fn email_shaped(name: &str) -> Strategy {
if name.to_ascii_lowercase().contains("email") {
Strategy::Email
} else {
Strategy::Redact
}
}
#[allow(clippy::too_many_lines)]
fn replacement_expr(
strategy: Strategy,
column: &Column,
token: &str,
unique: bool,
) -> Result<String, ScrubError> {
let limit = char_max_len(&column.ty);
let narrow = |overhead: usize| -> Result<String, ScrubError> {
bounded_token(token, limit, overhead, &column.name, unique)
};
if matches!(
strategy,
Strategy::Email | Strategy::Name | Strategy::Redact | Strategy::Phone
) {
require_type(column, strategy, is_texty(&column.ty))?;
}
Ok(match strategy {
Strategy::Auto => replacement_expr(auto_strategy(column)?, column, token, unique)?,
Strategy::Encrypted => {
return Err(ScrubError::StrategyTypeMismatch {
column: column.name.clone(),
strategy: strategy.as_str(),
detail: "an encrypted replacement is built in Rust, not in SQL".to_owned(),
});
}
Strategy::Email => {
let tok = narrow("scrubbed+".len() + SCRUB_EMAIL_DOMAIN.len())?;
format!("'scrubbed+' || {tok} || '{SCRUB_EMAIL_DOMAIN}'")
}
Strategy::Name => {
let tok = narrow("Scrubbed ".len())?;
format!("'Scrubbed ' || {tok}")
}
Strategy::Redact => {
let tok = narrow("[scrubbed:]".len())?;
format!("'[scrubbed:' || {tok} || ']'")
}
Strategy::Phone => {
const PHONE_DIGITS: usize = 10;
let overhead = "+1555".len();
if let Some(limit) = limit
&& limit < overhead + PHONE_DIGITS
{
return Err(ScrubError::ColumnTooNarrow {
column: column.name.clone(),
limit,
overhead,
floor: PHONE_DIGITS,
});
}
format!("'+1555' || translate(substr({token}, 1, {PHONE_DIGITS}), 'abcdef', '0123456')")
}
Strategy::Null => "NULL".to_owned(),
Strategy::Uuid => {
require_type(column, strategy, matches!(column.ty, ColumnType::Uuid))?;
format!("(substr({token}, 1, 32))::uuid")
}
Strategy::Bytes => {
require_type(column, strategy, matches!(column.ty, ColumnType::Bytes))?;
format!("decode({token}, 'hex')")
}
Strategy::Json => {
let json = "'{\"scrubbed\": true}'";
match &column.ty {
ColumnType::Json | ColumnType::Attachment => format!("{json}::jsonb"),
ColumnType::Opaque { pg_type } if pg_type == "json" => format!("{json}::json"),
ty if is_texty(ty) => {
const JSON_LITERAL_LEN: usize = 18;
if let Some(limit) = limit
&& limit < JSON_LITERAL_LEN
{
return Err(ScrubError::ColumnTooNarrow {
column: column.name.clone(),
limit,
overhead: JSON_LITERAL_LEN,
floor: 0,
});
}
json.to_owned()
}
other => {
return Err(ScrubError::StrategyTypeMismatch {
column: column.name.clone(),
strategy: strategy.as_str(),
detail: format!("{other:?}"),
});
}
}
}
Strategy::Zero => match &column.ty {
ColumnType::Bool => "false".to_owned(),
ColumnType::Int32
| ColumnType::Int64
| ColumnType::Float32
| ColumnType::Float64
| ColumnType::Decimal { .. } => "0".to_owned(),
ColumnType::Opaque { pg_type }
if pg_type.starts_with("numeric")
|| matches!(pg_type.as_str(), "int2" | "money" | "oid") =>
{
"0".to_owned()
}
other => {
return Err(ScrubError::StrategyTypeMismatch {
column: column.name.clone(),
strategy: strategy.as_str(),
detail: format!("{other:?}"),
});
}
},
Strategy::Epoch => match &column.ty {
ColumnType::Timestamp => "'1970-01-01 00:00:00'::timestamp".to_owned(),
ColumnType::TimestampTz => "'1970-01-01 00:00:00+00'::timestamptz".to_owned(),
ColumnType::Opaque { pg_type } if pg_type == "date" => "'1970-01-01'::date".to_owned(),
ColumnType::Opaque { pg_type } if pg_type == "time" => "'00:00:00'::time".to_owned(),
ColumnType::Opaque { pg_type } if pg_type == "timetz" => {
"'00:00:00+00'::timetz".to_owned()
}
other => {
return Err(ScrubError::StrategyTypeMismatch {
column: column.name.clone(),
strategy: strategy.as_str(),
detail: format!("{other:?}"),
});
}
},
})
}
fn require_type(column: &Column, strategy: Strategy, ok: bool) -> Result<(), ScrubError> {
if ok {
return Ok(());
}
Err(ScrubError::StrategyTypeMismatch {
column: column.name.clone(),
strategy: strategy.as_str(),
detail: format!("{:?}", column.ty),
})
}
fn bounded_token(
token: &str,
limit: Option<usize>,
overhead: usize,
column: &str,
unique: bool,
) -> Result<String, ScrubError> {
let (full, floor) = if unique {
(UNIQUE_TOKEN_HEX_LEN, MIN_UNIQUE_TOKEN_WIDTH)
} else {
(TOKEN_HEX_LEN, MIN_TOKEN_WIDTH)
};
let Some(limit) = limit else {
return Ok(token.to_owned());
};
let available = limit.saturating_sub(overhead);
if available >= full {
Ok(token.to_owned())
} else if available >= floor {
Ok(format!("substr({token}, 1, {available})"))
} else {
Err(ScrubError::ColumnTooNarrow {
column: column.to_owned(),
limit,
overhead,
floor,
})
}
}
fn assignment(column: &Column, value: &str, strategy: Strategy) -> String {
let ident = quote_ident(&column.name);
if column.nullable && strategy != Strategy::Null {
format!("{ident} = CASE WHEN {ident} IS NULL THEN NULL ELSE {value} END")
} else {
format!("{ident} = {value}")
}
}
fn qualified_ident(table: &str) -> String {
format!("\"public\".{}", quote_ident(table))
}
fn extract_anonymize_tables(src: &str) -> Result<BTreeSet<String>, ScrubError> {
use syn::visit::Visit as _;
let file = syn::parse_file(src).map_err(|e| ScrubError::SourceScan {
detail: e.to_string(),
})?;
let mut scan = AnonymizeScan::default();
scan.visit_file(&file);
if let Some(detail) = scan.unresolved.into_iter().next() {
return Err(ScrubError::UnresolvableAnonymize { detail });
}
Ok(scan.tables)
}
#[derive(Default)]
struct AnonymizeScan {
tables: BTreeSet<String>,
unresolved: Vec<String>,
}
impl<'ast> syn::visit::Visit<'ast> for AnonymizeScan {
fn visit_expr_call(&mut self, node: &'ast syn::ExprCall) {
if let syn::Expr::Path(path) = &*node.func {
let segments = &path.path.segments;
let is_anonymize = segments.len() >= 2
&& segments[segments.len() - 1].ident == "anonymize"
&& segments[segments.len() - 2].ident == "ModelRegistration";
if is_anonymize {
match node.args.first() {
Some(syn::Expr::Lit(syn::ExprLit {
lit: syn::Lit::Str(table),
..
})) => {
self.tables.insert(table.value());
}
_ => self.unresolved.push(quote::quote!(#node).to_string()),
}
}
}
syn::visit::visit_expr_call(self, node);
}
}
fn scan_anonymize_tables(root: &Path) -> Result<BTreeSet<String>, ScrubError> {
let mut out = BTreeSet::new();
if !root.is_dir() {
return Ok(out);
}
let mut stack = vec![root.to_path_buf()];
while let Some(dir) = stack.pop() {
let entries = std::fs::read_dir(&dir).map_err(|e| ScrubError::SourceScan {
detail: format!("{}: {e}", dir.display()),
})?;
for entry in entries {
let entry = entry.map_err(|e| ScrubError::SourceScan {
detail: format!("{}: {e}", dir.display()),
})?;
let path = entry.path();
let file_type = entry.file_type().map_err(|e| ScrubError::SourceScan {
detail: format!("{}: {e}", path.display()),
})?;
if file_type.is_dir() {
stack.push(path);
} else if file_type.is_file() && path.extension().is_some_and(|ext| ext == "rs") {
let src = std::fs::read_to_string(&path).map_err(|e| ScrubError::SourceScan {
detail: format!("{}: {e}", path.display()),
})?;
let found = extract_anonymize_tables(&src).map_err(|e| match e {
ScrubError::SourceScan { detail } => ScrubError::SourceScan {
detail: format!("{}: {detail}", path.display()),
},
ScrubError::UnresolvableAnonymize { detail } => {
ScrubError::UnresolvableAnonymize {
detail: format!("{} — {detail}", path.display()),
}
}
other => other,
})?;
out.extend(found);
}
}
}
Ok(out)
}
fn encrypted_columns(
project_root: &Path,
) -> Result<BTreeMap<String, BTreeMap<String, bool>>, ScrubError> {
let Some(path) = crate::schema::existing_models_path(project_root) else {
return Ok(BTreeMap::new());
};
crate::schema::parse::parse_encrypted_columns_path(&path).map_err(|e| ScrubError::SourceScan {
detail: e.to_string(),
})
}
fn guard_scrub_target(profile: &str, force: bool) -> Result<(), ScrubError> {
super::guard_destructive(profile, force).map_err(|_| ScrubError::ProductionRefused {
profile: profile.to_owned(),
})
}
fn same_database(a: &str, b: &str) -> bool {
let parts = |raw: &str| -> Option<(String, u16, String)> {
let parsed = url::Url::parse(raw).ok()?;
let host = parsed.host_str()?.to_ascii_lowercase();
let port = parsed.port().unwrap_or(5432);
let name = parsed
.path_segments()
.and_then(|mut s| s.next())
.map(str::to_owned)
.filter(|n| !n.is_empty())?;
Some((host, port, name))
};
match (parts(a), parts(b)) {
(Some(left), Some(right)) => left == right,
_ => false,
}
}
fn guard_configured_source(
artifact_profile: Option<&str>,
project_root: &Path,
targets: &[(String, String)],
allowed: bool,
) -> Result<(), ScrubError> {
if allowed {
return Ok(());
}
let mut candidates: Vec<String> = artifact_profile
.map(str::to_owned)
.into_iter()
.chain(profiles_with_config(project_root))
.filter(|p| !super::is_safe_destructive_profile(p))
.collect();
candidates.sort();
candidates.dedup();
for profile in candidates {
let table = migrate::read_autumn_toml_table_with_profile(Some(&profile));
let Some(declared) = migrate::resolve_primary_database_url_from_sources(
|_| Err(std::env::VarError::NotPresent),
table.as_ref(),
) else {
continue;
};
for (_, url) in targets {
if same_database(&declared, url) {
return Err(ScrubError::OverwritesConfiguredTarget {
profile,
database: parsed_db_name(url),
});
}
}
}
Ok(())
}
fn profiles_with_config(project_root: &Path) -> Vec<String> {
let Ok(entries) = std::fs::read_dir(project_root) else {
return Vec::new();
};
let mut out: Vec<String> = entries
.filter_map(Result::ok)
.filter_map(|entry| {
let name = entry.file_name().to_string_lossy().into_owned();
name.strip_prefix("autumn-")
.and_then(|rest| rest.strip_suffix(".toml"))
.map(str::to_owned)
})
.filter(|p| !p.is_empty())
.collect();
out.sort();
out.dedup();
out
}
fn parsed_db_name(url: &str) -> String {
url::Url::parse(url)
.ok()
.and_then(|u| {
u.path_segments()
.and_then(|mut s| s.next())
.map(str::to_owned)
})
.filter(|n| !n.is_empty())
.unwrap_or_else(|| "<unknown>".to_owned())
}
fn suggested_config_stanza(unclassified: &[String]) -> String {
use std::fmt::Write as _;
let mut by_table: BTreeMap<&str, Vec<&str>> = BTreeMap::new();
for entry in unclassified {
if let Some((table, column)) = entry.split_once('.') {
by_table.entry(table).or_default().push(column);
}
}
let mut out = String::from(
"# Paste into scrub.toml, then replace `auto` with an explicit strategy\n\
# (email/name/phone/redact/null/uuid/bytes/json/zero/epoch), or move the\n\
# column into that table's `safe = [...]` list if it holds no PII.\n",
);
for (table, columns) in by_table {
let _ = write!(out, "\n[tables.{table}.pii]\n");
for column in columns {
let _ = writeln!(out, "{column} = \"auto\"");
}
}
out
}
pub fn run(args: &ScrubArgs) {
eprintln!("\u{1F342} autumn db scrub\n");
if let Err(e) = scrub(args) {
eprintln!("\u{2717} {e}");
if let ScrubError::Unclassified { columns } = &e {
println!("{}", suggested_config_stanza(columns));
}
std::process::exit(1);
}
}
fn scrub(args: &ScrubArgs) -> Result<(), ScrubError> {
let profile = migrate::effective_profile(args.profile.as_deref());
let writes = !args.check && !args.dry_run;
if writes {
guard_scrub_target(&profile, args.force)?;
}
let sources = load_source_classification(args)?;
let targets = super::backup::resolve_all_target_urls(args.profile.as_deref())?;
let restored = if let Some(artifact) = &args.artifact {
let artifact_profile = super::backup::artifact_source_profile(artifact);
debug_assert!(
writes,
"--check/--dry-run must never reach an artifact restore"
);
eprintln!(
" \u{2139} Artifact provenance: {}.",
artifact_profile.as_deref().map_or_else(
|| "unknown (no manifest)".to_owned(),
|p| format!("{p:?} profile")
)
);
guard_configured_source(
artifact_profile.as_deref(),
Path::new("."),
&targets,
args.allow_source_overwrite,
)?;
eprintln!(
"\u{2500}\u{2500} restoring {} \u{2500}\u{2500}",
artifact.display()
);
super::backup::restore(&super::backup::RestoreArgs {
artifact: artifact.clone(),
profile: args.profile.clone(),
force: args.force,
shard: None,
offsite: false,
})?;
true
} else {
false
};
classify_and_apply(args, &profile, &targets, &sources).inspect_err(|_| {
if restored {
eprintln!(
"\n\u{26A0}\u{FE0F} The artifact was ALREADY RESTORED before this failure, so \
the target database now holds UNSCRUBBED data.\n \
Do not hand it to anyone: fix the problem below and re-run the same \
command, or drop the database."
);
}
})
}
struct SourceClassification {
config: ScrubConfig,
encrypted: BTreeMap<String, BTreeMap<String, bool>>,
anonymize: BTreeSet<String>,
roots: Vec<sample::SampleSpec>,
}
fn load_source_classification(args: &ScrubArgs) -> Result<SourceClassification, ScrubError> {
let config = load_config(args.config.as_deref())?;
check_purge_list(&config)?;
let project_root = Path::new(".");
let models_path = crate::schema::existing_models_path(project_root);
let declared_encrypted = config
.tables
.iter()
.any(|(_, rule)| !rule.encrypted.is_empty());
if models_path.is_none() && !declared_encrypted {
return Err(ScrubError::EncryptedMetadataUnavailable);
}
let mut encrypted = encrypted_columns(project_root)?;
for (table, rule) in &config.tables {
for (column, mode) in &rule.encrypted {
encrypted
.entry(table.clone())
.or_default()
.entry(column.clone())
.or_insert_with(|| mode.is_deterministic());
}
}
let anonymize = scan_anonymize_tables(&project_root.join("src"))?;
let encrypted_count: usize = encrypted.values().map(BTreeMap::len).sum();
eprintln!(
" \u{2139} Automatic classification: {encrypted_count} #[encrypted] column(s), \
{} GDPR anonymize registration(s).",
anonymize.len()
);
let roots = args
.sample
.iter()
.map(|spec| sample::parse_spec(spec))
.collect::<Result<Vec<_>, _>>()?;
if roots.is_empty() && sources_declare_sampling(&config) {
eprintln!(
" \u{2139} scrub.toml declares [sample] rules, but no --sample root was \
given \u{2014} the whole copy is kept."
);
}
Ok(SourceClassification {
config,
encrypted,
anonymize,
roots,
})
}
const fn sources_declare_sampling(config: &ScrubConfig) -> bool {
!config.sample.always_include.is_empty() || !config.sample.never_include.is_empty()
}
#[allow(clippy::too_many_lines)]
fn classify_and_apply(
args: &ScrubArgs,
profile: &str,
targets: &[(String, String)],
sources: &SourceClassification,
) -> Result<(), ScrubError> {
let mut plans = Vec::with_capacity(targets.len());
for (label, url) in targets {
let facts = probe_database_facts(url, label, &sources.config)?;
if !facts.other_schemas.is_empty() {
return Err(ScrubError::UnsupportedSchemas {
schemas: facts.other_schemas.iter().cloned().collect(),
});
}
let tables = introspect::introspect_postgres(url).map_err(|e| ScrubError::Introspect {
label: label.clone(),
detail: e.to_string(),
})?;
let unreachable = unreachable_tables(&facts, &tables);
if !unreachable.is_empty() {
return Err(ScrubError::InaccessibleTables {
tables: unreachable,
});
}
if !args.sample.is_empty() && !facts.legacy_inheritance.is_empty() {
return Err(ScrubError::LegacyInheritance {
tables: facts.legacy_inheritance,
});
}
if !facts.views_reading_by_name.is_empty() {
return Err(ScrubError::ViewReadsRelationsByName {
views: facts.views_reading_by_name,
});
}
if !facts.untraceable_view_functions.is_empty() {
return Err(ScrubError::UntraceableViewFunction {
views: facts.untraceable_view_functions,
});
}
let mut unreachable_views: Vec<String> = facts
.all_materialized_views
.iter()
.filter(|view| !facts.materialized_views.contains(view))
.cloned()
.collect();
if !unreachable_views.is_empty() {
unreachable_views.sort();
return Err(ScrubError::UnrefreshableViews {
views: unreachable_views,
});
}
let plan = build_plan(&ClassificationInputs {
tables: &tables,
config: &sources.config,
encrypted: &sources.encrypted,
anonymize_tables: &sources.anonymize,
facts: &facts,
})?;
let sampling = if sources.roots.is_empty() {
None
} else {
Some(build_sample_plan_for(args, sources, &tables, &facts)?)
};
let mut rls: Vec<String> = plan
.tables
.iter()
.map(|t| t.table.clone())
.chain(
purge_statements(&facts.framework_tables, &sources.config)
.into_iter()
.map(|(table, _)| table),
)
.chain(
sampling
.iter()
.flat_map(sample::SamplePlan::locked_tables)
.map(str::to_owned),
)
.filter(|t| facts.rls_tables.contains(t))
.collect();
rls.sort();
rls.dedup();
if !rls.is_empty() {
return Err(ScrubError::RowLevelSecurity { tables: rls });
}
if facts.replication_role != "origin" {
return Err(ScrubError::ReplicaSessionRole {
role: facts.replication_role,
});
}
let mut emptying_triggers: Vec<String> =
purge_statements(&facts.framework_tables, &sources.config)
.into_iter()
.map(|(table, _)| table)
.chain(
sampling
.iter()
.flat_map(|s| s.emptied_tables())
.map(|(table, _)| table.to_owned()),
)
.filter(|t| facts.delete_triggered_tables.contains(t))
.collect();
emptying_triggers.sort();
emptying_triggers.dedup();
if !emptying_triggers.is_empty() {
return Err(ScrubError::EmptyingTriggerLeak {
tables: emptying_triggers,
});
}
report_plan(label, &plan);
report_framework_tables(&facts.framework_tables, &sources.config);
report_triggers(&plan, sampling.as_ref(), &facts);
if let Some(sampling) = &sampling {
report_sample_plan(label, sampling);
}
plans.push((label, url, plan, facts, sampling));
}
if args.check {
eprintln!(
"\n\u{2713} Every column in `public` is classified \u{2014} no unclassified data can leak."
);
if !sources.roots.is_empty() {
eprintln!(
"\u{2713} Every table is covered by the sample \u{2014} no table would be \
emptied unannounced."
);
}
return Ok(());
}
if args.dry_run {
let unprintable: Vec<String> = plans
.iter()
.filter(|(_, url, _, _, _)| password_free_conninfo(url).is_none())
.map(|(label, _, _, _, _)| (*label).clone())
.collect();
if !unprintable.is_empty() {
return Err(ScrubError::UnprintableTarget {
targets: unprintable,
});
}
guard_scrub_target(profile, args.force).map_err(|_| {
ScrubError::UnprintableProductionTarget {
profile: profile.to_owned(),
}
})?;
let mut ambiguous: Vec<String> = plans
.iter()
.filter(|(_, _, _, facts, _)| {
let e = &facts.endpoint;
e.address.is_none() && e.port.is_none()
})
.map(|(label, _, _, _, _)| (*label).clone())
.collect();
ambiguous.sort();
if !ambiguous.is_empty() {
return Err(ScrubError::UnprintableAmbiguousTarget { targets: ambiguous });
}
let mut portless: Vec<String> = plans
.iter()
.filter(|(_, url, _, _, _)| {
password_free_conninfo(url)
.is_some_and(|conninfo| stated_host_and_port(&conninfo).1.is_none())
})
.map(|(label, _, _, _, _)| (*label).clone())
.collect();
portless.sort();
if !portless.is_empty() {
return Err(ScrubError::UnprintablePortlessTarget { targets: portless });
}
let mut addressed: Vec<String> = plans
.iter()
.filter(|(_, url, _, _, _)| {
password_free_conninfo(url).is_some_and(|c| endpoint_can_come_from_elsewhere(&c))
})
.map(|(label, _, _, _, _)| (*label).clone())
.collect();
addressed.sort();
if !addressed.is_empty() {
return Err(ScrubError::UnprintableHostaddrTarget { targets: addressed });
}
let mut multi: Vec<String> = plans
.iter()
.filter(|(_, url, _, _, _)| {
password_free_conninfo(url).is_some_and(|c| states_multiple_endpoints(&c))
})
.map(|(label, _, _, _, _)| (*label).clone())
.collect();
multi.sort();
if !multi.is_empty() {
return Err(ScrubError::UnprintableMultiHostTarget { targets: multi });
}
let mut encrypted_rewrites: Vec<String> = plans
.iter()
.flat_map(|(label, _, plan, _, _)| {
plan.tables.iter().flat_map(move |table| {
table
.encrypted
.iter()
.map(move |rewrite| format!("{label}: {}.{}", table.table, rewrite.column))
})
})
.collect();
encrypted_rewrites.sort();
if !encrypted_rewrites.is_empty() {
return Err(ScrubError::UnprintableEncryptedRewrite {
columns: encrypted_rewrites,
});
}
let mut sized: Vec<Option<BTreeMap<String, i64>>> = Vec::with_capacity(plans.len());
for (label, url, _, _, sampling) in &plans {
sized.push(match sampling {
Some(plan) => {
let mut conn = probe_connection(url, label, "size the sample")?;
Some(
sample::source_counts(&mut conn, plan)
.map_err(|e| ScrubError::Sql(e.to_string()))?,
)
}
None => None,
});
}
eprintln!(" \\set ON_ERROR_STOP on");
eprintln!(" \\set autumn_ok true");
for (index, (label, url, plan, facts, sampling)) in plans.iter().enumerate() {
let no_deferral = BTreeSet::new();
let deferred: &BTreeSet<String> =
sampling.as_ref().map_or(&no_deferral, |s| &s.purge_after);
let purges = purge_statements(&facts.framework_tables, &sources.config);
let phases = emptying_phases(&purges, deferred, sampling.as_ref());
eprintln!(" \\if :autumn_ok");
for line in psql_connect(label, url) {
eprintln!(" {line}");
}
if let Some(conninfo) = password_free_conninfo(url) {
for line in psql_connection_assertion(&conninfo, &facts.endpoint.database) {
eprintln!(" {line}");
}
}
eprintln!(" \\set autumn_scrubbed false");
eprintln!(" BEGIN;");
for statement in session_settings() {
eprintln!(" {statement};");
}
eprintln!(" {}", replication_role_assertion());
eprintln!(" {}", target_guard(&facts.endpoint));
for statement in lock_statements(plan, &purges, sampling.as_ref()) {
eprintln!(" {statement};");
}
for (_, statement) in &phases.before {
eprintln!(" {statement};");
}
if let Some(sampling) = sampling {
report_sample_sql(sampling, sized[index].as_ref());
}
for (_, statement) in &phases.after_sample {
eprintln!(" {statement};");
}
for table in &plan.tables {
if let Some(sql) = &table.sql {
eprintln!(" {sql};");
}
}
for (_, statement) in &phases.final_pass {
eprintln!(" {statement};");
}
for view in &facts.materialized_views {
eprintln!(" REFRESH MATERIALIZED VIEW {};", qualified_ident(view));
}
for view in &facts.unpopulated_views {
eprintln!(
" REFRESH MATERIALIZED VIEW {} WITH NO DATA;",
qualified_ident(view)
);
}
for (table, _) in &phases.final_pass {
eprintln!(" {};", emptiness_assertion(table));
}
if let Some(sampling) = sampling {
for (statement, comment) in sample_postcondition_statements(sampling) {
eprintln!(" {statement}; -- {comment}");
}
}
eprintln!(" SELECT true AS autumn_scrubbed \\gset");
eprintln!(" COMMIT;");
for line in post_commit_fence(&facts.endpoint) {
eprintln!(" {line}");
}
eprintln!(" \\endif");
eprintln!(" \\endif");
}
let sampled: Vec<_> = plans
.iter()
.filter_map(|(label, url, _, facts, sampling)| {
sampling.as_ref().map(|s| (label, url, facts, s))
})
.collect();
if !sampled.is_empty() {
eprintln!(" \\if :autumn_ok");
eprintln!(" \\set ON_ERROR_STOP off");
for (label, url, facts, sampling) in sampled {
let purges = purge_statements(&facts.framework_tables, &sources.config);
for line in psql_connect(label, url) {
eprintln!(" {line}");
}
eprintln!(" \\set autumn_here false");
eprintln!(
" SELECT ({} AND NOT ({})) AS autumn_here \\gset",
password_free_conninfo(url).map_or_else(
|| "false".to_owned(),
|conninfo| psql_connection_terms(&conninfo, &facts.endpoint.database),
),
endpoint_mismatch(&facts.endpoint)
);
eprintln!(" \\if :autumn_here");
eprintln!(" SET lock_timeout = '{COMPACT_LOCK_TIMEOUT}';");
for table in compacted_tables(sampling, &purged_tables(&purges)) {
eprintln!(" VACUUM (FULL, ANALYZE) {};", qualified_ident(table));
}
eprintln!(" \\endif");
}
eprintln!(" \\endif");
}
eprintln!("\n\u{2713} Dry run only \u{2014} nothing was written.");
return Ok(());
}
if plans
.iter()
.any(|(_, _, plan, _, _)| plan.tables.iter().any(|t| !t.encrypted.is_empty()))
{
let ring = resolve_key_ring(profile, Path::new("."))?;
autumn_web::encryption::install_key_ring(ring);
}
let mut committed: Vec<&str> = Vec::new();
let warn_committed = |committed: &[&str]| {
if !committed.is_empty() {
eprintln!(
"\n\u{26A0}\u{FE0F} Already committed before this failure: {}. \
Those databases ARE scrubbed; every later target is untouched and still \
holds real data.",
committed.join(", ")
);
}
};
let mut pending_compaction: Vec<PendingCompaction<'_>> = Vec::new();
for (label, url, plan, facts, sampling) in &plans {
let purges = purge_statements(&facts.framework_tables, &sources.config);
let (applied, sampled) = execute(
url,
plan,
&purges,
&ViewRefresh {
ordered: &facts.materialized_views,
all: &facts.all_materialized_views,
unpopulated: &facts.unpopulated_views,
},
sampling.as_ref(),
label,
)
.inspect_err(|_| warn_committed(&committed))?;
committed.push(label);
for (table, rows) in applied {
eprintln!(" \u{2713} {table}: {rows} row(s) scrubbed.");
}
if let (Some(sampling), Some(sampled)) = (sampling.as_ref(), sampled) {
report_sample_outcome(label, &sampled);
pending_compaction.push((
url.as_str(),
label.as_str(),
sampling,
purged_tables(&purges),
facts.all_materialized_views.clone(),
sampled.size_before,
));
}
}
if let Some(dir) = &args.output {
eprintln!("\u{2500}\u{2500} writing a scrubbed artifact \u{2500}\u{2500}");
super::backup::backup(&super::backup::BackupArgs {
profile: args.profile.clone(),
dir: Some(dir.clone()),
format: super::backup::BackupFormat::Custom,
keep: None,
target: super::backup::TargetSelector::All,
upload: false,
})?;
}
for (url, label, sampling, purged, refreshed, size_before) in pending_compaction {
report_reclaimed_size(url, label, sampling, &purged, &refreshed, size_before);
}
eprintln!("\n\u{2713} Scrub complete.");
Ok(())
}
fn build_sample_plan_for(
args: &ScrubArgs,
sources: &SourceClassification,
tables: &[Table],
facts: &DatabaseFacts,
) -> Result<sample::SamplePlan, ScrubError> {
let universe: Vec<(String, Vec<String>)> = tables
.iter()
.filter(|t| !facts.partitions.contains_key(&t.name))
.map(|t| (t.name.clone(), t.primary_key.clone()))
.collect();
let framework: BTreeSet<String> = facts.framework_tables.iter().cloned().collect();
let purged: BTreeSet<String> = purge_statements(&facts.framework_tables, &sources.config)
.into_iter()
.map(|(table, _)| table)
.collect();
Ok(sample::build_plan(&sample::SampleInputs {
roots: &sources.roots,
seed: args.seed,
rules: &sources.config.sample,
tables: &universe,
foreign_keys: &facts.foreign_keys,
framework_tables: &framework,
purged: &purged,
partitions: &facts.partitions,
})?)
}
fn report_sample_plan(label: &str, plan: &sample::SamplePlan) {
let roots: Vec<String> = plan
.tables
.iter()
.filter_map(|t| match t.role {
sample::SampleRole::Root(amount) => Some(match amount {
sample::SampleAmount::Percent(pct) => format!("{} {pct}%", t.table),
sample::SampleAmount::Count(n) => format!("{} {n} row(s)", t.table),
}),
_ => None,
})
.collect();
eprintln!(
" \u{2139} Sampling {label} from {}, seed {} \u{2014} the same seed against the same \
source selects the identical rows.",
roots.join(", "),
plan.seed,
);
}
fn report_sample_sql(plan: &sample::SamplePlan, counts: Option<&BTreeMap<String, i64>>) {
let empty = BTreeMap::new();
let counts = counts.unwrap_or(&empty);
for statement in plan.setup_statements() {
eprintln!(" {statement};");
}
for statement in plan.seed_statements(counts) {
eprintln!(" {statement};");
}
for statement in plan.index_statements() {
eprintln!(" {statement};");
}
eprintln!(" {};", plan.walk_loop_statement().replace('\n', "\n "));
for statement in plan.delete_statements() {
eprintln!(" {statement};");
}
for (constraint, statement) in plan.integrity_statements() {
eprintln!(
" {}; -- verifies {constraint}",
integrity_assertion(&statement)
);
}
for statement in sample_count_capture_statements(plan) {
eprintln!(" {statement};");
}
}
fn sample_count_capture_statements(plan: &sample::SamplePlan) -> Vec<String> {
if plan.tables.is_empty() {
return Vec::new();
}
let selects = plan
.tables
.iter()
.enumerate()
.map(|(index, table)| {
format!(
"SELECT {index} AS table_index, (SELECT count(*) FROM {}) AS n",
qualified_ident(&table.table)
)
})
.collect::<Vec<_>>()
.join(" UNION ALL ");
vec![format!(
"CREATE TEMPORARY TABLE {SAMPLE_COUNTS_TABLE} ON COMMIT DROP AS {selects}"
)]
}
const SAMPLE_COUNTS_TABLE: &str = "autumn_scrub_sample_counts_7c3e91d4a2";
fn sample_postcondition_statements(plan: &sample::SamplePlan) -> Vec<(String, String)> {
if plan.tables.is_empty() {
return Vec::new();
}
let tables = plan
.tables
.iter()
.map(|table| sample::comment_safe(&table.table))
.collect::<Vec<_>>()
.join(", ");
let mismatches = plan
.tables
.iter()
.enumerate()
.map(|(index, table)| {
format!(
"(SELECT count(*) FROM {}) <> \
(SELECT n FROM pg_temp.{SAMPLE_COUNTS_TABLE} WHERE table_index = {index})",
qualified_ident(&table.table)
)
})
.collect::<Vec<_>>()
.join(" OR ");
let body = format!(
"BEGIN IF {mismatches} THEN \
RAISE EXCEPTION 'a sampled table''s row count moved during the materialized-view refreshes'; \
END IF; END"
);
let tag = sample::dollar_tag(&body);
let mut statements = vec![(
format!("DO {tag} {body} {tag}"),
format!(
"sampled tables ({tables}): row counts still the sampled subset after the refreshes"
),
)];
for (constraint, statement) in plan.integrity_statements() {
statements.push((
integrity_assertion(&statement),
format!("re-verifies {constraint} after the refreshes"),
));
}
statements
}
fn report_sample_outcome(label: &str, outcome: &sample::SampleOutcome) {
eprintln!(" \u{2500}\u{2500} {label}: sampled rows \u{2500}\u{2500}");
for count in &outcome.counts {
eprintln!(
" {}: {} \u{2192} {} row(s) ({}, {})",
count.table,
count.before,
count.after,
percent_of(count.after, count.before),
count.role,
);
}
let before: i64 = outcome.counts.iter().map(|c| c.before).sum();
let after: i64 = outcome.counts.iter().map(|c| c.after).sum();
eprintln!(
" Total: {before} \u{2192} {after} row(s) ({} of the source), settled in {} pass(es).",
percent_of(after, before),
outcome.passes,
);
eprintln!(
" \u{2713} {} foreign key(s) re-verified \u{2014} every reference in the subset resolves.",
outcome.verified,
);
}
fn report_reclaimed_size(
url: &str,
label: &str,
plan: &sample::SamplePlan,
purged: &[String],
refreshed: &[String],
before: i64,
) {
let Ok(mut conn) = probe_connection(url, label, "compact the sampled tables") else {
warn_not_compacted(before, "could not connect to compact the sampled tables");
return;
};
if let Err(e) = sql_query(format!("SET lock_timeout = '{COMPACT_LOCK_TIMEOUT}'"))
.execute(&mut conn)
.map_err(|e| e.to_string())
{
warn_not_compacted(before, &e);
return;
}
for table in compacted_tables(plan, purged) {
if let Err(e) = sql_query(format!("VACUUM (FULL, ANALYZE) {}", qualified_ident(table)))
.execute(&mut conn)
{
warn_not_compacted(before, &format!("{table}: {e}"));
return;
}
}
let Ok(after) = sample::data_size(&mut conn, plan, &also_measured(purged, refreshed)) else {
warn_not_compacted(before, "could not measure the compacted size");
return;
};
eprintln!(
" Table size: {} \u{2192} {} ({} of the source).",
human_bytes(before),
human_bytes(after),
percent_of(after, before),
);
}
const COMPACT_LOCK_TIMEOUT: &str = "30s";
fn warn_not_compacted(before: i64, detail: &str) {
eprintln!(
" \u{26A0}\u{FE0F} The subset is committed, but the tables were NOT compacted \
({detail}).\n \
Deleting rows frees no disk on its own, so they still occupy {} \u{2014} close any \
other connection to this database and run `VACUUM (FULL, ANALYZE)` to reclaim it.",
human_bytes(before),
);
}
#[allow(clippy::cast_precision_loss)]
fn percent_of(part: i64, whole: i64) -> String {
if whole <= 0 {
return "n/a".to_owned();
}
format!("{:.1}%", part as f64 * 100.0 / whole as f64)
}
#[allow(clippy::cast_precision_loss)]
fn human_bytes(bytes: i64) -> String {
const UNITS: [&str; 5] = ["B", "kB", "MB", "GB", "TB"];
let mut value = bytes.max(0) as f64;
let mut unit = 0;
while value >= 1024.0 && unit + 1 < UNITS.len() {
value /= 1024.0;
unit += 1;
}
if unit == 0 {
format!("{bytes} B")
} else {
format!("{value:.1} {}", UNITS[unit])
}
}
fn report_triggers(plan: &ScrubPlan, sampling: Option<&sample::SamplePlan>, facts: &DatabaseFacts) {
let mut triggered: Vec<&str> = plan
.tables
.iter()
.map(|t| t.table.as_str())
.chain(
sampling
.into_iter()
.flat_map(sample::SamplePlan::subsetted_tables),
)
.filter(|t| facts.triggered_tables.contains(*t))
.collect();
triggered.sort_unstable();
triggered.dedup();
if triggered.is_empty() {
return;
}
eprintln!(
" \u{26A0}\u{FE0F} {} table(s) this run writes to carry user-defined triggers or \
rules: {}.\n \
An audit or history trigger copies the PRE-scrub row into another table as the \
rewrite or the sample's removals run, which can re-introduce real values; an \
`ON DELETE ... DO INSTEAD` rule can stop the sample removing anything at all. \
Check them, or disable them on the copy before scrubbing.",
triggered.len(),
triggered.join(", ")
);
}
fn report_plan(label: &str, plan: &ScrubPlan) {
eprintln!(
"\u{2500}\u{2500} {label} \u{2500}\u{2500}\n {} column(s) across {} table(s) \
classified as PII.",
plan.column_count(),
plan.tables.len()
);
for table in &plan.tables {
for column in &table.columns {
eprintln!(
" {}.{} \u{2192} {} ({})",
table.table,
column.column,
column.strategy.as_str(),
column.source.as_str()
);
}
}
}
const MV_REFRESH_CLOSURE: &str = "WITH RECURSIVE mv AS ( \
SELECT rel.oid, rel.relispopulated FROM pg_class rel \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
WHERE rel.relkind = 'm' \
), viewrule AS ( \
SELECT oid, ev_class FROM pg_rewrite WHERE rulename = '_RETURN' AND ev_type = '1' \
), dep_fn AS ( \
SELECT d.classid, d.objid, d.refobjid AS fn \
FROM pg_depend d \
WHERE d.classid IN ('pg_rewrite'::regclass, 'pg_proc'::regclass) \
AND d.refclassid = 'pg_proc'::regclass \
AND (d.classid = 'pg_proc'::regclass OR d.objid IN (SELECT oid FROM viewrule)) \
UNION \
SELECT d.classid, d.objid, o.oprcode \
FROM pg_depend d \
JOIN pg_operator o ON o.oid = d.refobjid \
WHERE d.classid IN ('pg_rewrite'::regclass, 'pg_proc'::regclass) \
AND d.refclassid = 'pg_operator'::regclass AND o.oprcode <> 0 \
AND (d.classid = 'pg_proc'::regclass OR d.objid IN (SELECT oid FROM viewrule)) \
UNION \
SELECT d.classid, d.objid, cd.refobjid \
FROM pg_depend d \
JOIN pg_constraint con ON con.contypid = d.refobjid \
JOIN pg_depend cd ON cd.classid = 'pg_constraint'::regclass \
AND cd.objid = con.oid AND cd.refclassid = 'pg_proc'::regclass \
WHERE d.classid IN ('pg_rewrite'::regclass, 'pg_proc'::regclass) \
AND d.refclassid = 'pg_type'::regclass \
AND (d.classid = 'pg_proc'::regclass OR d.objid IN (SELECT oid FROM viewrule)) \
), fn_reach AS ( \
SELECT df.objid AS rule, df.fn FROM dep_fn df \
WHERE df.classid = 'pg_rewrite'::regclass \
UNION \
SELECT r.rule, df.fn FROM fn_reach r \
JOIN dep_fn df ON df.classid = 'pg_proc'::regclass AND df.objid = r.fn \
), rel_edge AS ( \
SELECT DISTINCT r.ev_class AS dependent, d.refobjid AS source \
FROM pg_depend d \
JOIN viewrule r ON r.oid = d.objid \
WHERE d.classid = 'pg_rewrite'::regclass \
AND d.refclassid = 'pg_class'::regclass \
AND d.refobjid <> r.ev_class \
UNION \
SELECT DISTINCT r.ev_class AS dependent, fd.refobjid AS source \
FROM viewrule r \
JOIN fn_reach fr ON fr.rule = r.oid \
JOIN pg_depend fd ON fd.classid = 'pg_proc'::regclass \
AND fd.objid = fr.fn AND fd.refclassid = 'pg_class'::regclass \
WHERE fd.refobjid <> r.ev_class \
), reach AS ( \
SELECT m.oid AS dependent, e.source FROM mv m \
JOIN rel_edge e ON e.dependent = m.oid \
UNION \
SELECT h.dependent, e.source FROM reach h \
JOIN rel_edge e ON e.dependent = h.source \
WHERE h.source NOT IN (SELECT oid FROM mv) \
), edge AS ( \
SELECT DISTINCT dependent, source FROM reach \
WHERE source IN (SELECT oid FROM mv) \
), needed AS ( \
SELECT oid FROM mv WHERE relispopulated \
UNION \
SELECT e.source FROM edge e JOIN needed n ON n.oid = e.dependent \
), nedge AS ( \
SELECT dependent, source FROM edge WHERE dependent IN (SELECT oid FROM needed) \
)";
fn rules_reaching(extra: &str) -> String {
format!(
"SELECT DISTINCT rel.relname AS name FROM pg_rewrite r \
JOIN pg_class rel ON rel.oid = r.ev_class \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
WHERE r.rulename <> '_RETURN' AND r.ev_enabled IN ('O', 'A') {extra}"
)
}
fn triggers_reaching(extra: &str) -> String {
format!(
"WITH RECURSIVE fires AS ( \
SELECT t.tgrelid AS oid, (t.tgtype & 1) <> 0 AS by_row FROM pg_trigger t \
WHERE NOT t.tgisinternal AND t.tgenabled IN ('O', 'A') {extra}\
), ancestry AS ( \
SELECT oid, by_row FROM fires \
UNION \
SELECT i.inhparent, true FROM pg_inherits i \
JOIN ancestry a ON a.oid = i.inhrelid AND a.by_row \
) \
SELECT DISTINCT rel.relname AS name FROM ancestry a \
JOIN pg_class rel ON rel.oid = a.oid \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public'"
)
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ServerEndpoint {
pub database: String,
pub address: Option<String>,
pub port: Option<String>,
pub system_identifier: Option<String>,
pub server_port: Option<String>,
pub data_directory: Option<String>,
}
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct DatabaseFacts {
pub foreign_key_columns: BTreeSet<(String, String)>,
pub checked_columns: BTreeSet<(String, String)>,
pub generated_columns: BTreeSet<(String, String)>,
pub unique_columns: BTreeSet<(String, String)>,
pub nulls_not_distinct_columns: BTreeSet<(String, String)>,
pub partitions: BTreeMap<String, String>,
pub rls_tables: BTreeSet<String>,
pub legacy_inheritance: Vec<String>,
pub triggered_tables: BTreeSet<String>,
pub endpoint: ServerEndpoint,
pub replication_role: String,
pub delete_triggered_tables: BTreeSet<String>,
pub materialized_views: Vec<String>,
pub all_materialized_views: Vec<String>,
pub unpopulated_views: Vec<String>,
pub untraceable_view_functions: Vec<String>,
pub views_reading_by_name: Vec<String>,
pub other_schemas: BTreeSet<String>,
pub framework_tables: Vec<String>,
pub public_columns: BTreeSet<(String, String)>,
pub public_base_tables: BTreeSet<String>,
pub foreign_keys: Vec<sample::ForeignKeyConstraint>,
}
#[derive(diesel::QueryableByName)]
struct RowCount {
#[diesel(sql_type = diesel::sql_types::BigInt)]
n: i64,
}
#[derive(diesel::QueryableByName)]
struct NameRow {
#[diesel(sql_type = diesel::sql_types::Text)]
name: String,
}
#[derive(diesel::QueryableByName)]
struct PairRow {
#[diesel(sql_type = diesel::sql_types::Text)]
tbl: String,
#[diesel(sql_type = diesel::sql_types::Text)]
col: String,
}
#[derive(diesel::QueryableByName)]
struct ConstraintRow {
#[diesel(sql_type = diesel::sql_types::Text)]
name: String,
#[diesel(sql_type = diesel::sql_types::Text)]
child: String,
#[diesel(sql_type = diesel::sql_types::Array<diesel::sql_types::Text>)]
child_cols: Vec<String>,
#[diesel(sql_type = diesel::sql_types::Text)]
parent: String,
#[diesel(sql_type = diesel::sql_types::Array<diesel::sql_types::Text>)]
parent_cols: Vec<String>,
#[diesel(sql_type = diesel::sql_types::Bool)]
match_full: bool,
#[diesel(sql_type = diesel::sql_types::Bool)]
cloned: bool,
}
fn pair_rows(sql: &str, conn: &mut PgConnection) -> Result<BTreeMap<String, String>, ScrubError> {
let rows: Vec<PairRow> = sql_query(sql)
.load(conn)
.map_err(|e| ScrubError::Sql(e.to_string()))?;
Ok(rows.into_iter().map(|r| (r.tbl, r.col)).collect())
}
fn pair_set(rows: Vec<PairRow>) -> BTreeSet<(String, String)> {
rows.into_iter().map(|r| (r.tbl, r.col)).collect()
}
fn probe_connection(url: &str, label: &str, what: &str) -> Result<PgConnection, ScrubError> {
let mut conn = PgConnection::establish(url).map_err(|_| ScrubError::Introspect {
label: label.to_owned(),
detail: format!(
"could not connect to database {:?} to {what}",
parsed_db_name(url)
),
})?;
conn.batch_execute("SET search_path = pg_catalog, public")
.map_err(|e| ScrubError::Introspect {
label: label.to_owned(),
detail: format!("could not pin the search path to {what}: {e}"),
})?;
Ok(conn)
}
fn has_catalog_column(
conn: &mut PgConnection,
relation: &str,
column: &str,
) -> Result<bool, ScrubError> {
let rows: Vec<NameRow> = sql_query(format!(
"SELECT 'yes' AS name FROM pg_attribute \
WHERE attrelid = {}::regclass AND attname = {} AND NOT attisdropped",
quote_literal(relation),
quote_literal(column)
))
.load(conn)
.map_err(|e| ScrubError::Sql(e.to_string()))?;
Ok(!rows.is_empty())
}
#[allow(clippy::too_many_lines)]
fn probe_database_facts(
url: &str,
label: &str,
config: &ScrubConfig,
) -> Result<DatabaseFacts, ScrubError> {
let mut conn = probe_connection(url, label, "inspect its catalog")?;
let foreign_key_columns = pair_set(
sql_query(
"SELECT rel.relname AS tbl, att.attname AS col \
FROM pg_constraint c \
JOIN pg_class rel ON rel.oid = c.conrelid \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
CROSS JOIN LATERAL unnest(c.conkey) AS k(attnum) \
JOIN pg_attribute att ON att.attrelid = rel.oid AND att.attnum = k.attnum \
WHERE c.contype = 'f' \
UNION \
SELECT frel.relname AS tbl, fatt.attname AS col \
FROM pg_constraint c \
JOIN pg_class frel ON frel.oid = c.confrelid \
JOIN pg_namespace fns ON fns.oid = frel.relnamespace AND fns.nspname = 'public' \
CROSS JOIN LATERAL unnest(c.confkey) AS fk(attnum) \
JOIN pg_attribute fatt ON fatt.attrelid = frel.oid AND fatt.attnum = fk.attnum \
WHERE c.contype = 'f'",
)
.load(&mut conn)
.map_err(|e| ScrubError::Sql(e.to_string()))?,
);
let cloned = if has_catalog_column(&mut conn, "pg_constraint", "conparentid")? {
"c.conparentid <> 0"
} else {
"false"
};
let constraint_rows: Vec<ConstraintRow> = sql_query(format!(
"SELECT c.conname AS name, rel.relname AS child, frel.relname AS parent, \
(SELECT array_agg(att.attname::text ORDER BY k.ord) \
FROM unnest(c.conkey) WITH ORDINALITY AS k(attnum, ord) \
JOIN pg_attribute att ON att.attrelid = c.conrelid AND att.attnum = k.attnum) \
AS child_cols, \
(SELECT array_agg(att.attname::text ORDER BY k.ord) \
FROM unnest(c.confkey) WITH ORDINALITY AS k(attnum, ord) \
JOIN pg_attribute att ON att.attrelid = c.confrelid AND att.attnum = k.attnum) \
AS parent_cols, \
{cloned} AS cloned, \
c.confmatchtype = 'f' AS match_full \
FROM pg_constraint c \
JOIN pg_class rel ON rel.oid = c.conrelid \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
JOIN pg_class frel ON frel.oid = c.confrelid \
JOIN pg_namespace fns ON fns.oid = frel.relnamespace AND fns.nspname = 'public' \
WHERE c.contype = 'f'"
))
.load(&mut conn)
.map_err(|e| ScrubError::Sql(e.to_string()))?;
let foreign_keys: Vec<sample::ForeignKeyConstraint> = constraint_rows
.into_iter()
.filter(|row| !row.cloned)
.map(|row| sample::ForeignKeyConstraint {
name: row.name,
child_table: row.child,
child_columns: row.child_cols,
parent_table: row.parent,
parent_columns: row.parent_cols,
match_full: row.match_full,
})
.collect();
let checked_columns = pair_set(
sql_query(
"SELECT rel.relname AS tbl, att.attname AS col \
FROM pg_constraint c \
JOIN pg_class rel ON rel.oid = c.conrelid \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
CROSS JOIN LATERAL unnest(c.conkey) AS k(attnum) \
JOIN pg_attribute att ON att.attrelid = rel.oid AND att.attnum = k.attnum \
WHERE c.contype = 'c'",
)
.load(&mut conn)
.map_err(|e| ScrubError::Sql(e.to_string()))?,
);
let mut generated_predicates: Vec<&str> = Vec::new();
if has_catalog_column(&mut conn, "pg_attribute", "attgenerated")? {
generated_predicates.push("att.attgenerated <> ''");
}
if has_catalog_column(&mut conn, "pg_attribute", "attidentity")? {
generated_predicates.push("att.attidentity = 'a'");
}
let generated_columns = if generated_predicates.is_empty() {
BTreeSet::new()
} else {
pair_set(
sql_query(format!(
"SELECT rel.relname AS tbl, att.attname AS col \
FROM pg_attribute att \
JOIN pg_class rel ON rel.oid = att.attrelid \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
WHERE att.attnum > 0 AND NOT att.attisdropped AND ({})",
generated_predicates.join(" OR ")
))
.load(&mut conn)
.map_err(|e| ScrubError::Sql(e.to_string()))?,
)
};
let key_slice = if has_catalog_column(&mut conn, "pg_index", "indnkeyatts")? {
"i.indkey[0:i.indnkeyatts-1]"
} else {
"i.indkey"
};
let mut unique_columns = pair_set(
sql_query(format!(
"SELECT rel.relname AS tbl, att.attname AS col \
FROM pg_index i \
JOIN pg_class rel ON rel.oid = i.indrelid \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
JOIN pg_attribute att ON att.attrelid = rel.oid \
AND att.attnum = ANY({key_slice}) \
WHERE i.indisunique AND att.attnum > 0"
))
.load(&mut conn)
.map_err(|e| ScrubError::Sql(e.to_string()))?,
);
unique_columns.extend(pair_set(
sql_query(
"SELECT rel.relname AS tbl, att.attname AS col \
FROM pg_index i \
JOIN pg_class rel ON rel.oid = i.indrelid \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
JOIN pg_depend d ON d.objid = i.indexrelid AND d.refobjid = i.indrelid \
JOIN pg_attribute att ON att.attrelid = rel.oid AND att.attnum = d.refobjsubid \
WHERE i.indisunique AND d.refobjsubid > 0 \
AND (i.indexprs IS NOT NULL OR i.indpred IS NOT NULL)",
)
.load(&mut conn)
.map_err(|e| ScrubError::Sql(e.to_string()))?,
));
let nulls_not_distinct_columns =
if has_catalog_column(&mut conn, "pg_index", "indnullsnotdistinct")? {
pair_set(
sql_query(format!(
"SELECT rel.relname AS tbl, att.attname AS col \
FROM pg_index i \
JOIN pg_class rel ON rel.oid = i.indrelid \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
JOIN pg_attribute att ON att.attrelid = rel.oid \
AND att.attnum = ANY({key_slice}) \
WHERE i.indisunique AND i.indnullsnotdistinct AND att.attnum > 0"
))
.load(&mut conn)
.map_err(|e| ScrubError::Sql(e.to_string()))?,
)
} else {
BTreeSet::new()
};
let names = |q: &str, conn: &mut PgConnection| -> Result<Vec<String>, ScrubError> {
let rows: Vec<NameRow> = sql_query(q)
.load(conn)
.map_err(|e| ScrubError::Sql(e.to_string()))?;
Ok(rows.into_iter().map(|r| r.name).collect())
};
let partitions: BTreeMap<String, String> =
if has_catalog_column(&mut conn, "pg_class", "relispartition")? {
pair_rows(
"WITH RECURSIVE up AS ( \
SELECT rel.oid, rel.oid AS leaf FROM pg_class rel \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
WHERE rel.relispartition AND rel.relkind IN ('r', 'p', 'f') \
UNION ALL \
SELECT i.inhparent, up.leaf FROM pg_inherits i JOIN up ON up.oid = i.inhrelid \
) \
SELECT DISTINCT leafrel.relname AS tbl, rootrel.relname AS col \
FROM up \
JOIN pg_class leafrel ON leafrel.oid = up.leaf \
JOIN pg_class rootrel ON rootrel.oid = up.oid \
WHERE NOT rootrel.relispartition",
&mut conn,
)?
} else {
BTreeMap::new()
};
let legacy_inheritance: Vec<String> = {
let partition_filter = if has_catalog_column(&mut conn, "pg_class", "relispartition")? {
"AND NOT child.relispartition"
} else {
""
};
let rows: Vec<NameRow> = sql_query(format!(
"SELECT child.relname || ' (inherits ' || parent.relname || ')' AS name \
FROM pg_inherits i \
JOIN pg_class child ON child.oid = i.inhrelid \
JOIN pg_namespace cns ON cns.oid = child.relnamespace AND cns.nspname = 'public' \
JOIN pg_class parent ON parent.oid = i.inhparent \
JOIN pg_namespace pns ON pns.oid = parent.relnamespace AND pns.nspname = 'public' \
WHERE child.relkind IN ('r', 'p', 'f') {partition_filter}"
))
.load(&mut conn)
.map_err(|e| ScrubError::Sql(e.to_string()))?;
let mut names: Vec<String> = rows.into_iter().map(|r| r.name).collect();
names.sort();
names
};
let rls_tables = names(
"SELECT rel.relname AS name FROM pg_class rel \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
WHERE rel.relrowsecurity",
&mut conn,
)?
.into_iter()
.collect();
let mut triggered_tables: BTreeSet<String> = names(&triggers_reaching(""), &mut conn)?
.into_iter()
.collect();
triggered_tables.extend(names(&rules_reaching(""), &mut conn)?);
let replication_role = names(
"SELECT pg_catalog.current_setting('session_replication_role') AS name",
&mut conn,
)?
.into_iter()
.next()
.unwrap_or_else(|| "origin".to_owned());
let (address, port) = pair_rows(
"SELECT coalesce(pg_catalog.inet_server_addr()::text, '') AS tbl, \
coalesce(pg_catalog.inet_server_port()::text, '') AS col",
&mut conn,
)?
.into_iter()
.next()
.unwrap_or_default();
let optional = |query: &str, conn: &mut PgConnection| -> Option<String> {
names(query, conn)
.ok()
.and_then(|rows| rows.into_iter().next())
.filter(|value| !value.is_empty())
};
let setting_of = |parameter: &str| {
format!(
"SELECT coalesce((SELECT setting FROM pg_catalog.pg_settings \
WHERE name = {}), '') AS name",
quote_literal(parameter),
)
};
let endpoint = ServerEndpoint {
database: names("SELECT pg_catalog.current_database() AS name", &mut conn)?
.into_iter()
.next()
.unwrap_or_default(),
address: (!address.is_empty()).then_some(address),
port: (!port.is_empty()).then_some(port),
system_identifier: optional(
"SELECT system_identifier::text AS name FROM pg_catalog.pg_control_system()",
&mut conn,
),
server_port: optional(
"SELECT pg_catalog.current_setting('port') AS name",
&mut conn,
),
data_directory: optional(&setting_of("data_directory"), &mut conn),
};
let mut delete_triggered_tables: BTreeSet<String> =
names(&triggers_reaching("AND (t.tgtype & 8) <> 0 "), &mut conn)?
.into_iter()
.collect();
delete_triggered_tables.extend(names(&rules_reaching("AND r.ev_type = '4' "), &mut conn)?);
let other_schemas = names(
"SELECT DISTINCT ns.nspname AS name FROM pg_class rel \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace \
WHERE rel.relkind IN ('r', 'p', 'f', 'm') \
AND ns.nspname NOT IN ('public', 'information_schema') \
AND ns.nspname NOT LIKE 'pg\\_%'",
&mut conn,
)?
.into_iter()
.collect();
let materialized_views = names(
&format!(
"{MV_REFRESH_CLOSURE}, depth AS ( \
SELECT oid, 0 AS lvl FROM needed \
WHERE oid NOT IN (SELECT dependent FROM nedge) \
UNION ALL \
SELECT e.dependent, d.lvl + 1 FROM nedge e JOIN depth d ON d.oid = e.source \
WHERE d.lvl < 32 \
) \
SELECT rel.relname AS name \
FROM (SELECT oid, max(lvl) AS lvl FROM depth GROUP BY oid) o \
JOIN pg_class rel ON rel.oid = o.oid ORDER BY o.lvl, rel.relname"
),
&mut conn,
)?;
let all_materialized_views = names(
&format!(
"{MV_REFRESH_CLOSURE} \
SELECT rel.relname AS name FROM needed n \
JOIN pg_class rel ON rel.oid = n.oid ORDER BY rel.relname"
),
&mut conn,
)?;
let opaque_body = if has_catalog_column(&mut conn, "pg_proc", "prosqlbody")? {
"p.prokind <> 'a' AND p.prosqlbody IS NULL"
} else {
"p.prokind <> 'a'"
};
let untraceable_view_functions = names(
&format!(
"{MV_REFRESH_CLOSURE}, node AS ( \
SELECT oid AS root, oid AS relation FROM needed \
UNION \
SELECT h.dependent, h.source FROM reach h \
JOIN needed nd ON nd.oid = h.dependent \
) \
SELECT DISTINCT root.relname || ' via ' || p.proname AS name \
FROM node n \
JOIN pg_class root ON root.oid = n.root \
JOIN viewrule rw ON rw.ev_class = n.relation \
JOIN fn_reach fr ON fr.rule = rw.oid \
JOIN pg_proc p ON p.oid = fr.fn \
JOIN pg_namespace pn ON pn.oid = p.pronamespace \
WHERE pn.nspname NOT IN ('pg_catalog', 'information_schema') \
AND {opaque_body} \
ORDER BY name"
),
&mut conn,
)?;
let views_reading_by_name = names(
&format!(
"{MV_REFRESH_CLOSURE}, node AS ( \
SELECT oid AS root, oid AS relation FROM needed \
UNION \
SELECT h.dependent, h.source FROM reach h \
JOIN needed nd ON nd.oid = h.dependent \
) \
SELECT DISTINCT root.relname || ' via ' || p.proname AS name \
FROM node n \
JOIN pg_class root ON root.oid = n.root \
JOIN viewrule rw ON rw.ev_class = n.relation \
JOIN pg_rewrite rr ON rr.oid = rw.oid \
JOIN pg_proc p ON p.proname IN ( \
'query_to_xml', 'query_to_xmlschema', 'query_to_xml_and_xmlschema', \
'schema_to_xml', 'schema_to_xmlschema', 'schema_to_xml_and_xmlschema', \
'database_to_xml', 'database_to_xmlschema', \
'database_to_xml_and_xmlschema') \
JOIN pg_namespace pn ON pn.oid = p.pronamespace AND pn.nspname = 'pg_catalog' \
WHERE rr.ev_action ~ (':funcid ' || p.oid || '\\y') \
ORDER BY name"
),
&mut conn,
)?;
let unpopulated_views = names(
"SELECT rel.relname AS name FROM pg_class rel \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
WHERE rel.relkind = 'm' AND NOT rel.relispopulated ORDER BY rel.relname",
&mut conn,
)?;
let wanted = probe_table_names(config)
.iter()
.map(|t| quote_literal(t))
.collect::<Vec<_>>()
.join(", ");
let framework_tables = names(
&format!(
"SELECT rel.relname AS name FROM pg_class rel \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
WHERE rel.relkind IN ('r', 'p') AND rel.relname IN ({wanted}) \
ORDER BY rel.relname"
),
&mut conn,
)?;
let public_base_tables = names(
"SELECT rel.relname AS name FROM pg_class rel \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
WHERE rel.relkind IN ('r', 'p', 'f')",
&mut conn,
)?
.into_iter()
.collect();
let public_columns = pair_set(
sql_query(
"SELECT rel.relname AS tbl, att.attname AS col \
FROM pg_attribute att \
JOIN pg_class rel ON rel.oid = att.attrelid \
JOIN pg_namespace ns ON ns.oid = rel.relnamespace AND ns.nspname = 'public' \
WHERE rel.relkind IN ('r', 'p', 'f') AND att.attnum > 0 AND NOT att.attisdropped",
)
.load(&mut conn)
.map_err(|e| ScrubError::Sql(e.to_string()))?,
);
Ok(DatabaseFacts {
foreign_key_columns,
checked_columns,
generated_columns,
unique_columns,
nulls_not_distinct_columns,
partitions,
rls_tables,
legacy_inheritance,
triggered_tables,
endpoint,
replication_role,
delete_triggered_tables,
materialized_views,
all_materialized_views,
unpopulated_views,
untraceable_view_functions,
views_reading_by_name,
other_schemas,
framework_tables,
public_columns,
public_base_tables,
foreign_keys,
})
}
fn probe_table_names(config: &ScrubConfig) -> BTreeSet<String> {
FRAMEWORK_PAYLOAD_TABLES
.iter()
.map(|t| (*t).to_owned())
.chain(config.framework.purge.iter().cloned())
.collect()
}
fn report_framework_tables(present: &[String], config: &ScrubConfig) {
if present.is_empty() {
return;
}
let (purged, kept): (Vec<&String>, Vec<&String>) = present
.iter()
.partition(|t| config.framework.purge.contains(t));
for table in purged {
eprintln!(" {table} \u{2192} emptied (framework, [framework] purge)");
}
let kept: Vec<&&String> = kept
.iter()
.filter(|t| FRAMEWORK_PAYLOAD_TABLES.contains(&t.as_str()))
.collect();
if kept.is_empty() {
return;
}
eprintln!(
" \u{26A0}\u{FE0F} {} framework-owned table(s) are NOT scrubbed and may carry app-supplied \
payloads (queued jobs, offline-sync rows, experiment assignments):\n{}\n \
Add them to `[framework] purge = [...]` in {SCRUB_CONFIG_FILE} to empty them, or \
empty them yourself.",
kept.len(),
kept.iter()
.map(|t| format!(" - {t}"))
.collect::<Vec<_>>()
.join("\n"),
);
}
fn emptying_phases<'a>(
purges: &'a [(String, String)],
deferred: &BTreeSet<String>,
sampling: Option<&'a sample::SamplePlan>,
) -> EmptyingPhases<'a> {
let owned = |t: &'a (String, String)| (t.0.as_str(), t.1.clone());
let mut final_pass: Vec<(&str, String)> = purges.iter().map(owned).collect();
final_pass.extend(
sampling
.map(sample::SamplePlan::emptied_tables)
.unwrap_or_default(),
);
EmptyingPhases {
before: purges
.iter()
.filter(|(t, _)| !deferred.contains(t))
.map(owned)
.collect(),
after_sample: purges
.iter()
.filter(|(t, _)| deferred.contains(t))
.map(owned)
.collect(),
final_pass,
}
}
fn purged_tables(purges: &[(String, String)]) -> Vec<String> {
purges.iter().map(|(table, _)| table.clone()).collect()
}
type PendingCompaction<'a> = (
&'a str,
&'a str,
&'a sample::SamplePlan,
Vec<String>,
Vec<String>,
i64,
);
fn also_measured(purged: &[String], refreshed: &[String]) -> Vec<String> {
let mut names: Vec<String> = purged.iter().chain(refreshed).cloned().collect();
names.sort_unstable();
names.dedup();
names
}
fn compacted_tables<'a>(plan: &'a sample::SamplePlan, purged: &'a [String]) -> Vec<&'a str> {
let mut tables: Vec<&str> = plan
.subsetted_tables()
.into_iter()
.chain(purged.iter().map(String::as_str))
.collect();
tables.sort_unstable();
tables.dedup();
tables
}
fn psql_connect(label: &str, url: &str) -> Vec<String> {
password_free_conninfo(url).map_or_else(
|| {
vec![
format!(
"\\echo '{label}: no printable connection string (keyword form) \
-- connect to this target yourself, then delete the \\quit below'"
),
"\\quit".to_owned(),
]
},
|conninfo| {
vec![
"\\setenv PGHOSTADDR".to_owned(),
"\\setenv PGSERVICE".to_owned(),
format!(
"\\connect -reuse-previous=off {}",
quote_psql_arg(&conninfo)
),
]
},
)
}
fn endpoint_can_come_from_elsewhere(conninfo: &str) -> bool {
if std::env::var("PGHOSTADDR").is_ok_and(|v| !v.is_empty())
|| std::env::var("PGSERVICE").is_ok_and(|v| !v.is_empty())
{
return true;
}
let Ok(parsed) = url::Url::parse(conninfo) else {
return false;
};
crate::pg::parse_raw_query_pairs(parsed.query().unwrap_or(""))
.into_iter()
.any(|(key, value)| (key == "hostaddr" || key == "service") && !value.is_empty())
}
fn states_multiple_endpoints(conninfo: &str) -> bool {
let (host, port) = stated_host_and_port(conninfo);
let parts =
|value: Option<String>| value.map_or(0, |v| v.split(',').filter(|p| !p.is_empty()).count());
parts(host) > 1 || parts(port) > 1
}
fn stated_host_and_port(conninfo: &str) -> (Option<String>, Option<String>) {
let Ok(parsed) = url::Url::parse(conninfo) else {
return (None, None);
};
let mut host = parsed
.host_str()
.filter(|h| !h.is_empty())
.map(str::to_owned);
let mut port = parsed.port().map(|p| p.to_string());
for (key, value) in crate::pg::parse_raw_query_pairs(parsed.query().unwrap_or("")) {
if value.is_empty() {
continue;
}
match key.as_str() {
"host" => host = Some(value),
"port" => port = Some(value),
_ => {}
}
}
(host, port)
}
fn psql_connection_terms(conninfo: &str, database: &str) -> String {
let equals = |variable: &str, value: &str| {
format!(
":'{variable}' OPERATOR(pg_catalog.=) {}",
quote_literal(value)
)
};
let database_is = equals("DBNAME", database);
let (Some(host), port) = stated_host_and_port(conninfo) else {
return "false".to_owned();
};
let hosts: Vec<&str> = host.split(',').filter(|part| !part.is_empty()).collect();
let ports: Vec<&str> = port
.as_deref()
.map(|p| p.split(',').filter(|part| !part.is_empty()).collect())
.unwrap_or_default();
if hosts.is_empty() {
return "false".to_owned();
}
let any_host = || {
hosts
.iter()
.map(|h| equals("HOST", h))
.collect::<Vec<_>>()
.join(" OR ")
};
let endpoint = match (ports.len(), hosts.len()) {
(0, _) => return "false".to_owned(),
(1, _) => format!("({}) AND {}", any_host(), equals("PORT", ports[0])),
(p, h) if p == h => hosts
.iter()
.zip(&ports)
.map(|(host, port)| format!("({} AND {})", equals("HOST", host), equals("PORT", port)))
.collect::<Vec<_>>()
.join(" OR "),
_ => return "false".to_owned(),
};
format!("{database_is} AND ({endpoint})")
}
fn psql_connection_assertion(conninfo: &str, database: &str) -> Vec<String> {
vec![
"\\set autumn_ok false".to_owned(),
format!(
"SELECT ({}) AS autumn_ok \\gset",
psql_connection_terms(conninfo, database)
),
"\\if :autumn_ok".to_owned(),
]
}
fn target_guard(endpoint: &ServerEndpoint) -> String {
let literal =
|value: Option<&String>| value.map_or_else(|| "NULL".to_owned(), |v| quote_literal(v));
let setting = |parameter: &str| {
format!(
"coalesce((SELECT setting FROM pg_catalog.pg_settings WHERE name = {}), '')",
quote_literal(parameter),
)
};
let datadir = setting("data_directory");
let sysid_call = "(SELECT system_identifier::text FROM pg_catalog.pg_control_system())";
let body = format!(
" BEGIN IF {mismatch} THEN \
RAISE EXCEPTION {message}, {name}, {addr}, {port}, {sysid}, {server_port}, \
{datadir_want}, \
pg_catalog.current_database(), pg_catalog.inet_server_addr()::text, \
pg_catalog.inet_server_port()::text, {sysid_seen}, \
pg_catalog.current_setting('port'), {datadir_seen}; END IF; END ",
mismatch = endpoint_mismatch(endpoint),
sysid_seen = if endpoint.system_identifier.is_some() {
sysid_call.to_owned()
} else {
"NULL".to_owned()
},
datadir_seen = if endpoint.data_directory.is_some() {
format!("nullif({datadir}, '')")
} else {
"NULL".to_owned()
},
name = quote_literal(&endpoint.database),
addr = literal(endpoint.address.as_ref()),
port = literal(endpoint.port.as_ref()),
sysid = literal(endpoint.system_identifier.as_ref()),
server_port = literal(endpoint.server_port.as_ref()),
datadir_want = literal(endpoint.data_directory.as_ref()),
message = quote_literal(
"this block is for % at %:% (cluster %, port %, data directory %), but the \
session is on % at %:% (cluster %, port %, data directory %) — the \
\\connect above did not take effect (psql keeps the previous connection \
when one fails)"
),
);
let tag = sample::dollar_tag(&body);
format!("DO {tag}{body}{tag};")
}
fn replication_role_assertion() -> String {
let body = " BEGIN IF pg_catalog.current_setting('session_replication_role') <> 'origin' \
THEN RAISE EXCEPTION 'this block was planned against \
session_replication_role = origin, but this session is in % — a \
replica-only trigger fires here that the plan never inspected'\
, pg_catalog.current_setting('session_replication_role'); END IF; END "
.to_owned();
let tag = sample::dollar_tag(&body);
format!("DO {tag}{body}{tag};")
}
fn endpoint_mismatch(endpoint: &ServerEndpoint) -> String {
let literal =
|value: Option<&String>| value.map_or_else(|| "NULL".to_owned(), |v| quote_literal(v));
let mut terms = vec![
format!(
"pg_catalog.current_database() <> {}",
quote_literal(&endpoint.database)
),
format!(
"pg_catalog.inet_server_addr()::text IS DISTINCT FROM {}",
literal(endpoint.address.as_ref())
),
format!(
"pg_catalog.inet_server_port()::text IS DISTINCT FROM {}",
literal(endpoint.port.as_ref())
),
format!(
"pg_catalog.current_setting('port') IS DISTINCT FROM {}",
literal(endpoint.server_port.as_ref())
),
];
if let Some(sysid) = endpoint.system_identifier.as_ref() {
terms.push(format!(
"(SELECT system_identifier::text FROM pg_catalog.pg_control_system()) \
IS DISTINCT FROM {}",
quote_literal(sysid)
));
}
if let Some(datadir) = endpoint.data_directory.as_ref() {
terms.push(format!(
"nullif(coalesce((SELECT setting FROM pg_catalog.pg_settings \
WHERE name = 'data_directory'), ''), '') IS DISTINCT FROM {}",
quote_literal(datadir)
));
}
terms.join(" OR ")
}
fn post_commit_fence(endpoint: &ServerEndpoint) -> Vec<String> {
vec![
"\\set autumn_commit_error :ERROR".to_owned(),
"SET search_path = pg_catalog, public;".to_owned(),
"\\set autumn_ok false".to_owned(),
format!(
"SELECT (:'autumn_scrubbed'::bool AND NOT :'autumn_commit_error'::bool \
AND NOT ({})) AS autumn_ok \\gset",
endpoint_mismatch(endpoint),
),
]
}
const SECRET_KEYWORDS: [&str; 3] = ["password", "sslpassword", "oauth_client_secret"];
const PRINTABLE_KEYWORDS: [&str; 47] = [
"application_name",
"channel_binding",
"client_encoding",
"connect_timeout",
"dbname",
"fallback_application_name",
"gssdelegation",
"gssencmode",
"gsslib",
"host",
"hostaddr",
"keepalives",
"keepalives_count",
"keepalives_idle",
"keepalives_interval",
"krbsrvname",
"load_balance_hosts",
"max_protocol_version",
"min_protocol_version",
"oauth_client_id",
"oauth_discovery_url",
"oauth_issuer",
"oauth_scope",
"options",
"passfile",
"port",
"replication",
"require_auth",
"requirepeer",
"requiressl",
"service",
"ssl_max_protocol_version",
"ssl_min_protocol_version",
"sslcert",
"sslcertmode",
"sslcompression",
"sslcrl",
"sslcrldir",
"sslkey",
"sslkeylogfile",
"sslmode",
"sslnegotiation",
"sslrootcert",
"sslsni",
"target_session_attrs",
"tcp_user_timeout",
"user",
];
fn password_free_conninfo(url: &str) -> Option<String> {
let mut parsed = url::Url::parse(url).ok()?;
parsed.set_password(None).ok()?;
let pairs = crate::pg::parse_raw_query_pairs(parsed.query().unwrap_or(""));
if pairs
.iter()
.any(|(key, _)| !is_secret_keyword(key) && !is_printable_keyword(key))
{
return None;
}
let kept: Vec<(String, String)> = pairs
.into_iter()
.filter(|(key, _)| !is_secret_keyword(key))
.collect();
if kept.is_empty() {
parsed.set_query(None);
} else {
let query = kept
.iter()
.map(|(key, value)| {
format!(
"{}={}",
crate::pg::query_value_token(key),
crate::pg::query_value_token(value),
)
})
.collect::<Vec<_>>()
.join("&");
parsed.set_query(Some(&query));
}
Some(parsed.into())
}
fn is_secret_keyword(key: &str) -> bool {
SECRET_KEYWORDS
.iter()
.any(|secret| key.eq_ignore_ascii_case(secret))
}
fn is_printable_keyword(key: &str) -> bool {
PRINTABLE_KEYWORDS
.iter()
.any(|known| key.eq_ignore_ascii_case(known))
}
fn quote_psql_arg(value: &str) -> String {
let escaped: String = value
.chars()
.flat_map(|c| {
let escape = matches!(c, '"' | '\\');
escape.then_some('\\').into_iter().chain(std::iter::once(c))
})
.collect();
format!("\"{escaped}\"")
}
fn session_settings() -> Vec<String> {
[
"SET LOCAL search_path = pg_catalog, public",
"SET LOCAL standard_conforming_strings = on",
"SET LOCAL DateStyle = 'ISO, YMD'",
"SET LOCAL IntervalStyle = 'iso_8601'",
"SET LOCAL TimeZone = 'UTC'",
"SET LOCAL bytea_output = 'hex'",
"SET LOCAL lc_monetary = 'C'",
"SET LOCAL extra_float_digits = 3",
]
.into_iter()
.map(str::to_owned)
.collect()
}
fn lock_statements(
plan: &ScrubPlan,
purges: &[(String, String)],
sampling: Option<&sample::SamplePlan>,
) -> Vec<String> {
let mut locked: Vec<&str> = plan
.tables
.iter()
.map(|t| t.table.as_str())
.chain(purges.iter().map(|(table, _)| table.as_str()))
.chain(
sampling
.into_iter()
.flat_map(sample::SamplePlan::locked_tables),
)
.collect();
locked.sort_unstable();
locked.dedup();
locked
.into_iter()
.map(|table| {
format!(
"LOCK TABLE {} IN SHARE ROW EXCLUSIVE MODE",
qualified_ident(table)
)
})
.collect()
}
fn integrity_assertion(check: &str) -> String {
let body = format!(
"BEGIN IF ({check}) > 0 THEN \
RAISE EXCEPTION 'a foreign key this run checked does not resolve'; \
END IF; END"
);
let tag = sample::dollar_tag(&body);
format!("DO {tag} {body} {tag}")
}
fn emptiness_assertion(table: &str) -> String {
let body = format!(
"BEGIN IF EXISTS (SELECT 1 FROM {}) THEN \
RAISE EXCEPTION 'a table this run promised would be empty still holds rows'; \
END IF; END",
qualified_ident(table)
);
let tag = sample::dollar_tag(&body);
format!("DO {tag} {body} {tag}")
}
struct EmptyingPhases<'a> {
before: Vec<(&'a str, String)>,
after_sample: Vec<(&'a str, String)>,
final_pass: Vec<(&'a str, String)>,
}
fn purge_statements(present: &[String], config: &ScrubConfig) -> Vec<(String, String)> {
present
.iter()
.filter(|t| config.framework.purge.contains(t))
.map(|t| (t.clone(), format!("DELETE FROM {}", qualified_ident(t))))
.collect()
}
fn resolve_key_ring(
profile: &str,
project_root: &Path,
) -> Result<autumn_web::encryption::KeyRing, ScrubError> {
let store = autumn_web::credentials::load_credentials(profile, project_root).map_err(|e| {
ScrubError::EncryptionKeyUnavailable {
profile: profile.to_owned(),
detail: e.to_string(),
}
})?;
match autumn_web::encryption::key_ring_from_credentials(&store) {
Ok(Some(ring)) => Ok(ring),
Ok(None) => Err(ScrubError::EncryptionKeyUnavailable {
profile: profile.to_owned(),
detail: format!(
"`{}.primary_key` is not configured",
autumn_web::encryption::CREDENTIALS_NAMESPACE
),
}),
Err(e) => Err(ScrubError::EncryptionKeyUnavailable {
profile: profile.to_owned(),
detail: e.to_string(),
}),
}
}
const ENCRYPTED_BATCH_ROWS: usize = 500;
#[derive(diesel::QueryableByName)]
struct RowTokenRow {
#[diesel(sql_type = diesel::sql_types::Text)]
row_key: String,
#[diesel(sql_type = diesel::sql_types::Text)]
token: String,
}
fn rewrite_encrypted_column(
conn: &mut PgConnection,
table: &str,
row_key: &str,
rewrite: &EncryptedRewrite,
) -> Result<usize, diesel::result::Error> {
use autumn_web::encryption::{Mode, encrypt_text};
let ident = quote_ident(&rewrite.column);
let seed = format!("{} || '|' || ({row_key})", quote_literal(&rewrite.column));
let rows: Vec<RowTokenRow> = sql_query(format!(
"SELECT ({row_key}) AS row_key, md5({seed}) || md5(({seed}) || '#2') AS token \
FROM {} WHERE {ident} IS NOT NULL",
qualified_ident(table),
))
.load(conn)?;
let mode = if rewrite.deterministic {
Mode::Deterministic
} else {
Mode::Randomized
};
let mut updated = 0;
for chunk in rows.chunks(ENCRYPTED_BATCH_ROWS) {
let mut values = Vec::with_capacity(chunk.len());
for row in chunk {
let plaintext = match rewrite.shape {
Strategy::Email => format!("scrubbed+{}{SCRUB_EMAIL_DOMAIN}", row.token),
_ => format!("[scrubbed:{}]", row.token),
};
let envelope = encrypt_text(mode, &plaintext).map_err(|e| {
diesel::result::Error::QueryBuilderError(
format!(
"could not encrypt a replacement for {table}.{}: {e}",
rewrite.column
)
.into(),
)
})?;
values.push(format!(
"({}, {})",
quote_literal(&row.row_key),
quote_literal(&envelope)
));
}
updated += sql_query(format!(
"UPDATE {} AS t SET {ident} = v.val FROM (VALUES {}) AS v(k, val) \
WHERE ({row_key}) = v.k",
qualified_ident(table),
values.join(", ")
))
.execute(conn)?;
}
Ok(updated)
}
type Applied = (Vec<(String, usize)>, Option<sample::SampleOutcome>);
struct ViewRefresh<'a> {
ordered: &'a [String],
all: &'a [String],
unpopulated: &'a [String],
}
fn verify_sample_survived_refreshes(
conn: &mut PgConnection,
settled: &sample::SampleOutcome,
plan: &sample::SamplePlan,
) -> Result<(), sample::SampleFailure> {
let mut moved = Vec::new();
for count in &settled.counts {
let now: RowCount = sql_query(format!(
"SELECT count(*) AS n FROM {}",
qualified_ident(&count.table)
))
.get_result(conn)?;
if now.n != count.after {
moved.push(format!(
"{} ({} row(s), sampled to {})",
count.table, now.n, count.after
));
}
}
if !moved.is_empty() {
moved.sort();
return Err(sample::SampleFailure::Refused(
sample::SampleError::SampleMutatedByRefresh { tables: moved },
));
}
let mut violations = Vec::new();
for (label, sql) in &plan.integrity_statements() {
let orphans: RowCount = sql_query(sql).get_result(conn)?;
if orphans.n > 0 {
violations.push(format!("{label}: {} unresolved reference(s)", orphans.n));
}
}
if !violations.is_empty() {
violations.sort();
return Err(sample::SampleFailure::Refused(
sample::SampleError::IntegrityViolation { violations },
));
}
Ok(())
}
fn execute(
url: &str,
plan: &ScrubPlan,
purges: &[(String, String)],
views: &ViewRefresh<'_>,
sampling: Option<&sample::SamplePlan>,
label: &str,
) -> Result<Applied, ScrubError> {
if plan.tables.is_empty() && purges.is_empty() && views.ordered.is_empty() && sampling.is_none()
{
return Ok((Vec::new(), None));
}
let mut conn = probe_connection(url, label, "apply the scrub")?;
let mut counts = Vec::with_capacity(plan.tables.len());
let mut outcome = None;
conn.transaction::<_, sample::SampleFailure, _>(|conn| {
conn.batch_execute(&session_settings().join("; "))?;
for statement in lock_statements(plan, purges, sampling) {
sql_query(statement).execute(conn)?;
}
let no_deferral = BTreeSet::new();
let deferred: &BTreeSet<String> = sampling.map_or(&no_deferral, |s| &s.purge_after);
let phases = emptying_phases(purges, deferred, sampling);
let mut purged_rows: BTreeMap<&str, usize> = BTreeMap::new();
for (table, statement) in &phases.before {
let rows = sql_query(statement).execute(conn)?;
*purged_rows.entry(*table).or_default() += rows;
}
if let Some(sampling) = sampling {
outcome = Some(sample::apply(
conn,
sampling,
&also_measured(&purged_tables(purges), views.all),
)?);
}
for (table, statement) in &phases.after_sample {
let rows = sql_query(statement).execute(conn)?;
*purged_rows.entry(*table).or_default() += rows;
}
for table in &plan.tables {
if let Some(sql) = &table.sql {
let rows = sql_query(sql).execute(conn)?;
counts.push((table.table.clone(), rows));
}
for rewrite in &table.encrypted {
let rows = rewrite_encrypted_column(conn, &table.table, &table.row_key, rewrite)?;
counts.push((format!("{}.{}", table.table, rewrite.column), rows));
}
}
for (table, statement) in &phases.final_pass {
let rows = sql_query(statement).execute(conn)?;
*purged_rows.entry(*table).or_default() += rows;
}
for (table, rows) in purged_rows {
counts.push((format!("{table} (emptied)"), rows));
}
for view in views.ordered {
sql_query(format!(
"REFRESH MATERIALIZED VIEW {}",
qualified_ident(view)
))
.execute(conn)?;
counts.push((format!("{view} (materialized view refreshed)"), 0));
}
for view in views.unpopulated {
sql_query(format!(
"REFRESH MATERIALIZED VIEW {} WITH NO DATA",
qualified_ident(view)
))
.execute(conn)?;
counts.push((format!("{view} (materialized view left unpopulated)"), 0));
}
let mut refilled = Vec::new();
for (table, _) in &phases.final_pass {
let row: RowCount = sql_query(format!(
"SELECT count(*) AS n FROM {}",
qualified_ident(table)
))
.get_result(conn)?;
if row.n > 0 {
refilled.push(format!("{table} ({} row(s))", row.n));
}
}
if !refilled.is_empty() {
refilled.sort();
refilled.dedup();
return Err(sample::SampleFailure::Refused(
sample::SampleError::NotEmptied { tables: refilled },
));
}
if let (Some(settled), Some(plan)) = (outcome.as_ref(), sampling) {
verify_sample_survived_refreshes(conn, settled, plan)?;
}
Ok(())
})
.map_err(|e| match e {
sample::SampleFailure::Refused(refusal) => ScrubError::from(refusal),
sample::SampleFailure::Db(error) => ScrubError::Sql(error.to_string()),
})?;
Ok((counts, outcome))
}
#[cfg(test)]
mod tests {
use super::{emptiness_assertion, integrity_assertion, rules_reaching, triggers_reaching};
#[test]
fn an_encrypted_rewrite_withholds_the_whole_printed_script() {
let refusal = ScrubError::UnprintableEncryptedRewrite {
columns: vec!["control: users.api_token".to_owned()],
}
.to_string();
assert!(
refusal.contains("cannot print a runnable script"),
"it must refuse the script, not annotate it: {refusal}"
);
assert!(
refusal.contains("users.api_token"),
"and name the column that cannot be printed: {refusal}"
);
assert!(
refusal.contains("without --dry-run"),
"and say what to run instead: {refusal}"
);
assert!(
refusal.contains("plan above is complete"),
"and keep the reported plan standing: {refusal}"
);
assert!(
!refusal.to_lowercase().contains("primary_key")
&& !refusal.contains("deterministic_key")
&& !refusal.contains("key_derivation_salt"),
"no key material may reach the refusal: {refusal}"
);
}
#[test]
fn the_compaction_after_commit_is_fenced_by_psql() {
let endpoint = super::ServerEndpoint {
database: "app".to_owned(),
address: None,
port: None,
system_identifier: Some("7682669380557907941".to_owned()),
server_port: Some("5435".to_owned()),
data_directory: Some("/tmp/pgd3".to_owned()),
};
let fence = super::post_commit_fence(&endpoint);
assert!(
fence
.iter()
.any(|line| line == "SET search_path = pg_catalog, public;"),
"the fence must re-pin the path it needs: {fence:?}"
);
assert!(
fence
.iter()
.any(|line| line.contains("\\gset") && line.contains("AS autumn_ok")),
"psql, not the server, has to decide this one: {fence:?}"
);
assert!(
!fence.iter().any(|line| line == "\\if :autumn_ok"),
"the fence sets the flag; the compaction pass reads it: {fence:?}"
);
assert_eq!(
fence
.iter()
.filter(|line| *line == "\\set autumn_ok false")
.count(),
1,
"the flag must be cleared before the gset that may not run: {fence:?}"
);
assert!(
fence
.iter()
.any(|line| line.contains("autumn_scrubbed") && line.contains("::bool")),
"the fence must require the transaction to have succeeded: {fence:?}"
);
assert_eq!(
fence.first().map(String::as_str),
Some("\\set autumn_commit_error :ERROR"),
"the commit's own error state has to be read before it is lost: {fence:?}"
);
let mismatch = super::endpoint_mismatch(&endpoint);
assert!(
fence.iter().any(|line| line.contains(&mismatch)),
"the fence must ask exactly what the guard asks: {fence:?}"
);
assert!(
super::target_guard(&endpoint).contains(&mismatch),
"and the guard must ask it too: {mismatch}"
);
}
#[test]
fn the_target_guard_pins_the_whole_endpoint_and_qualifies_every_call() {
let here = super::ServerEndpoint {
database: "app_copy".to_owned(),
address: Some("10.0.0.2/32".to_owned()),
port: Some("5432".to_owned()),
system_identifier: Some("7682669380557907941".to_owned()),
server_port: Some("5432".to_owned()),
data_directory: Some("/var/lib/postgresql/16/main".to_owned()),
};
let guard = super::target_guard(&here);
assert!(
guard.contains("pg_catalog.current_database() <> 'app_copy'")
&& guard
.contains("pg_catalog.inet_server_addr()::text IS DISTINCT FROM '10.0.0.2/32'")
&& guard.contains("pg_catalog.inet_server_port()::text IS DISTINCT FROM '5432'"),
"the guard must pin the whole endpoint: {guard}"
);
assert!(
guard.contains("IS DISTINCT FROM '7682669380557907941'"),
"the cluster identity must be pinned too, or two socket clusters are \
indistinguishable: {guard}"
);
assert!(
guard.contains("pg_catalog.current_setting('port') IS DISTINCT FROM '5432'"),
"the server's configured port must be pinned: {guard}"
);
assert!(
guard.contains("IS DISTINCT FROM '/var/lib/postgresql/16/main'"),
"the data directory must be pinned, or a clone passes: {guard}"
);
assert!(
!guard.contains("current_setting('data_directory')")
&& guard.contains("FROM pg_catalog.pg_settings WHERE name = 'data_directory'"),
"a restricted setting must degrade to NULL, not raise: {guard}"
);
for call in [
"current_database()",
"inet_server_addr()",
"inet_server_port()",
"pg_control_system()",
] {
for (at, _) in guard.match_indices(call) {
assert!(
guard[..at].ends_with("pg_catalog."),
"every catalog call must be qualified: {call} at {at} in {guard}"
);
}
}
assert!(
!guard.contains("$s"),
"RAISE has no positional format specifiers: {guard}"
);
assert_eq!(
guard.matches('%').count(),
12,
"six placeholders for the target endpoint, six for the session's: {guard}"
);
let socket = super::target_guard(&super::ServerEndpoint {
database: "app_copy".to_owned(),
..super::ServerEndpoint::default()
});
assert!(
socket.contains("IS DISTINCT FROM NULL"),
"a socket target pins NULL explicitly: {socket}"
);
assert_eq!(
socket.matches("IS DISTINCT FROM NULL").count(),
3,
"address, port and the configured port are compared even when NULL — \
for a socket target that IS the answer: {socket}"
);
assert!(
!socket.contains("pg_control_system()) IS DISTINCT FROM")
&& !socket.contains("'data_directory'), ''), '') IS DISTINCT FROM"),
"an unreadable discriminator must be dropped, not guessed at: {socket}"
);
assert!(
!socket.contains("pg_catalog.pg_control_system()"),
"and the message must not call for it either: {socket}"
);
let hostile = super::target_guard(&super::ServerEndpoint {
database: "it's".to_owned(),
..super::ServerEndpoint::default()
});
assert!(
hostile.contains("'it''s'"),
"a quote in the database name must not break out of the literal: {hostile}"
);
}
#[test]
fn the_reconnect_is_proved_from_psql_own_variables() {
let lines = super::psql_connection_assertion(
"postgres://postgres@127.0.0.1:25433/dry_t2",
"dry_t2",
);
assert_eq!(
lines.first().map(String::as_str),
Some("\\set autumn_ok false"),
"the flag must be cleared before the gset that may not run: {lines:?}"
);
let probe = lines
.iter()
.find(|l| l.contains("\\gset"))
.unwrap_or_else(|| panic!("psql, not the server, has to decide this: {lines:?}"));
assert!(
probe.contains(":'PORT' OPERATOR(pg_catalog.=) '25433'")
&& probe.contains(":'HOST' OPERATOR(pg_catalog.=) '127.0.0.1'")
&& probe.contains(":'DBNAME' OPERATOR(pg_catalog.=) 'dry_t2'"),
"every stated component must be pinned to psql's own value: {probe}"
);
assert!(
!probe.contains(" = '"),
"the comparison must not resolve through search_path: {probe}"
);
assert_eq!(
lines.last().map(String::as_str),
Some("\\if :autumn_ok"),
"and the block has to sit inside it: {lines:?}"
);
}
#[test]
fn a_query_pair_overrides_the_authority_it_duplicates() {
let overridden = super::psql_connection_assertion(
"postgres://postgres@a.example:5432/app?port=6000",
"app",
)
.into_iter()
.find(|l| l.contains("\\gset"))
.expect("the probe must be emitted");
assert!(
overridden.contains(":'PORT' OPERATOR(pg_catalog.=) '6000'")
&& !overridden.contains("'5432'"),
"the query pair wins over the authority: {overridden}"
);
}
#[test]
fn a_failover_list_keeps_its_host_port_pairing() {
let paired = super::psql_connection_terms(
"postgres://postgres@a.example/app?host=a.example,b.example&port=5432,5433",
"app",
);
assert!(
paired.contains(
"(:'HOST' OPERATOR(pg_catalog.=) 'a.example' AND :'PORT' OPERATOR(pg_catalog.=) '5432')"
) && paired.contains(
"(:'HOST' OPERATOR(pg_catalog.=) 'b.example' AND :'PORT' OPERATOR(pg_catalog.=) '5433')"
),
"each host must carry its OWN port: {paired}"
);
let broadcast = super::psql_connection_terms(
"postgres://postgres@a.example/app?host=a.example,b.example&port=5432",
"app",
);
assert!(
broadcast.contains(":'HOST' OPERATOR(pg_catalog.=) 'a.example'")
&& broadcast.contains(":'HOST' OPERATOR(pg_catalog.=) 'b.example'")
&& broadcast.matches(":'PORT'").count() == 1,
"one port must cover every host: {broadcast}"
);
let mismatched = super::psql_connection_terms(
"postgres://postgres@a.example/app?host=a.example,b.example&port=1,2,3",
"app",
);
assert_eq!(
mismatched, "false",
"a pairing libpq itself rejects must not be guessed at: {mismatched}"
);
}
#[test]
fn a_conninfo_without_a_host_proves_nothing() {
assert_eq!(
super::psql_connection_terms("postgres:///app", "app"),
"false",
"the one discriminator this proof exists for cannot be optional"
);
assert_eq!(
super::psql_connection_terms("postgres://db.internal/app", "app"),
"false",
"accepting any port for the host drops the discriminator"
);
}
#[test]
fn the_pasting_session_must_be_in_the_origin_replication_role() {
let assertion = super::replication_role_assertion();
assert!(
assertion
.contains("pg_catalog.current_setting('session_replication_role') <> 'origin'")
&& assertion.contains("RAISE EXCEPTION"),
"the script must assert what the command refuses to plan without: {assertion}"
);
assert!(
!assertion.contains("SET LOCAL"),
"it must not try to reset a setting an ordinary role cannot: {assertion}"
);
}
#[test]
fn the_measured_set_covers_purges_and_refreshed_views_once_each() {
let purged = vec!["autumn_jobs".to_owned(), "shared".to_owned()];
let refreshed = vec!["shared".to_owned(), "big_report".to_owned()];
assert_eq!(
super::also_measured(&purged, &refreshed),
vec![
"autumn_jobs".to_owned(),
"big_report".to_owned(),
"shared".to_owned()
],
"every relation the run touches outside the plan, and each one once"
);
assert!(
super::also_measured(&[], &[]).is_empty(),
"and an unsampled target with neither measures nothing extra"
);
}
#[test]
fn the_connect_boundary_carries_the_whole_endpoint_and_no_password() {
let lines = super::psql_connect(
"control",
"postgres://scrubby:hunter2@db1.internal:6543/app",
);
assert_eq!(
&lines[..2],
&[r"\setenv PGHOSTADDR", r"\setenv PGSERVICE"],
"the boundary must clear libpq's environment-routed endpoint knobs first: {lines:?}"
);
let line = lines.join("\n");
assert!(
lines[2].starts_with(r#"\connect -reuse-previous=off ""#),
"the boundary must inherit nothing from the previous connection: {line}"
);
assert!(
!line.contains("hunter2"),
"the password must never reach the printed script: {line}"
);
assert!(
line.contains("db1.internal") && line.contains("6543") && line.contains("app"),
"and the endpoint must survive: {line}"
);
let overridden = super::psql_connect(
"control",
"postgres://authority/app?host=queryhost&dbname=copy",
)
.join("\n");
assert!(
overridden.contains("host=queryhost") && overridden.contains("dbname=copy"),
"an override must survive into the printed boundary: {overridden}"
);
let spaced = super::psql_connect(
"control",
"postgres://u@h/app?options=-c%20search_path%3Dpg_catalog&password=hunter2",
)
.join("\n");
assert!(
!spaced.contains('+'),
"a space must be percent-encoded, never written as `+`: {spaced}"
);
assert!(
spaced.contains("%20"),
"and it must still be there, not dropped: {spaced}"
);
assert!(
!spaced.contains("hunter2") && !spaced.contains("password"),
"while the credential is still removed: {spaced}"
);
for keyword in [
"host=db2.internal dbname=app password=hunter2",
"host=db2.internal dbname=app password = hunter2",
"host=db2 password='two words' dbname=app",
] {
let line = super::psql_connect("control", keyword).join("\n");
assert!(
!line.contains("hunter2") && !line.contains("two words"),
"no keyword form may print its password: {line}"
);
assert!(
!line.contains("\\connect"),
"and none may claim to move the session: {line}"
);
assert!(
line.contains("control"),
"the operator must still be told which target it was: {line}"
);
assert!(
line.contains("\\quit"),
"a boundary that cannot be printed must halt psql: {line}"
);
}
let in_query =
super::psql_connect("control", "postgres://bob@db/app?password=secret2").join("\n");
assert!(
!in_query.contains("secret2"),
"a query-string password must be stripped too: {in_query}"
);
let both = super::psql_connect(
"control",
"postgres://alice:secret1@db/app?password=secret2&sslpassword=k3y&application_name=x",
)
.join("\n");
assert!(
!both.contains("secret1") && !both.contains("secret2") && !both.contains("k3y"),
"every credential form must go: {both}"
);
assert!(
both.contains("application_name=x"),
"and every non-secret parameter must stay: {both}"
);
}
#[test]
fn oauth_client_secret_is_stripped_not_printed() {
let printed = super::psql_connect(
"control",
"postgres://u@db/app?oauth_client_secret=s3cr3t&oauth_issuer=https%3A%2F%2Fidp.example",
)
.join("\n");
assert!(
printed.contains("\\connect"),
"the target must stay printable: {printed}"
);
assert!(
!printed.contains("s3cr3t"),
"the OAuth client secret must never reach the printed script: {printed}"
);
assert!(
printed.contains("oauth_issuer="),
"the non-secret OAuth parameters must survive: {printed}"
);
let upper = super::password_free_conninfo("postgres://u@db/app?OAUTH_CLIENT_SECRET=s3cr3t")
.expect("a URI must stay printable");
assert!(
!upper.contains("s3cr3t"),
"OAUTH_CLIENT_SECRET must strip like oauth_client_secret: {upper}"
);
}
#[test]
fn an_unknown_keyword_refuses_the_target() {
assert!(
super::password_free_conninfo("postgres://u@db/app?widgets=1").is_none(),
"an unrecognised keyword must make the conninfo unprintable"
);
let refused = super::psql_connect("control", "postgres://u@db/app?widgets=1").join("\n");
assert!(
!refused.contains("\\connect"),
"and the boundary must not claim to move the session: {refused}"
);
assert!(
refused.contains("\\quit"),
"a boundary that cannot be printed must halt psql: {refused}"
);
}
#[test]
fn supported_non_secret_libpq_options_print() {
let url = "postgres://u@db/app?sslcertmode=allow&sslsni=1&min_protocol_version=3.0\
&max_protocol_version=latest&oauth_client_id=cli";
let conninfo =
super::password_free_conninfo(url).expect("supported non-secret options must print");
for key in [
"sslcertmode",
"sslsni",
"min_protocol_version",
"max_protocol_version",
"oauth_client_id",
] {
assert!(conninfo.contains(key), "{key} must be kept: {conninfo}");
}
assert!(
super::password_free_conninfo("postgres://u@db/app?scram_client_key=AAAA").is_none(),
"a SCRAM key is a credential, not a printable option"
);
}
#[test]
fn the_allowlist_keeps_ordinary_parameters() {
let printed = super::password_free_conninfo(
"postgres://u@db/app?sslmode=verify-full&sslrootcert=%2Fetc%2Fca.pem&connect_timeout=10&target_session_attrs=read-write&application_name=autumn-scrub&hostaddr=10.0.0.5&keepalives_idle=60&tcp_user_timeout=5000&passfile=%2Frun%2Fsecrets%2Fpgpass&options=-c%20search_path%3Dapp",
)
.expect("ordinary libpq parameters must stay printable");
for keyword in [
"sslmode=",
"sslrootcert=",
"connect_timeout=",
"target_session_attrs=",
"application_name=",
"hostaddr=",
"keepalives_idle=",
"tcp_user_timeout=",
"passfile=",
"options=",
] {
assert!(
printed.contains(keyword),
"a known-non-secret parameter must survive: {printed}"
);
}
}
#[test]
fn keyword_matching_is_case_insensitive() {
let printed = super::password_free_conninfo(
"postgres://u@db/app?PASSWORD=s3cr3t&HOST=db.internal&AppLICATION_Name=x",
)
.expect("a URI must stay printable");
assert!(
!printed.contains("s3cr3t"),
"PASSWORD must strip like password: {printed}"
);
assert!(
printed.contains("HOST=db.internal"),
"HOST must survive like host: {printed}"
);
assert!(
printed.contains("AppLICATION_Name=x"),
"the key keeps its own spelling: {printed}"
);
}
#[test]
fn the_printed_sample_reproves_its_postconditions_after_refresh() {
use super::sample::{
ForeignKeyConstraint, SampleAmount, SampleInputs, SampleRules, SampleSpec, build_plan,
};
use std::collections::{BTreeMap, BTreeSet};
let tables = vec![
("users".to_owned(), vec!["id".to_owned()]),
("orders".to_owned(), vec!["id".to_owned()]),
];
let keys = vec![ForeignKeyConstraint {
name: "orders_user_fk".to_owned(),
child_table: "orders".to_owned(),
child_columns: vec!["user_id".to_owned()],
parent_table: "users".to_owned(),
parent_columns: vec!["id".to_owned()],
match_full: false,
}];
let roots = vec![SampleSpec {
table: "users".to_owned(),
amount: SampleAmount::Count(10),
}];
let rules = SampleRules::default();
let plan = build_plan(&SampleInputs {
roots: &roots,
seed: 7,
rules: &rules,
tables: &tables,
foreign_keys: &keys,
framework_tables: &BTreeSet::new(),
purged: &BTreeSet::new(),
partitions: &BTreeMap::new(),
})
.expect("the fixture plan must build");
let capture = super::sample_count_capture_statements(&plan);
assert_eq!(capture.len(), 1, "one capture statement: {capture:?}");
let capture = &capture[0];
assert!(
capture.starts_with(&format!(
"CREATE TEMPORARY TABLE {} ON COMMIT DROP AS ",
super::SAMPLE_COUNTS_TABLE
)),
"the counts must live in a temp table, not psql variables — psql does not \
interpolate variables inside dollar-quoted DO blocks: {capture}"
);
for table in &plan.tables {
assert!(
capture.contains(&format!(
"(SELECT count(*) FROM {})",
super::qualified_ident(&table.table)
)),
"every sampled table must be captured: {capture}"
);
}
let postconditions = super::sample_postcondition_statements(&plan);
let integrity: Vec<String> = plan
.integrity_statements()
.iter()
.map(|(_, statement)| super::integrity_assertion(statement))
.collect();
assert_eq!(
postconditions.len(),
1 + integrity.len(),
"one count check plus the same FK statements the sample printed: {postconditions:?}"
);
let (block, comment) = &postconditions[0];
assert!(
block.contains("RAISE EXCEPTION")
&& block.contains(&format!("pg_temp.{}", super::SAMPLE_COUNTS_TABLE)),
"the count check must abort the transaction on mismatch, reading the captured \
counts rather than print-time values: {block}"
);
for table in &plan.tables {
assert!(
block.contains(&super::qualified_ident(&table.table)),
"every sampled table must be re-counted: {block}"
);
}
assert!(
comment.contains("after the refreshes"),
"the comment must say when it runs: {comment}"
);
for (index, expected) in integrity.iter().enumerate() {
let (statement, comment) = &postconditions[index + 1];
assert_eq!(
statement, expected,
"the FK half must re-run the printed statements, not new ones"
);
assert!(
comment.contains("after the refreshes"),
"and say so: {comment}"
);
}
}
#[test]
fn an_unsampled_target_prints_no_postconditions() {
use super::sample::{SampleInputs, SampleRules, build_plan};
use std::collections::{BTreeMap, BTreeSet};
let rules = SampleRules::default();
let empty = build_plan(&SampleInputs {
roots: &[],
seed: 7,
rules: &rules,
tables: &[],
foreign_keys: &[],
framework_tables: &BTreeSet::new(),
purged: &BTreeSet::new(),
partitions: &BTreeMap::new(),
})
.expect("an empty plan must build");
assert!(empty.tables.is_empty());
assert!(super::sample_count_capture_statements(&empty).is_empty());
assert!(super::sample_postcondition_statements(&empty).is_empty());
}
#[test]
fn the_session_pins_cover_every_rendering_the_row_key_depends_on() {
let pinned = super::session_settings().join("; ");
for setting in [
"search_path",
"standard_conforming_strings",
"DateStyle",
"IntervalStyle",
"TimeZone",
"bytea_output",
"extra_float_digits",
"lc_monetary",
] {
assert!(
pinned.contains(setting),
"{setting} changes how a value renders, so it must be pinned: {pinned}"
);
}
assert!(
super::session_settings()
.iter()
.all(|s| s.starts_with("SET LOCAL ")),
"and every pin must be transaction-scoped: {pinned}"
);
}
#[test]
fn rules_are_asked_for_by_the_same_rules_but_without_the_walk() {
let sql = rules_reaching("AND r.ev_type = '4' ");
assert!(
!sql.contains("pg_inherits"),
"a rule fires only on the relation named, so no walk: {sql}"
);
assert!(
sql.contains("r.ev_enabled IN ('O', 'A')"),
"a disabled rule cannot fire either: {sql}"
);
assert!(
sql.contains("r.rulename <> '_RETURN'"),
"every view's own SELECT rule must be excluded: {sql}"
);
assert!(
sql.contains("AND r.ev_type = '4'"),
"the caller's event filter must reach the query: {sql}"
);
assert!(
!rules_reaching("").contains("ev_type"),
"and an empty filter must not smuggle one in"
);
}
#[test]
fn only_row_triggers_propagate_up_the_inheritance_tree() {
let sql = triggers_reaching("AND (t.tgtype & 8) <> 0 ");
assert!(
sql.contains("(t.tgtype & 1) <> 0 AS by_row"),
"the seed must record whether each trigger is row-level: {sql}"
);
assert!(
sql.contains("JOIN ancestry a ON a.oid = i.inhrelid AND a.by_row"),
"and only row-level ones may propagate to an ancestor: {sql}"
);
assert!(
sql.contains("t.tgenabled IN ('O', 'A')"),
"a disabled trigger must not count: {sql}"
);
assert!(
sql.contains("AND (t.tgtype & 8) <> 0"),
"the caller's event filter must reach the seed: {sql}"
);
assert!(
!triggers_reaching("").contains("tgtype & 8"),
"and an empty filter must not smuggle one in"
);
}
#[test]
fn an_assertion_delimiter_cannot_be_closed_by_an_identifier() {
let check = r#"SELECT count(*) FROM "public"."us$$ers""#;
let sql = integrity_assertion(check);
assert!(
!sql.starts_with("DO $$ "),
"the delimiter must come from the body, not a constant: {sql}"
);
assert!(sql.contains(check), "the check itself must survive: {sql}");
let tag = sql
.split_whitespace()
.nth(1)
.expect("the block must open with `DO <tag>`");
assert_eq!(
sql.matches(tag).count(),
2,
"the tag must appear exactly twice — opening and closing: {sql}"
);
assert!(sql.ends_with(tag), "and must close the block: {sql}");
}
#[test]
fn an_emptiness_assertion_delimiter_widens_the_same_way() {
let sql = emptiness_assertion("jobs$autumn_walk$queue");
let tag = sql
.split_whitespace()
.nth(1)
.expect("the block must open with `DO <tag>`");
assert_ne!(
tag, "$autumn_walk$",
"a table name carrying the default tag must push it wider: {sql}"
);
assert_eq!(sql.matches(tag).count(), 2, "opened and closed once: {sql}");
}
use std::collections::{BTreeMap, BTreeSet};
use autumn_schema_core::{Backend, Column, ColumnType, ForeignKey, Index, Table};
use super::*;
fn text_col(name: &str) -> Column {
Column::new(name, ColumnType::Text)
}
fn pk_col(name: &str) -> Column {
let mut c = Column::new(name, ColumnType::Int64);
c.primary_key = true;
c
}
fn users_table() -> Table {
let mut t = Table::new("users", Backend::Postgres);
t.primary_key = vec!["id".to_owned()];
let mut email = text_col("email");
email.unique = true;
let mut bio = text_col("bio");
bio.nullable = true;
t.columns = vec![
pk_col("id"),
email,
text_col("full_name"),
bio,
Column::new("created_at", ColumnType::Timestamp),
];
t
}
fn empty_encrypted() -> BTreeMap<String, BTreeMap<String, bool>> {
BTreeMap::new()
}
fn encrypted_columns_of(
table: &str,
columns: &[&str],
) -> BTreeMap<String, BTreeMap<String, bool>> {
BTreeMap::from([(
table.to_owned(),
columns.iter().map(|c| ((*c).to_owned(), false)).collect(),
)])
}
fn no_anonymize() -> BTreeSet<String> {
BTreeSet::new()
}
fn plan_for(
tables: &[Table],
config: &ScrubConfig,
encrypted: &BTreeMap<String, BTreeMap<String, bool>>,
anonymize: &BTreeSet<String>,
) -> Result<ScrubPlan, ScrubError> {
plan_with_facts(
tables,
config,
encrypted,
anonymize,
&DatabaseFacts::default(),
)
}
fn plan_with_facts(
tables: &[Table],
config: &ScrubConfig,
encrypted: &BTreeMap<String, BTreeMap<String, bool>>,
anonymize: &BTreeSet<String>,
facts: &DatabaseFacts,
) -> Result<ScrubPlan, ScrubError> {
build_plan(&ClassificationInputs {
tables,
config,
encrypted,
anonymize_tables: anonymize,
facts,
})
}
fn sql_of(plan: &ScrubPlan, table: &str) -> String {
plan.tables
.iter()
.find(|t| t.table == table)
.unwrap_or_else(|| panic!("no plan for {table}"))
.sql
.clone()
.unwrap_or_else(|| panic!("{table} has no SQL statement"))
}
#[test]
fn config_parses_the_sample_rules() {
let config = parse_config_str(
r#"
[sample]
always_include = ["countries"]
never_include = ["audit_logs"]
"#,
)
.unwrap();
assert_eq!(config.sample.always_include, vec!["countries".to_owned()]);
assert_eq!(config.sample.never_include, vec!["audit_logs".to_owned()]);
assert!(sources_declare_sampling(&config));
assert!(!sources_declare_sampling(&ScrubConfig::default()));
}
#[test]
fn an_unknown_sample_key_is_refused_rather_than_ignored() {
let err = parse_config_str(
r#"
[sample]
allways_include = ["countries"]
"#,
)
.unwrap_err();
assert!(matches!(err, ScrubError::Config { .. }));
}
#[test]
fn percentages_report_n_a_rather_than_dividing_by_zero() {
assert_eq!(percent_of(0, 0), "n/a");
assert_eq!(percent_of(2, 200), "1.0%");
assert_eq!(percent_of(200, 200), "100.0%");
}
#[test]
fn byte_sizes_read_at_human_scale() {
assert_eq!(human_bytes(0), "0 B");
assert_eq!(human_bytes(512), "512 B");
assert_eq!(human_bytes(1024), "1.0 kB");
assert_eq!(human_bytes(1024 * 1024 * 3), "3.0 MB");
}
#[test]
fn config_parses_defaults_safe_and_pii() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at"]
[tables.users]
safe = ["role"]
[tables.users.pii]
email = "email"
full_name = "name"
"#,
)
.expect("config must parse");
assert_eq!(config.defaults.safe_columns, vec!["id", "created_at"]);
let users = config.tables.get("users").expect("users rule");
assert_eq!(users.safe, vec!["role"]);
assert_eq!(users.pii.get("email"), Some(&Strategy::Email));
assert_eq!(users.pii.get("full_name"), Some(&Strategy::Name));
}
#[test]
fn config_rejects_unknown_strategy() {
let err = parse_config_str(
r#"
[tables.users.pii]
email = "obfuscate"
"#,
)
.expect_err("unknown strategy must be rejected");
assert!(
err.to_string().contains("obfuscate"),
"error should name the bad strategy: {err}"
);
}
#[test]
fn config_rejects_unknown_keys() {
let err = parse_config_str(
r#"
[tables.users]
saf = ["role"]
"#,
)
.expect_err("a typo'd key must not be silently ignored");
assert!(err.to_string().contains("saf"), "error: {err}");
}
#[test]
fn empty_config_parses_to_default() {
assert_eq!(parse_config_str("").unwrap(), ScrubConfig::default());
}
#[test]
fn unclassified_columns_are_refused_and_listed() {
let tables = vec![users_table()];
let err = plan_for(
&tables,
&ScrubConfig::default(),
&empty_encrypted(),
&no_anonymize(),
)
.expect_err("an all-unclassified schema must be refused");
let ScrubError::Unclassified { columns } = &err else {
panic!("expected Unclassified, got {err:?}");
};
assert_eq!(
columns,
&vec![
"users.bio".to_owned(),
"users.created_at".to_owned(),
"users.email".to_owned(),
"users.full_name".to_owned(),
"users.id".to_owned(),
]
);
let rendered = err.to_string();
assert!(rendered.contains("users.email"), "{rendered}");
assert!(rendered.contains(SCRUB_CONFIG_FILE), "{rendered}");
}
#[test]
fn a_newly_added_column_flips_a_previously_clean_config_to_failure() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at"]
[tables.users]
safe = []
[tables.users.pii]
email = "email"
full_name = "name"
bio = "redact"
"#,
)
.unwrap();
let tables = vec![users_table()];
plan_for(&tables, &config, &empty_encrypted(), &no_anonymize())
.expect("the fully-declared schema must pass");
let mut with_new_column = users_table();
with_new_column.columns.push(text_col("ssn"));
let err = plan_for(
&[with_new_column],
&config,
&empty_encrypted(),
&no_anonymize(),
)
.expect_err("a new undeclared column must fail the scrub");
assert!(matches!(
err,
ScrubError::Unclassified { ref columns } if columns == &vec!["users.ssn".to_owned()]
));
}
#[test]
fn stale_config_entries_are_refused() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "email", "full_name", "bio"]
[tables.users.pii]
emial = "email"
"#,
)
.unwrap();
let err = plan_for(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
)
.expect_err("a config naming a column that no longer exists must fail");
assert!(
matches!(err, ScrubError::StaleConfig { ref entries } if entries.contains(&"users.emial".to_owned())),
"got {err:?}"
);
}
#[test]
fn stale_config_table_is_refused() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "email", "full_name", "bio"]
[tables.legacy_users]
safe = ["x"]
"#,
)
.unwrap();
let err = plan_for(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
)
.expect_err("a config table that no longer exists must fail");
assert!(
matches!(err, ScrubError::StaleConfig { ref entries } if entries.contains(&"legacy_users".to_owned())),
"got {err:?}"
);
}
#[test]
fn a_column_cannot_be_both_safe_and_pii() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "full_name", "bio"]
[tables.users]
safe = ["email"]
[tables.users.pii]
email = "email"
"#,
)
.unwrap();
let err = plan_for(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
)
.expect_err("a contradictory declaration must fail");
assert!(matches!(err, ScrubError::Contradiction { .. }), "{err:?}");
}
#[test]
fn encrypted_columns_are_pii_without_any_declaration() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "full_name", "bio"]
[tables.users]
safe = []
"#,
)
.unwrap();
let encrypted = encrypted_columns_of("users", &["email"]);
let plan = plan_for(&[users_table()], &config, &encrypted, &no_anonymize())
.expect("an #[encrypted] column needs no declaration");
let column = plan
.column("users", "email")
.expect("email must be in the plan");
assert_eq!(column.source, ClassSource::Encrypted);
}
#[test]
fn safe_cannot_override_an_encrypted_column() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "full_name", "bio"]
[tables.users]
safe = ["email"]
"#,
)
.unwrap();
let encrypted = encrypted_columns_of("users", &["email"]);
let err = plan_for(&[users_table()], &config, &encrypted, &no_anonymize())
.expect_err("marking an #[encrypted] column safe must be refused");
assert!(
matches!(err, ScrubError::SafeOverridesEncrypted { ref columns } if columns == &vec!["users.email".to_owned()]),
"got {err:?}"
);
}
#[test]
fn gdpr_anonymize_table_classifies_its_columns_as_pii() {
let config = parse_config_str(
r#"
[tables.users]
safe = ["id", "created_at"]
"#,
)
.unwrap();
let anonymize = BTreeSet::from(["users".to_owned()]);
let plan = plan_for(&[users_table()], &config, &empty_encrypted(), &anonymize)
.expect("a GDPR-anonymize table needs no per-column declaration");
for column in ["email", "full_name", "bio"] {
let decision = plan
.column("users", column)
.unwrap_or_else(|| panic!("{column} must be scrubbed"));
assert_eq!(decision.source, ClassSource::GdprAnonymize);
}
assert!(plan.column("users", "id").is_none());
assert!(plan.column("users", "created_at").is_none());
}
#[test]
fn the_global_safe_list_may_not_narrow_a_gdpr_anonymize_table() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "full_name"]
"#,
)
.unwrap();
let anonymize = BTreeSet::from(["users".to_owned()]);
let plan = plan_for(&[users_table()], &config, &empty_encrypted(), &anonymize).unwrap();
assert!(
plan.column("users", "full_name").is_some(),
"a global safe_columns entry must not narrow an anonymize registration"
);
assert!(plan.column("users", "id").is_none());
}
#[test]
fn a_table_the_role_cannot_see_is_a_refusal_not_a_clean_report() {
let facts = DatabaseFacts {
public_base_tables: BTreeSet::from([
"users".to_owned(),
"secrets".to_owned(),
"autumn_jobs".to_owned(),
]),
..DatabaseFacts::default()
};
assert_eq!(
unreachable_tables(&facts, &[users_table()]),
vec!["secrets".to_owned()]
);
let visible = DatabaseFacts {
public_base_tables: BTreeSet::from(["users".to_owned()]),
..DatabaseFacts::default()
};
assert!(unreachable_tables(&visible, &[users_table()]).is_empty());
}
#[test]
fn a_unique_uuid_column_still_gets_a_castable_32_hex_value() {
let column = Column::new("token", ColumnType::Uuid);
let expr = replacement_expr(Strategy::Uuid, &column, "TOK", true).unwrap();
assert_eq!(expr, "(substr(TOK, 1, 32))::uuid");
}
#[test]
fn a_purge_target_under_rls_is_refused_like_any_other_write() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "email", "full_name", "bio"]
[framework]
purge = ["autumn_jobs"]
"#,
)
.unwrap();
let facts = DatabaseFacts {
framework_tables: vec!["autumn_jobs".to_owned()],
rls_tables: BTreeSet::from(["autumn_jobs".to_owned()]),
..DatabaseFacts::default()
};
let plan = plan_with_facts(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
&facts,
)
.unwrap();
assert!(plan.tables.is_empty());
let purges = purge_statements(&facts.framework_tables, &config);
assert_eq!(purges.len(), 1);
assert!(
facts.rls_tables.contains(&purges[0].0),
"a purge target under RLS must reach the refusal"
);
}
#[test]
fn column_level_privilege_gaps_are_refused_too() {
let facts = DatabaseFacts {
public_base_tables: BTreeSet::from(["users".to_owned()]),
public_columns: BTreeSet::from([
("users".to_owned(), "id".to_owned()),
("users".to_owned(), "email".to_owned()),
("users".to_owned(), "ssn".to_owned()),
]),
..DatabaseFacts::default()
};
assert_eq!(
unreachable_tables(&facts, &[users_table()]),
vec!["users.ssn".to_owned()]
);
}
#[test]
fn every_payload_bearing_framework_table_is_listed() {
for table in [
"_autumn_ledger_revisions",
"_autumn_ledger_high_water",
"_autumn_version_history",
"api_tokens",
"autumn_experiment_assignments",
"autumn_experiment_changes",
"autumn_experiment_overrides",
"autumn_feature_flags",
"autumn_jobs",
"autumn_repository_commit_hooks",
"autumn_runtime_config_changes",
"autumn_runtime_config_values",
"autumn_search_documents",
"autumn_sync_rows",
"feature_flag_changes",
] {
assert!(
FRAMEWORK_PAYLOAD_TABLES.contains(&table),
"{table} carries app data the classification never sees"
);
}
}
#[test]
fn encrypted_columns_may_be_declared_when_there_is_no_model_source() {
let config = parse_config_str(
r#"
[tables.users.encrypted]
api_token = "randomized"
email = "deterministic"
"#,
)
.unwrap();
let users = config.tables.get("users").unwrap();
assert_eq!(
users.encrypted.get("email"),
Some(&EncryptionMode::Deterministic),
"the mode cannot be guessed: re-encrypting a deterministic column in \
randomized mode leaves ciphertext the app can no longer equality-query"
);
assert!(!users.encrypted["api_token"].is_deterministic());
}
#[test]
fn the_feature_flag_identity_tables_are_payload_carriers() {
for table in ["feature_flag_changes", "autumn_feature_flags"] {
assert!(
FRAMEWORK_PAYLOAD_TABLES.contains(&table),
"{table} carries actor identities the classification never sees"
);
assert!(is_framework_table(table));
}
}
#[test]
fn a_check_constrained_column_is_refused_rather_than_guessed_at() {
let facts = DatabaseFacts {
checked_columns: BTreeSet::from([("users".to_owned(), "bio".to_owned())]),
..DatabaseFacts::default()
};
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "email", "full_name"]
[tables.users.pii]
bio = "redact"
"#,
)
.unwrap();
let err = plan_with_facts(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
&facts,
)
.expect_err("no fabricated value can be proven to satisfy an arbitrary CHECK");
assert!(
matches!(err, ScrubError::CheckConstrainedColumn { ref column } if column == "users.bio"),
"got {err:?}"
);
}
#[test]
fn the_referenced_side_of_a_foreign_key_is_refused() {
let facts = DatabaseFacts {
foreign_key_columns: BTreeSet::from([("users".to_owned(), "email".to_owned())]),
..DatabaseFacts::default()
};
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "full_name", "bio"]
[tables.users.pii]
email = "email"
"#,
)
.unwrap();
let err = plan_with_facts(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
&facts,
)
.expect_err("rewriting a referenced natural key breaks its children");
assert!(
matches!(err, ScrubError::PiiOnKeyColumn { ref columns } if columns == &vec!["users.email".to_owned()]),
"got {err:?}"
);
}
#[test]
fn a_generated_column_is_never_rewritten() {
let facts = DatabaseFacts {
generated_columns: BTreeSet::from([("users".to_owned(), "full_name".to_owned())]),
..DatabaseFacts::default()
};
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "email", "bio"]
"#,
)
.unwrap();
let plan = plan_with_facts(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
&facts,
)
.expect("a generated column needs no declaration");
assert!(plan.column("users", "full_name").is_none());
}
#[test]
fn a_partition_is_scrubbed_through_its_parent_not_twice() {
let mut partition = users_table();
partition.name = "users_2026_01".to_owned();
let facts = DatabaseFacts {
partitions: BTreeMap::from([("users_2026_01".to_owned(), "users".to_owned())]),
..DatabaseFacts::default()
};
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at"]
[tables.users.pii]
email = "email"
full_name = "name"
bio = "redact"
"#,
)
.unwrap();
let plan = plan_with_facts(
&[users_table(), partition],
&config,
&empty_encrypted(),
&no_anonymize(),
&facts,
)
.expect("a partition needs no declaration of its own");
assert_eq!(
plan.tables.len(),
1,
"the parent UPDATE already covers the partition's rows"
);
assert_eq!(plan.tables[0].table, "users");
}
#[test]
fn null_is_refused_on_a_nulls_not_distinct_unique_column() {
let mut t = users_table();
t.columns[3].unique = true; let facts = DatabaseFacts {
nulls_not_distinct_columns: BTreeSet::from([("users".to_owned(), "bio".to_owned())]),
..DatabaseFacts::default()
};
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "email", "full_name"]
[tables.users.pii]
bio = "null"
"#,
)
.unwrap();
let err = plan_with_facts(&[t], &config, &empty_encrypted(), &no_anonymize(), &facts)
.expect_err("a second NULL is a violation under NULLS NOT DISTINCT");
assert!(
matches!(err, ScrubError::NonUniqueStrategy { .. }),
"{err:?}"
);
}
#[test]
fn a_unique_column_gets_a_wider_token_and_a_higher_floor() {
let table = users_table();
let unique_token = token_expr(&table, "email", true);
assert_eq!(
unique_token.matches("md5(").count(),
2,
"a unique column needs more entropy than one md5's 32 hex characters: \
{unique_token}"
);
assert!(token_expr(&table, "bio", false).contains("md5("));
let narrow = Column::new(
"code",
ColumnType::Opaque {
pg_type: "varchar(19)".to_owned(),
},
);
assert!(replacement_expr(Strategy::Redact, &narrow, "TOK", false).is_ok());
assert!(matches!(
replacement_expr(Strategy::Redact, &narrow, "TOK", true),
Err(ScrubError::ColumnTooNarrow { .. })
));
}
#[test]
fn a_text_strategy_is_refused_on_a_non_character_column() {
for (name, ty) in [
("age", ColumnType::Int32),
("seen_at", ColumnType::TimestampTz),
(
"ip",
ColumnType::Opaque {
pg_type: "inet".to_owned(),
},
),
] {
let column = Column::new(name, ty);
for strategy in [
Strategy::Email,
Strategy::Name,
Strategy::Redact,
Strategy::Phone,
] {
assert!(
matches!(
replacement_expr(strategy, &column, "TOK", false),
Err(ScrubError::StrategyTypeMismatch { .. })
),
"{strategy:?} on {name} must be refused at plan time, not at apply time"
);
}
}
}
#[test]
fn the_null_assignment_carries_no_untyped_case() {
let mut column = Column::new("token", ColumnType::Uuid);
column.nullable = true;
assert_eq!(
assignment(&column, "NULL", Strategy::Null),
r#""token" = NULL"#
);
assert!(
assignment(&column, "X", Strategy::Uuid).contains("CASE WHEN"),
"every other strategy still preserves NULLs"
);
}
#[test]
fn auto_strategy_covers_the_everyday_opaque_pii_types() {
for (pg_type, expected) in [
("date", Strategy::Epoch),
("time", Strategy::Epoch),
("int2", Strategy::Zero),
("numeric(12,2)", Strategy::Zero),
] {
let column = Column::new(
"value",
ColumnType::Opaque {
pg_type: pg_type.to_owned(),
},
);
assert_eq!(
auto_strategy(&column).unwrap(),
expected,
"`{pg_type}` is an everyday PII column type and needs a usable strategy"
);
assert!(replacement_expr(expected, &column, "TOK", false).is_ok());
}
}
#[test]
fn purge_never_accepts_schema_bookkeeping() {
for table in NEVER_PURGEABLE_TABLES {
let config =
parse_config_str(&format!("[framework]\npurge = [\"{table}\"]\n")).unwrap();
assert!(
matches!(
check_purge_list(&config),
Err(ScrubError::PurgeSchemaBookkeeping { .. })
),
"emptying {table} would make the copy un-migratable or un-routable"
);
}
}
#[test]
fn declaring_a_framework_table_points_at_the_right_mechanism() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "email", "full_name", "bio"]
[tables.api_tokens.pii]
token = "redact"
"#,
)
.unwrap();
let err = plan_for(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
)
.expect_err("a framework table cannot be declared column-by-column");
assert!(
matches!(err, ScrubError::FrameworkTableDeclared { ref tables } if tables == &vec!["api_tokens".to_owned()]),
"got {err:?}"
);
}
#[test]
fn safe_may_narrow_a_gdpr_anonymize_table() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at"]
[tables.users]
safe = ["full_name"]
"#,
)
.unwrap();
let anonymize = BTreeSet::from(["users".to_owned()]);
let plan = plan_for(&[users_table()], &config, &empty_encrypted(), &anonymize).unwrap();
assert!(plan.column("users", "full_name").is_none());
assert!(plan.column("users", "email").is_some());
}
#[test]
fn explicit_pii_wins_over_the_auto_strategy() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "email", "bio"]
[tables.users.pii]
full_name = "redact"
"#,
)
.unwrap();
let plan = plan_for(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
)
.unwrap();
let column = plan.column("users", "full_name").unwrap();
assert_eq!(column.strategy, Strategy::Redact);
assert_eq!(column.source, ClassSource::Config);
}
#[test]
fn a_plaintext_strategy_is_refused_on_an_encrypted_column() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "full_name", "bio"]
[tables.users.pii]
email = "redact"
"#,
)
.unwrap();
let err = plan_for(
&[users_table()],
&config,
&encrypted_columns_of("users", &["email"]),
&no_anonymize(),
)
.expect_err("plaintext must never be written into an #[encrypted] column");
assert!(
matches!(err, ScrubError::PlaintextIntoEncrypted { ref columns } if columns[0].starts_with("users.email")),
"got {err:?}"
);
}
#[test]
fn an_encrypted_column_resolves_to_a_re_encryption_and_carries_its_mode() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "full_name", "bio"]
"#,
)
.unwrap();
let encrypted = BTreeMap::from([(
"users".to_owned(),
BTreeMap::from([("email".to_owned(), true)]),
)]);
let plan = plan_for(&[users_table()], &config, &encrypted, &no_anonymize()).unwrap();
assert_eq!(
plan.column("users", "email").unwrap().strategy,
Strategy::Encrypted
);
let table = &plan.tables[0];
assert!(
table.sql.is_none(),
"an encrypted rewrite is not expressible as SQL: {:?}",
table.sql
);
assert_eq!(table.encrypted.len(), 1);
assert!(
table.encrypted[0].deterministic,
"a deterministic column must be re-encrypted deterministically, or equality \
lookups against it stop matching"
);
assert_eq!(table.encrypted[0].shape, Strategy::Email);
}
#[test]
fn null_may_still_be_declared_on_a_nullable_encrypted_column() {
let mut t = users_table();
t.columns[1].nullable = true;
t.columns[1].unique = false;
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "full_name", "bio"]
[tables.users.pii]
email = "null"
"#,
)
.unwrap();
let plan = plan_for(
&[t],
&config,
&encrypted_columns_of("users", &["email"]),
&no_anonymize(),
)
.expect("NULL is a valid, readable value for an encrypted column");
assert_eq!(
plan.column("users", "email").unwrap().strategy,
Strategy::Null
);
}
#[test]
fn pii_on_a_primary_key_is_refused() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["created_at", "email", "full_name", "bio"]
[tables.users.pii]
id = "zero"
"#,
)
.unwrap();
let err = plan_for(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
)
.expect_err("scrubbing a primary key would break referencing rows");
assert!(
matches!(err, ScrubError::PiiOnKeyColumn { ref columns } if columns == &vec!["users.id".to_owned()]),
"got {err:?}"
);
}
#[test]
fn pii_on_a_foreign_key_is_refused() {
let mut posts = Table::new("posts", Backend::Postgres);
posts.primary_key = vec!["id".to_owned()];
let mut author = Column::new("author_id", ColumnType::Int64);
author.references = Some(ForeignKey::new("users", "id"));
posts.columns = vec![pk_col("id"), author];
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id"]
[tables.posts.pii]
author_id = "zero"
"#,
)
.unwrap();
let err = plan_for(&[posts], &config, &empty_encrypted(), &no_anonymize())
.expect_err("scrubbing a foreign key would break referential integrity");
assert!(
matches!(err, ScrubError::PiiOnKeyColumn { ref columns } if columns == &vec!["posts.author_id".to_owned()]),
"got {err:?}"
);
}
#[test]
fn a_gdpr_anonymize_table_never_auto_classifies_its_key_columns() {
let mut posts = Table::new("posts", Backend::Postgres);
posts.primary_key = vec!["id".to_owned()];
let mut author = Column::new("author_id", ColumnType::Int64);
author.references = Some(ForeignKey::new("users", "id"));
posts.columns = vec![pk_col("id"), author, text_col("body")];
let plan = plan_for(
&[posts],
&ScrubConfig::default(),
&empty_encrypted(),
&BTreeSet::from(["posts".to_owned()]),
)
.expect("key columns are structurally safe under a table-level inference");
assert!(plan.column("posts", "id").is_none());
assert!(plan.column("posts", "author_id").is_none());
assert!(plan.column("posts", "body").is_some());
}
#[test]
fn null_strategy_is_refused_on_a_not_null_column() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "email", "bio"]
[tables.users.pii]
full_name = "null"
"#,
)
.unwrap();
let err = plan_for(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
)
.expect_err("NULL into a NOT NULL column must be refused at plan time");
assert!(matches!(err, ScrubError::NullOnNotNull { .. }), "{err:?}");
}
#[test]
fn null_strategy_is_allowed_on_a_nullable_column() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "email", "full_name"]
[tables.users.pii]
bio = "null"
"#,
)
.unwrap();
let plan = plan_for(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
)
.unwrap();
assert_eq!(
plan.column("users", "bio").unwrap().strategy,
Strategy::Null
);
}
#[test]
fn a_non_injective_strategy_is_refused_on_a_unique_column() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "full_name", "bio"]
[tables.users.pii]
email = "json"
"#,
)
.unwrap();
let err = plan_for(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
)
.expect_err("a constant replacement would violate the unique index");
assert!(
matches!(err, ScrubError::NonUniqueStrategy { .. }),
"got {err:?}"
);
}
#[test]
fn row_key_uses_the_primary_key_when_present() {
assert_eq!(row_key_expr(&users_table()), r#"coalesce("id"::text, '')"#);
}
#[test]
fn row_key_falls_back_to_ctid_without_a_primary_key() {
let mut t = Table::new("legacy", Backend::Postgres);
t.columns = vec![text_col("note")];
assert_eq!(row_key_expr(&t), "ctid::text");
}
#[test]
fn row_key_concatenates_a_composite_primary_key() {
let mut t = Table::new("memberships", Backend::Postgres);
t.primary_key = vec!["user_id".to_owned(), "team_id".to_owned()];
let mut user_id = Column::new("user_id", ColumnType::Int64);
user_id.primary_key = true;
let mut team_id = Column::new("team_id", ColumnType::Int64);
team_id.primary_key = true;
t.columns = vec![user_id, team_id];
assert_eq!(row_key_expr(&t), r#"ROW("user_id", "team_id")::text"#);
}
#[test]
fn token_is_salted_per_column_so_two_columns_never_match() {
let table = users_table();
assert_ne!(
token_expr(&table, "email", false),
token_expr(&table, "full_name", false),
"two PII columns of one row must not receive the same fake value"
);
}
#[test]
fn email_expression_is_unique_per_row_and_uses_a_reserved_domain() {
let expr = replacement_expr(Strategy::Email, &text_col("email"), "TOK", false).unwrap();
assert!(expr.contains("TOK"), "must vary per row: {expr}");
assert!(
expr.contains("@example.invalid"),
"must use a reserved, undeliverable domain: {expr}"
);
}
#[test]
fn varchar_length_bounds_the_generated_value() {
let column = Column::new(
"email",
ColumnType::Opaque {
pg_type: "varchar(40)".to_owned(),
},
);
let expr = replacement_expr(Strategy::Email, &column, "TOK", false).unwrap();
assert!(
expr.contains("substr(TOK, 1, 15)"),
"token must be narrowed to fit varchar(40): {expr}"
);
}
#[test]
fn a_too_narrow_column_is_refused_rather_than_silently_truncated() {
let column = Column::new(
"email",
ColumnType::Opaque {
pg_type: "varchar(28)".to_owned(),
},
);
let err = replacement_expr(Strategy::Email, &column, "TOK", false)
.expect_err("a column too narrow for a unique fake must be refused");
assert!(matches!(err, ScrubError::ColumnTooNarrow { .. }), "{err:?}");
}
#[test]
fn char_length_is_parsed_from_the_opaque_pg_type() {
assert_eq!(
char_max_len(&ColumnType::Opaque {
pg_type: "varchar(64)".to_owned()
}),
Some(64)
);
assert_eq!(
char_max_len(&ColumnType::Opaque {
pg_type: "char(2)".to_owned()
}),
Some(2)
);
assert_eq!(char_max_len(&ColumnType::Text), None);
assert_eq!(
char_max_len(&ColumnType::Opaque {
pg_type: "citext".to_owned()
}),
None
);
}
#[test]
fn auto_strategy_is_derived_from_the_column_type() {
assert_eq!(
auto_strategy(&text_col("email")).unwrap(),
Strategy::Email,
"an email-named text column gets a syntactically valid address"
);
assert_eq!(auto_strategy(&text_col("bio")).unwrap(), Strategy::Redact);
assert_eq!(
auto_strategy(&Column::new("token", ColumnType::Uuid)).unwrap(),
Strategy::Uuid
);
assert_eq!(
auto_strategy(&Column::new("blob", ColumnType::Bytes)).unwrap(),
Strategy::Bytes
);
assert_eq!(
auto_strategy(&Column::new("meta", ColumnType::Json)).unwrap(),
Strategy::Json
);
assert_eq!(
auto_strategy(&Column::new("age", ColumnType::Int32)).unwrap(),
Strategy::Zero
);
assert_eq!(
auto_strategy(&Column::new("seen_at", ColumnType::TimestampTz)).unwrap(),
Strategy::Epoch
);
}
#[test]
fn auto_strategy_refuses_to_guess_for_a_closed_set_or_exotic_type() {
assert!(matches!(
auto_strategy(&Column::new(
"status",
ColumnType::Enum {
variants: vec!["draft".to_owned()]
}
)),
Err(ScrubError::NoAutoStrategy { .. })
));
assert!(matches!(
auto_strategy(&Column::new(
"addr",
ColumnType::Opaque {
pg_type: "inet".to_owned()
}
)),
Err(ScrubError::NoAutoStrategy { .. })
));
}
#[test]
fn update_preserves_nulls_and_batches_a_table_into_one_statement() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at"]
[tables.users.pii]
email = "email"
full_name = "name"
bio = "redact"
"#,
)
.unwrap();
let plan = plan_for(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
)
.unwrap();
assert_eq!(plan.tables.len(), 1, "one statement per table");
let sql = sql_of(&plan, "users");
assert!(sql.starts_with(r#"UPDATE "public"."users" SET "#), "{sql}");
assert!(
sql.contains(r#""bio" = CASE WHEN "bio" IS NULL THEN NULL ELSE"#),
"a nullable column must keep its NULLs: {sql}"
);
assert!(
!sql.contains(r#""full_name" = CASE"#),
"a NOT NULL column needs no CASE: {sql}"
);
assert!(sql.contains(r#""email" = "#) && sql.contains(r#""full_name" = "#));
}
#[test]
fn a_table_with_no_pii_produces_no_statement() {
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "email", "full_name", "bio"]
"#,
)
.unwrap();
let plan = plan_for(
&[users_table()],
&config,
&empty_encrypted(),
&no_anonymize(),
)
.unwrap();
assert!(plan.tables.is_empty(), "nothing to scrub, nothing emitted");
}
#[test]
fn identifiers_with_quotes_are_escaped_in_the_statement() {
let mut t = Table::new(r#"we"ird"#, Backend::Postgres);
t.primary_key = vec!["id".to_owned()];
t.columns = vec![pk_col("id"), text_col(r#"na"me"#)];
let mut config = ScrubConfig::default();
config.defaults.safe_columns = vec!["id".to_owned()];
config.tables.insert(
r#"we"ird"#.to_owned(),
TableRule {
safe: Vec::new(),
pii: BTreeMap::from([(r#"na"me"#.to_owned(), Strategy::Redact)]),
encrypted: BTreeMap::new(),
},
);
let plan = plan_for(&[t], &config, &empty_encrypted(), &no_anonymize()).unwrap();
let sql = sql_of(&plan, r#"we"ird"#);
assert!(
sql.starts_with(r#"UPDATE "public"."we""ird" SET "na""me" = "#),
"{sql}"
);
}
#[test]
fn anonymize_registrations_are_extracted_from_source() {
let src = r#"
use autumn_web::gdpr::{GdprRegistry, ModelRegistration};
fn registry() -> GdprRegistry {
GdprRegistry::new()
.register(ModelRegistration::hard_delete("posts"))
.register(ModelRegistration::anonymize("comments"))
.register(ModelRegistration::retain("invoices", "legal hold"))
.register(autumn_web::gdpr::ModelRegistration::anonymize("profiles"))
}
"#;
let found = extract_anonymize_tables(src).expect("source must parse");
assert_eq!(
found,
BTreeSet::from(["comments".to_owned(), "profiles".to_owned()])
);
}
#[test]
fn a_commented_out_registration_is_not_extracted() {
let src = r#"
fn registry() {
// ModelRegistration::anonymize("ghosts")
let _ = ModelRegistration::anonymize("comments");
}
"#;
let found = extract_anonymize_tables(src).unwrap();
assert_eq!(found, BTreeSet::from(["comments".to_owned()]));
}
#[test]
fn a_non_literal_registration_argument_is_reported_not_ignored() {
let src = "
fn registry() {
let _ = ModelRegistration::anonymize(table_name());
}
";
let err = extract_anonymize_tables(src)
.expect_err("a table name the scanner cannot resolve must not pass silently");
assert!(
matches!(err, ScrubError::UnresolvableAnonymize { .. }),
"{err:?}"
);
}
#[test]
fn scrub_refuses_a_production_profile_without_force() {
for profile in ["prod", "production", "staging"] {
assert!(
matches!(
guard_scrub_target(profile, false),
Err(ScrubError::ProductionRefused { .. })
),
"{profile} must be refused"
);
assert!(guard_scrub_target(profile, true).is_ok());
}
for profile in ["dev", "development", "test"] {
assert!(guard_scrub_target(profile, false).is_ok());
}
}
#[test]
fn same_database_compares_host_port_and_name_ignoring_credentials() {
assert!(same_database(
"postgres://app:pw@db.example.com:5432/myapp",
"postgres://readonly:other@db.example.com:5432/myapp"
));
assert!(!same_database(
"postgres://app:pw@db.example.com:5432/myapp",
"postgres://app:pw@db.example.com:5432/myapp_staging"
));
assert!(!same_database(
"postgres://app:pw@db.example.com:5432/myapp",
"postgres://app:pw@staging.example.com:5432/myapp"
));
assert!(!same_database("not a url", "not a url"));
}
#[test]
fn a_bare_artifact_with_no_manifest_still_gets_the_source_guard() {
let dir = tempfile::tempdir().unwrap();
std::fs::write(
dir.path().join("autumn-prod.toml"),
"[database]\nprimary_url = \"postgres://app:pw@db.example.com:5432/myapp\"\n",
)
.unwrap();
std::fs::write(dir.path().join("autumn-dev.toml"), "[database]\n").unwrap();
assert_eq!(
profiles_with_config(dir.path()),
vec!["dev".to_owned(), "prod".to_owned()],
"every profile overlay is a candidate, not just the one an artifact names"
);
assert!(
!profiles_with_config(dir.path()).is_empty(),
"the guard has candidates to check even with unknown provenance"
);
}
#[test]
fn the_source_guard_waiver_is_not_force() {
let targets = vec![(
"control".to_owned(),
"postgres://app:pw@db.example.com:5432/myapp".to_owned(),
)];
let empty = tempfile::tempdir().unwrap();
assert!(
guard_configured_source(Some("prod"), empty.path(), &targets, false).is_ok(),
"no config overlay means nothing to compare against"
);
assert!(
guard_configured_source(Some("prod"), empty.path(), &targets, true).is_ok(),
"--allow-source-overwrite is the waiver"
);
}
#[test]
fn errors_never_leak_credentials() {
let err = ScrubError::ProductionRefused {
profile: "prod".to_owned(),
};
let rendered = err.to_string();
assert!(rendered.contains("prod"));
assert!(!rendered.contains("postgres://"));
assert!(!rendered.contains("hunter2"));
}
#[test]
fn check_report_prints_a_paste_ready_stanza_for_unclassified_columns() {
let stanza = suggested_config_stanza(&[
"users.email".to_owned(),
"users.full_name".to_owned(),
"posts.body".to_owned(),
]);
assert!(stanza.contains("[tables.users.pii]"), "{stanza}");
assert!(stanza.contains("email = \"auto\""), "{stanza}");
assert!(stanza.contains("[tables.posts.pii]"), "{stanza}");
assert!(stanza.contains("body = \"auto\""), "{stanza}");
}
#[test]
fn per_column_errors_are_reported_with_their_table() {
let mut t = Table::new("orders", Backend::Postgres);
t.primary_key = vec!["id".to_owned()];
t.columns = vec![
pk_col("id"),
Column::new(
"status",
ColumnType::Enum {
variants: vec!["draft".to_owned(), "paid".to_owned()],
},
),
];
let err = plan_for(
&[t],
&parse_config_str(
r#"
[defaults]
safe_columns = ["id"]
[tables.orders.pii]
status = "auto"
"#,
)
.unwrap(),
&empty_encrypted(),
&no_anonymize(),
)
.expect_err("a closed-set column has no generic fake");
assert!(
matches!(err, ScrubError::NoAutoStrategy { ref column, .. } if column == "orders.status"),
"the error must name the table too: {err:?}"
);
}
#[test]
fn a_strategy_that_cannot_produce_the_column_type_is_refused() {
let err = replacement_expr(Strategy::Uuid, &text_col("note"), "TOK", false)
.expect_err("a uuid cannot be written into a text column");
assert!(
matches!(err, ScrubError::StrategyTypeMismatch { .. }),
"{err:?}"
);
assert!(
matches!(
replacement_expr(Strategy::Zero, &text_col("note"), "TOK", false),
Err(ScrubError::StrategyTypeMismatch { .. })
),
"zero is meaningless for a text column"
);
assert!(
matches!(
replacement_expr(
Strategy::Epoch,
&Column::new("n", ColumnType::Int64),
"TOK",
false
),
Err(ScrubError::StrategyTypeMismatch { .. })
),
"epoch is meaningless for an integer column"
);
}
#[test]
fn typed_strategies_render_their_casts() {
assert_eq!(
replacement_expr(
Strategy::Uuid,
&Column::new("t", ColumnType::Uuid),
"TOK",
false
)
.unwrap(),
"(substr(TOK, 1, 32))::uuid"
);
assert_eq!(
replacement_expr(
Strategy::Bytes,
&Column::new("b", ColumnType::Bytes),
"TOK",
false
)
.unwrap(),
"decode(TOK, 'hex')"
);
assert_eq!(
replacement_expr(
Strategy::Zero,
&Column::new("ok", ColumnType::Bool),
"TOK",
false
)
.unwrap(),
"false"
);
assert_eq!(
replacement_expr(
Strategy::Epoch,
&Column::new("at", ColumnType::TimestampTz),
"TOK",
false
)
.unwrap(),
"'1970-01-01 00:00:00+00'::timestamptz"
);
assert_eq!(
replacement_expr(Strategy::Null, &text_col("bio"), "TOK", false).unwrap(),
"NULL"
);
}
#[test]
fn phone_produces_digits_and_refuses_a_column_that_cannot_hold_them() {
let expr = replacement_expr(Strategy::Phone, &text_col("phone"), "TOK", false).unwrap();
assert!(
expr.contains("translate("),
"must map hex onto digits: {expr}"
);
assert!(expr.starts_with("'+1555'"), "{expr}");
let narrow = Column::new(
"phone",
ColumnType::Opaque {
pg_type: "varchar(8)".to_owned(),
},
);
assert!(matches!(
replacement_expr(Strategy::Phone, &narrow, "TOK", false),
Err(ScrubError::ColumnTooNarrow { .. })
));
}
#[test]
fn a_partial_unique_index_still_constrains_the_rows_it_covers() {
let mut t = users_table();
t.columns[1].unique = false;
let mut index = Index::new("idx_users_email", vec!["email".to_owned()], true);
index.is_partial = true;
t.indexes = vec![index];
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "full_name", "bio"]
[tables.users.pii]
email = "json"
"#,
)
.unwrap();
let err = plan_for(&[t], &config, &empty_encrypted(), &no_anonymize())
.expect_err("a partial unique index still constrains the rows it covers");
assert!(
matches!(err, ScrubError::NonUniqueStrategy { .. }),
"{err:?}"
);
}
#[test]
fn a_composite_unique_index_constrains_each_of_its_members() {
let mut t = Table::new("cards", Backend::Postgres);
t.primary_key = vec!["id".to_owned()];
t.columns = vec![
pk_col("id"),
Column::new("user_id", ColumnType::Int64),
Column::new("last4", ColumnType::Int32),
];
let mut config = ScrubConfig::default();
config.defaults.safe_columns = vec!["id".to_owned(), "user_id".to_owned()];
config.tables.insert(
"cards".to_owned(),
TableRule {
safe: Vec::new(),
pii: BTreeMap::from([("last4".to_owned(), Strategy::Zero)]),
encrypted: BTreeMap::new(),
},
);
let facts = DatabaseFacts {
unique_columns: BTreeSet::from([
("cards".to_owned(), "user_id".to_owned()),
("cards".to_owned(), "last4".to_owned()),
]),
..DatabaseFacts::default()
};
let err = plan_with_facts(&[t], &config, &empty_encrypted(), &no_anonymize(), &facts)
.expect_err("a constant in one member of a composite unique key collides");
assert!(
matches!(err, ScrubError::NonUniqueStrategy { .. }),
"{err:?}"
);
}
#[test]
fn a_missing_config_file_is_not_an_error_but_an_explicit_one_is() {
let dir = tempfile::tempdir().unwrap();
let absent = dir.path().join("scrub.toml");
assert_eq!(
load_config_at(None, &absent).unwrap(),
ScrubConfig::default()
);
assert!(matches!(
load_config_at(Some(&absent), &absent),
Err(ScrubError::Config { .. })
));
}
#[test]
fn purge_only_accepts_framework_owned_tables() {
let config = parse_config_str(
r#"
[framework]
purge = ["autumn_jobs", "users"]
"#,
)
.unwrap();
let err = check_purge_list(&config)
.expect_err("emptying a user table must never hide behind `purge`");
assert!(
matches!(err, ScrubError::PurgeNotFrameworkTable { ref tables } if tables == &vec!["users".to_owned()]),
"got {err:?}"
);
let ok = parse_config_str(
r#"
[framework]
purge = ["autumn_jobs", "_autumn_ledger_revisions", "_autumn_ledger_high_water"]
"#,
)
.unwrap();
assert!(check_purge_list(&ok).is_ok());
}
#[test]
fn purging_one_ledger_table_without_the_other_is_refused() {
for (listed, missing) in [
("_autumn_ledger_revisions", "_autumn_ledger_high_water"),
("_autumn_ledger_high_water", "_autumn_ledger_revisions"),
] {
let config = parse_config_str(&format!(
"[framework]\npurge = [\"autumn_jobs\", \"{listed}\"]\n"
))
.unwrap();
let err = check_purge_list(&config)
.expect_err("emptying one ledger table alone must be refused");
assert!(
matches!(
err,
ScrubError::PurgeLedgerTablesUnpaired {
listed: ref got_listed,
missing: ref got_missing,
} if got_listed == listed && got_missing == missing
),
"got {err:?}"
);
assert!(err.to_string().contains(missing), "{err}");
}
for purge in [
r#"["autumn_jobs"]"#,
r#"["_autumn_ledger_revisions", "_autumn_ledger_high_water"]"#,
] {
let ok = parse_config_str(&format!("[framework]\npurge = {purge}\n")).unwrap();
assert!(check_purge_list(&ok).is_ok(), "{purge}");
}
}
#[test]
fn purge_statements_cover_only_the_opted_in_tables_that_exist() {
let config = parse_config_str(
r#"
[framework]
purge = ["autumn_jobs", "autumn_sync_rows"]
"#,
)
.unwrap();
let present = vec!["autumn_job_tracking".to_owned(), "autumn_jobs".to_owned()];
let statements = purge_statements(&present, &config);
assert_eq!(
statements,
vec![(
"autumn_jobs".to_owned(),
r#"DELETE FROM "public"."autumn_jobs""#.to_owned()
)]
);
}
#[test]
fn no_purge_declaration_empties_nothing() {
let present = vec!["autumn_jobs".to_owned()];
assert!(purge_statements(&present, &ScrubConfig::default()).is_empty());
}
#[test]
fn the_probe_covers_the_built_in_list_plus_every_purge_entry() {
let config = parse_config_str(
r#"
[framework]
purge = ["autumn_custom_outbox"]
"#,
)
.unwrap();
let probed = probe_table_names(&config);
assert!(
probed.contains("autumn_custom_outbox"),
"a purge entry outside the built-in list must still be probed, or it \
would be accepted and then silently do nothing"
);
for table in FRAMEWORK_PAYLOAD_TABLES {
assert!(probed.contains(*table));
}
}
#[test]
fn framework_payload_tables_are_all_framework_owned() {
for table in FRAMEWORK_PAYLOAD_TABLES {
assert!(
is_framework_table(table),
"{table} must be filtered out of the classified universe"
);
}
for table in UNPREFIXED_FRAMEWORK_TABLES {
assert!(is_framework_table(table));
}
assert!(
FRAMEWORK_PAYLOAD_TABLES.contains(&"api_tokens"),
"production API tokens in a staging copy are a live credential leak"
);
assert!(!FRAMEWORK_PAYLOAD_TABLES.contains(&"autumn_migration_checksums"));
assert!(!FRAMEWORK_PAYLOAD_TABLES.contains(&"_autumn_shard_map"));
}
#[test]
fn index_backed_uniqueness_is_recognized() {
let mut t = users_table();
t.columns[1].unique = false;
t.indexes = vec![Index::new(
"idx_users_email",
vec!["email".to_owned()],
true,
)];
let config = parse_config_str(
r#"
[defaults]
safe_columns = ["id", "created_at", "full_name", "bio"]
[tables.users.pii]
email = "json"
"#,
)
.unwrap();
let err = plan_for(&[t], &config, &empty_encrypted(), &no_anonymize())
.expect_err("a unique index must be honored like a unique column");
assert!(
matches!(err, ScrubError::NonUniqueStrategy { .. }),
"{err:?}"
);
}
}