use crate::{Dialect, RelError, RelationalBackend, Value};
use std::fmt;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SweepReport {
pub deleted: u64,
pub batches: u32,
pub complete: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum SweepError {
BadName(String),
ZeroBound,
NullCutoff,
Rel(RelError),
Dialect(String),
}
impl fmt::Display for SweepError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
SweepError::BadName(n) => write!(f, "not a plain identifier: {n:?}"),
SweepError::ZeroBound => f.write_str("batch and max_batches must be at least 1"),
SweepError::NullCutoff => f.write_str("the cutoff must not be NULL"),
SweepError::Rel(e) => write!(f, "sweep: {e}"),
SweepError::Dialect(d) => write!(f, "sweep: no statement for dialect {d}"),
}
}
}
impl std::error::Error for SweepError {}
fn ident(s: &str) -> bool {
!s.is_empty()
&& !s.starts_with(|c: char| c.is_ascii_digit())
&& s.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
}
fn table_name(s: &str) -> bool {
match s.split_once('.') {
Some((schema, table)) => ident(schema) && ident(table),
None => ident(s),
}
}
fn statement(dialect: &Dialect, table: &str, column: &str, batch: u32) -> Option<String> {
Some(match dialect {
Dialect::Postgres => format!(
"DELETE FROM {table} WHERE ctid IN \
(SELECT ctid FROM {table} WHERE {column} < $1 LIMIT {batch})"
),
Dialect::Sqlite => format!(
"DELETE FROM {table} WHERE rowid IN \
(SELECT rowid FROM {table} WHERE {column} < ? LIMIT {batch})"
),
Dialect::MySql | Dialect::MariaDb => {
format!("DELETE FROM {table} WHERE {column} < ? LIMIT {batch}")
}
_ => return None,
})
}
pub async fn sweep<B: RelationalBackend>(
db: &B,
table: &str,
column: &str,
cutoff: &Value,
batch: u32,
max_batches: u32,
) -> Result<SweepReport, SweepError> {
if !table_name(table) {
return Err(SweepError::BadName(table.to_string()));
}
if !ident(column) {
return Err(SweepError::BadName(column.to_string()));
}
if batch == 0 || max_batches == 0 {
return Err(SweepError::ZeroBound);
}
if *cutoff == Value::Null {
return Err(SweepError::NullCutoff);
}
let dialect = db.dialect();
let sql = statement(&dialect, table, column, batch)
.ok_or_else(|| SweepError::Dialect(dialect.to_string()))?;
let mut report = SweepReport {
deleted: 0,
batches: 0,
complete: false,
};
while report.batches < max_batches {
let n = db
.execute(&sql, std::slice::from_ref(cutoff))
.await
.map_err(SweepError::Rel)?;
report.batches += 1;
report.deleted += n;
if n == 0 {
report.complete = true;
break;
}
}
Ok(report)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn names_are_plain_identifiers_only() {
for ok in ["sessions", "_t1", "public.sessions", "A_b"] {
assert!(table_name(ok), "{ok}");
}
for bad in [
"",
"1t",
"t; DROP TABLE t",
"a.b.c",
".t",
"t.",
"a b",
"t\"",
"t-1",
"é",
] {
assert!(!table_name(bad), "{bad:?}");
}
assert!(!ident("a.b"));
}
#[test]
fn each_dialect_has_a_statement_with_one_parameter() {
for d in [
Dialect::Postgres,
Dialect::Sqlite,
Dialect::MySql,
Dialect::MariaDb,
] {
let s = statement(&d, "t", "c", 7).unwrap();
assert!(s.contains("LIMIT 7"), "{s}");
assert_eq!(s.matches("c < ").count(), 1);
}
assert!(statement(&Dialect::Other("x".into()), "t", "c", 1).is_none());
}
}