1use crate::error::Error;
4use sea_orm::{DbBackend, DbErr, Statement};
5use sea_orm_migration::prelude::*;
6
7pub 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
59pub 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
105pub 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
145pub 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}