use pgsafe::{lint_sql, Finding, Fix, LintOptions};
fn apply(sql: &str, fix: &Fix) -> String {
let mut out = sql.to_string();
let mut edits = fix.edits.clone();
edits.sort_by_key(|e| std::cmp::Reverse(e.start));
for e in edits {
out.replace_range(e.start as usize..e.end as usize, &e.replacement);
}
out
}
fn fix_for(sql: &str, rule: &str) -> (Vec<Finding>, Option<Fix>) {
let fs = lint_sql(sql, &LintOptions::default()).unwrap();
let fix = fs
.iter()
.find(|f| f.rule_id == rule)
.and_then(|f| f.fix.clone());
(fs, fix)
}
fn assert_clears(sql: &str, rule: &str) {
let (_, fix) = fix_for(sql, rule);
let fix = fix.unwrap_or_else(|| panic!("{rule} produced no fix for: {sql}"));
let fixed = apply(sql, &fix);
let after = lint_sql(&fixed, &LintOptions::default())
.unwrap_or_else(|e| panic!("fixed SQL did not parse: {fixed}\n{e}"));
assert!(
after.iter().all(|f| f.rule_id != rule),
"fix did not clear {rule}: {fixed}"
);
}
#[test]
fn add_index_fix_clears_and_parses() {
assert_clears("CREATE INDEX idx ON t (col);", "add-index-non-concurrent");
}
#[test]
fn require_timeout_fix_clears_and_parses() {
assert_clears("ALTER TABLE t ADD COLUMN c int;", "require-timeout");
}
#[test]
fn add_index_fix_clears_unique_index() {
assert_clears(
"CREATE UNIQUE INDEX idx ON t (col);",
"add-index-non-concurrent",
);
}
#[test]
fn require_timeout_first_finding_has_fix_subsequent_do_not() {
let sql = "CREATE INDEX i ON t (a);\nDROP TABLE other;";
let fs = lint_sql(sql, &LintOptions::default()).unwrap();
let timeout_findings: Vec<_> = fs
.iter()
.filter(|f| f.rule_id == "require-timeout")
.collect();
assert!(
timeout_findings.len() >= 2,
"expected ≥2 require-timeout findings, got {}",
timeout_findings.len()
);
assert!(
timeout_findings[0].fix.is_some(),
"first require-timeout finding must carry the fix"
);
assert!(
timeout_findings[1].fix.is_none(),
"subsequent require-timeout findings must not carry a fix"
);
}
#[test]
fn require_timeout_single_fix_clears_all_findings() {
let sql = "CREATE INDEX i ON t (a);\nDROP TABLE other;";
let fs = lint_sql(sql, &LintOptions::default()).unwrap();
let fix = fs
.iter()
.find(|f| f.rule_id == "require-timeout")
.and_then(|f| f.fix.clone())
.expect("require-timeout fix present on first finding");
let fixed = apply(sql, &fix);
let after = lint_sql(&fixed, &LintOptions::default())
.unwrap_or_else(|e| panic!("fixed SQL did not parse: {e}"));
assert!(
after.iter().all(|f| f.rule_id != "require-timeout"),
"single fix should clear all require-timeout findings; fixed SQL:\n{fixed}"
);
}
#[test]
fn require_timeout_fix_inserts_before_first_flagged_statement() {
let sql = "SELECT 1;\nCREATE INDEX i ON t (a);";
let (_, fix) = fix_for(sql, "require-timeout");
let fix = fix.expect("require-timeout fix present");
let first_edit = fix
.edits
.first()
.expect("timeout fix has at least one edit");
assert!(
first_edit.start > 0,
"prologue should precede the flagged statement, not byte 0"
);
let at = first_edit.start as usize;
assert!(
sql[at..].starts_with("CREATE INDEX"),
"prologue anchored at: {:?}",
&sql[at..]
);
}
#[test]
fn require_timeout_prologue_lands_above_leading_directive() {
let sql = "-- pgsafe:ignore add-index-non-concurrent built in a maintenance window\n\
CREATE INDEX idx_users_email ON users (email);";
let fs = lint_sql(sql, &LintOptions::default()).unwrap();
let rt_fix = fs
.iter()
.find(|f| f.rule_id == "require-timeout")
.and_then(|f| f.fix.clone())
.expect("require-timeout finding must carry a fix");
let fixed = apply(sql, &rt_fix);
let after = lint_sql(&fixed, &LintOptions::default())
.unwrap_or_else(|e| panic!("fixed SQL did not parse:\n{fixed}\n{e}"));
assert!(
after.iter().all(|f| f.rule_id != "require-timeout"),
"require-timeout must be cleared after the fix; fixed SQL:\n{fixed}"
);
assert!(
after.iter().all(|f| f.rule_id != "suppression-unused"),
"suppression-unused must not fire; the directive must still attach to the CREATE INDEX; \
fixed SQL:\n{fixed}"
);
let ainc = after
.iter()
.find(|f| f.rule_id == "add-index-non-concurrent")
.expect("add-index-non-concurrent must still be present (suppressed)");
assert!(
ainc.is_suppressed(),
"add-index-non-concurrent must remain suppressed; fixed SQL:\n{fixed}"
);
}
#[test]
fn require_timeout_prologue_lands_above_leading_directive_non_first_stmt() {
let sql = "CREATE TABLE t (data json); -- pgsafe:ignore prefer-jsonb ok\n\
-- pgsafe:ignore add-index-non-concurrent reason\n\
CREATE INDEX i ON existing (col);";
let fs = lint_sql(sql, &LintOptions::default()).unwrap();
let rt_fix = fs
.iter()
.find(|f| f.rule_id == "require-timeout")
.and_then(|f| f.fix.clone())
.expect("require-timeout finding must carry a fix");
let fixed = apply(sql, &rt_fix);
let after = lint_sql(&fixed, &LintOptions::default())
.unwrap_or_else(|e| panic!("fixed SQL did not parse:\n{fixed}\n{e}"));
assert!(
after.iter().all(|f| f.rule_id != "require-timeout"),
"require-timeout must be cleared after the fix; fixed SQL:\n{fixed}"
);
assert!(
after.iter().all(|f| f.rule_id != "suppression-unused"),
"suppression-unused must not fire; fixed SQL:\n{fixed}"
);
let ainc = after
.iter()
.find(|f| f.rule_id == "add-index-non-concurrent")
.expect("add-index-non-concurrent must still be present (suppressed)");
assert!(
ainc.is_suppressed(),
"add-index-non-concurrent must remain suppressed; fixed SQL:\n{fixed}"
);
let pjsonb = after
.iter()
.find(|f| f.rule_id == "prefer-jsonb")
.expect("prefer-jsonb must still be present (suppressed)");
assert!(
pjsonb.is_suppressed(),
"prefer-jsonb must remain suppressed; fixed SQL:\n{fixed}"
);
}
#[test]
fn require_timeout_prologue_lands_above_multiline_block_comment_directive() {
let sql = "CREATE TABLE t (data json); -- pgsafe:ignore prefer-jsonb ok\n\
/* pgsafe:ignore add-index-non-concurrent\n\
\x20 reason */\n\
CREATE INDEX i ON existing (col);";
let fs = lint_sql(sql, &LintOptions::default()).unwrap();
let rt_fix = fs
.iter()
.find(|f| f.rule_id == "require-timeout")
.and_then(|f| f.fix.clone())
.expect("require-timeout finding must carry a fix");
let fixed = apply(sql, &rt_fix);
let after = lint_sql(&fixed, &LintOptions::default())
.unwrap_or_else(|e| panic!("fixed SQL did not parse:\n{fixed}\n{e}"));
assert!(
after.iter().all(|f| f.rule_id != "require-timeout"),
"require-timeout must be cleared after the fix; fixed SQL:\n{fixed}"
);
assert!(
after.iter().all(|f| f.rule_id != "suppression-unused"),
"suppression-unused must not fire; fixed SQL:\n{fixed}"
);
let ainc = after
.iter()
.find(|f| f.rule_id == "add-index-non-concurrent")
.expect("add-index-non-concurrent must still be present (suppressed)");
assert!(
ainc.is_suppressed(),
"add-index-non-concurrent must remain suppressed; fixed SQL:\n{fixed}"
);
}
#[test]
fn drop_index_fix_clears() {
assert_clears("DROP INDEX idx;", "drop-index-non-concurrent");
}
#[test]
fn reindex_fix_clears() {
assert_clears("REINDEX INDEX idx;", "reindex-non-concurrent");
}
#[test]
fn detach_partition_fix_clears() {
assert_clears(
"ALTER TABLE p DETACH PARTITION p1;",
"detach-partition-non-concurrent",
);
}
#[test]
fn add_check_fix_clears() {
assert_clears(
"ALTER TABLE t ADD CONSTRAINT ck CHECK (a > 0);",
"add-check-without-not-valid",
);
}
#[test]
fn add_fk_fix_clears() {
assert_clears(
"ALTER TABLE t ADD CONSTRAINT fk FOREIGN KEY (a) REFERENCES u (id);",
"add-fk-without-not-valid",
);
}
#[test]
fn prefer_jsonb_fix_clears() {
assert_clears("CREATE TABLE t (data json);", "prefer-jsonb");
}
#[test]
fn prefer_bigint_fix_clears() {
assert_clears(
"CREATE TABLE t (id integer PRIMARY KEY);",
"prefer-bigint-primary-key",
);
}