Skip to main content

ferro_migration/
backfill.rs

1//! Backend-portable backfill helpers for SeaORM migrations.
2
3use crate::error::Error;
4use sea_orm::{DbBackend, DbErr, Statement};
5use sea_orm_migration::prelude::*;
6
7/// Backfill a column with random hex strings of `hex_len` characters.
8///
9/// `hex_len` must be even (each random byte produces two hex digits). Odd
10/// values return an error.
11///
12/// SQLite uses `lower(hex(randomblob(N/2)))`; Postgres uses
13/// `encode(gen_random_bytes(N/2), 'hex')` (requires `pgcrypto`; see README).
14/// MySQL is not supported.
15///
16/// Only rows where the column IS NULL or empty string are updated.
17pub async fn backfill_random_hex(
18    manager: &SchemaManager<'_>,
19    table: &str,
20    column: &str,
21    hex_len: u32,
22) -> Result<(), DbErr> {
23    let backend = manager.get_database_backend();
24    let sql = sql_for_random_hex(backend, table, column, hex_len)?;
25    manager
26        .get_connection()
27        .execute(Statement::from_string(backend, sql))
28        .await
29        .map(|_| ())
30}
31
32pub(crate) fn sql_for_random_hex(
33    backend: DbBackend,
34    table: &str,
35    column: &str,
36    hex_len: u32,
37) -> Result<String, Error> {
38    if hex_len % 2 != 0 {
39        return Err(Error::UnsupportedBackend(format!(
40            "hex_len must be even, got {hex_len}"
41        )));
42    }
43    let byte_len = hex_len / 2;
44    match backend {
45        DbBackend::Sqlite => Ok(format!(
46            "UPDATE \"{table}\" SET \"{column}\" = lower(hex(randomblob({byte_len}))) \
47             WHERE \"{column}\" IS NULL OR \"{column}\" = ''"
48        )),
49        DbBackend::Postgres => Ok(format!(
50            "UPDATE \"{table}\" SET \"{column}\" = encode(gen_random_bytes({byte_len}), 'hex') \
51             WHERE \"{column}\" IS NULL OR \"{column}\" = ''"
52        )),
53        DbBackend::MySql => Err(Error::UnsupportedBackend(
54            "backfill_random_hex: MySQL not supported".into(),
55        )),
56    }
57}
58
59/// Backfill a column with random UUID v4 strings.
60///
61/// SQLite emits a UUID-shaped hex string assembled from `randomblob`.
62/// Postgres uses `gen_random_uuid()::text` (Postgres 13+, core distribution).
63///
64/// Only rows where the column IS NULL or empty string are updated.
65pub async fn backfill_random_uuid(
66    manager: &SchemaManager<'_>,
67    table: &str,
68    column: &str,
69) -> Result<(), DbErr> {
70    let backend = manager.get_database_backend();
71    let sql = sql_for_random_uuid(backend, table, column)?;
72    manager
73        .get_connection()
74        .execute(Statement::from_string(backend, sql))
75        .await
76        .map(|_| ())
77}
78
79pub(crate) fn sql_for_random_uuid(
80    backend: DbBackend,
81    table: &str,
82    column: &str,
83) -> Result<String, Error> {
84    match backend {
85        DbBackend::Sqlite => Ok(format!(
86            "UPDATE \"{table}\" SET \"{column}\" = \
87             lower(hex(randomblob(4))) || '-' || \
88             lower(hex(randomblob(2))) || '-4' || \
89             substr(lower(hex(randomblob(2))), 2) || '-' || \
90             substr('89ab', abs(random()) % 4 + 1, 1) || \
91             substr(lower(hex(randomblob(2))), 2) || '-' || \
92             lower(hex(randomblob(6))) \
93             WHERE \"{column}\" IS NULL OR \"{column}\" = ''"
94        )),
95        DbBackend::Postgres => Ok(format!(
96            "UPDATE \"{table}\" SET \"{column}\" = gen_random_uuid()::text \
97             WHERE \"{column}\" IS NULL OR \"{column}\" = ''"
98        )),
99        DbBackend::MySql => Err(Error::UnsupportedBackend(
100            "backfill_random_uuid: MySQL not supported".into(),
101        )),
102    }
103}
104
105/// Backfill a column with the current timestamp.
106///
107/// SQLite uses `CURRENT_TIMESTAMP`; Postgres uses `now()`.
108/// MySQL is not supported.
109///
110/// Only rows where the column IS NULL are updated.
111pub async fn backfill_current_timestamp(
112    manager: &SchemaManager<'_>,
113    table: &str,
114    column: &str,
115) -> Result<(), DbErr> {
116    let backend = manager.get_database_backend();
117    let sql = sql_for_current_timestamp(backend, table, column)?;
118    manager
119        .get_connection()
120        .execute(Statement::from_string(backend, sql))
121        .await
122        .map(|_| ())
123}
124
125pub(crate) fn sql_for_current_timestamp(
126    backend: DbBackend,
127    table: &str,
128    column: &str,
129) -> Result<String, Error> {
130    match backend {
131        DbBackend::Sqlite => Ok(format!(
132            "UPDATE \"{table}\" SET \"{column}\" = CURRENT_TIMESTAMP \
133             WHERE \"{column}\" IS NULL"
134        )),
135        DbBackend::Postgres => Ok(format!(
136            "UPDATE \"{table}\" SET \"{column}\" = now() \
137             WHERE \"{column}\" IS NULL"
138        )),
139        DbBackend::MySql => Err(Error::UnsupportedBackend(
140            "backfill_current_timestamp: MySQL not supported".into(),
141        )),
142    }
143}
144
145/// General-purpose backfill escape hatch: caller provides a closure that
146/// produces backend-specific SQL.
147///
148/// The closure receives the current `DbBackend` and must return the SQL string
149/// to execute, or `Err(DbErr)` if the backend is unsupported.
150pub async fn backfill<F>(manager: &SchemaManager<'_>, sql_fn: F) -> Result<(), DbErr>
151where
152    F: FnOnce(DbBackend) -> Result<String, DbErr>,
153{
154    let backend = manager.get_database_backend();
155    let sql = sql_fn(backend)?;
156    manager
157        .get_connection()
158        .execute(Statement::from_string(backend, sql))
159        .await
160        .map(|_| ())
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    #[test]
168    fn random_hex_sqlite_emits_randomblob() {
169        let sql = sql_for_random_hex(DbBackend::Sqlite, "bookings", "checkin_token", 16).unwrap();
170        assert!(sql.contains("lower(hex(randomblob(8)))"));
171        assert!(sql.contains("\"bookings\""));
172        assert!(sql.contains("\"checkin_token\""));
173        assert!(sql.contains("IS NULL OR \"checkin_token\" = ''"));
174    }
175
176    #[test]
177    fn random_hex_postgres_emits_gen_random_bytes() {
178        let sql = sql_for_random_hex(DbBackend::Postgres, "bookings", "checkin_token", 16).unwrap();
179        assert!(sql.contains("encode(gen_random_bytes(8), 'hex')"));
180        assert!(sql.contains("\"bookings\""));
181    }
182
183    #[test]
184    fn random_hex_mysql_returns_unsupported() {
185        let err = sql_for_random_hex(DbBackend::MySql, "t", "c", 16).unwrap_err();
186        assert!(matches!(err, Error::UnsupportedBackend(ref msg) if msg.contains("MySQL")));
187    }
188
189    #[test]
190    fn random_hex_odd_hex_len_returns_error() {
191        let err = sql_for_random_hex(DbBackend::Sqlite, "t", "c", 5).unwrap_err();
192        assert!(matches!(err, Error::UnsupportedBackend(ref msg) if msg.contains("even")));
193        let err = sql_for_random_hex(DbBackend::Postgres, "t", "c", 5).unwrap_err();
194        assert!(matches!(err, Error::UnsupportedBackend(ref msg) if msg.contains("even")));
195    }
196
197    #[test]
198    fn random_uuid_sqlite_uses_randomblob_segments() {
199        let sql = sql_for_random_uuid(DbBackend::Sqlite, "users", "external_id").unwrap();
200        assert!(sql.contains("randomblob(4)"));
201        assert!(sql.contains("'-4'"));
202        assert!(sql.contains("substr('89ab'"));
203    }
204
205    #[test]
206    fn random_uuid_postgres_uses_gen_random_uuid() {
207        let sql = sql_for_random_uuid(DbBackend::Postgres, "users", "external_id").unwrap();
208        assert!(sql.contains("gen_random_uuid()::text"));
209    }
210
211    #[test]
212    fn current_timestamp_sqlite_uses_constant() {
213        let sql = sql_for_current_timestamp(DbBackend::Sqlite, "t", "created_at").unwrap();
214        assert!(sql.contains("CURRENT_TIMESTAMP"));
215    }
216
217    #[test]
218    fn current_timestamp_postgres_uses_now() {
219        let sql = sql_for_current_timestamp(DbBackend::Postgres, "t", "created_at").unwrap();
220        assert!(sql.contains("now()"));
221    }
222}