1use std::fmt;
14
15use sha2::{Digest, Sha256};
16
17pub use headgate_sql::PostgresNamespace;
18
19mod mysql;
20mod postgres;
21mod schema;
22
23pub use mysql::{
24 DEFAULT_MYSQL_LOCK_NAMESPACE, MysqlValidation, adopt_mysql, adopt_mysql_with_lock_namespace,
25 applied_mysql, migrate_mysql, migrate_mysql_with_lock_namespace, mysql_migration_lock_name,
26 validate_mysql,
27};
28pub use postgres::{
29 PostgresValidation, adopt_postgres, adopt_postgres_in_schema, applied_postgres,
30 applied_postgres_in_schema, migrate_postgres, migrate_postgres_in_schema, validate_postgres,
31 validate_postgres_in_schema,
32};
33
34#[derive(Clone, Copy, Debug, Eq, PartialEq)]
37pub enum Backend {
38 Postgres,
39 Mysql,
40}
41
42impl Backend {
43 pub const fn as_str(self) -> &'static str {
44 match self {
45 Self::Postgres => "postgres",
46 Self::Mysql => "mysql",
47 }
48 }
49}
50
51impl fmt::Display for Backend {
52 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53 f.write_str(self.as_str())
54 }
55}
56
57#[derive(Clone, Copy, Debug, Eq, PartialEq)]
58pub enum Direction {
59 Up,
60 Down,
61}
62
63impl Direction {
64 pub const fn as_str(self) -> &'static str {
65 match self {
66 Self::Up => "up",
67 Self::Down => "down",
68 }
69 }
70}
71
72impl fmt::Display for Direction {
73 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
74 f.write_str(self.as_str())
75 }
76}
77
78#[derive(Clone, Copy, Debug)]
81pub struct Migration {
82 pub version: u32,
83 pub name: &'static str,
84 pub up_sql: &'static str,
85 pub down_sql: &'static str,
86 pub online_safe: bool,
89}
90
91const POSTGRES_MIGRATIONS: &[Migration] = &[
92 Migration {
93 version: 1,
94 name: "initial_schema",
95 up_sql: include_str!("../migrations/postgres/0001_init.up.sql"),
96 down_sql: include_str!("../migrations/postgres/0001_init.down.sql"),
97 online_safe: false,
100 },
101 Migration {
102 version: 2,
103 name: "enqueue_backpressure",
104 up_sql: include_str!("../migrations/postgres/0002_enqueue_backpressure.up.sql"),
105 down_sql: include_str!("../migrations/postgres/0002_enqueue_backpressure.down.sql"),
106 online_safe: false,
108 },
109 Migration {
110 version: 3,
111 name: "job_results",
112 up_sql: include_str!("../migrations/postgres/0003_job_results.up.sql"),
113 down_sql: include_str!("../migrations/postgres/0003_job_results.down.sql"),
114 online_safe: true,
115 },
116 Migration {
117 version: 4,
118 name: "mid_run_output",
119 up_sql: include_str!("../migrations/postgres/0004_mid_run_output.up.sql"),
120 down_sql: include_str!("../migrations/postgres/0004_mid_run_output.down.sql"),
121 online_safe: true,
122 },
123 Migration {
124 version: 5,
125 name: "job_progress",
126 up_sql: include_str!("../migrations/postgres/0005_job_progress.up.sql"),
127 down_sql: include_str!("../migrations/postgres/0005_job_progress.down.sql"),
128 online_safe: true,
129 },
130 Migration {
131 version: 6,
132 name: "periodic_origin",
133 up_sql: include_str!("../migrations/postgres/0006_periodic_origin.up.sql"),
134 down_sql: include_str!("../migrations/postgres/0006_periodic_origin.down.sql"),
135 online_safe: true,
136 },
137 Migration {
138 version: 7,
139 name: "scheduler_events",
140 up_sql: include_str!("../migrations/postgres/0007_scheduler_events.up.sql"),
141 down_sql: include_str!("../migrations/postgres/0007_scheduler_events.down.sql"),
142 online_safe: true,
143 },
144 Migration {
145 version: 8,
146 name: "pending_state",
147 up_sql: include_str!("../migrations/postgres/0008_pending_tags_metrics.up.sql"),
148 down_sql: include_str!("../migrations/postgres/0008_pending_tags_metrics.down.sql"),
149 online_safe: false,
151 },
152 Migration {
153 version: 9,
154 name: "pending_tags_metrics",
155 up_sql: include_str!("../migrations/postgres/0009_pending_tags_metrics.up.sql"),
156 down_sql: include_str!("../migrations/postgres/0009_pending_tags_metrics.down.sql"),
157 online_safe: false,
158 },
159 Migration {
160 version: 10,
161 name: "sticky_routing",
162 up_sql: include_str!("../migrations/postgres/0010_sticky_routing.up.sql"),
163 down_sql: include_str!("../migrations/postgres/0010_sticky_routing.down.sql"),
164 online_safe: false,
165 },
166 Migration {
167 version: 11,
168 name: "partitioned_archive",
169 up_sql: include_str!("../migrations/postgres/0011_partitioned_archive.up.sql"),
170 down_sql: include_str!("../migrations/postgres/0011_partitioned_archive.down.sql"),
171 online_safe: true,
172 },
173 Migration {
174 version: 12,
175 name: "worker_control_state",
176 up_sql: include_str!("../migrations/postgres/0012_worker_control_state.up.sql"),
177 down_sql: include_str!("../migrations/postgres/0012_worker_control_state.down.sql"),
178 online_safe: true,
179 },
180];
181
182const MYSQL_MIGRATIONS: &[Migration] = &[
183 Migration {
184 version: 1,
185 name: "initial_schema",
186 up_sql: include_str!("../migrations/mysql/0001_init.up.sql"),
187 down_sql: include_str!("../migrations/mysql/0001_init.down.sql"),
188 online_safe: false,
189 },
190 Migration {
191 version: 2,
192 name: "enqueue_backpressure",
193 up_sql: include_str!("../migrations/mysql/0002_enqueue_backpressure.up.sql"),
194 down_sql: include_str!("../migrations/mysql/0002_enqueue_backpressure.down.sql"),
195 online_safe: false,
196 },
197 Migration {
198 version: 3,
199 name: "job_results",
200 up_sql: include_str!("../migrations/mysql/0003_job_results.up.sql"),
201 down_sql: include_str!("../migrations/mysql/0003_job_results.down.sql"),
202 online_safe: true,
203 },
204 Migration {
205 version: 4,
206 name: "mid_run_output",
207 up_sql: include_str!("../migrations/mysql/0004_mid_run_output.up.sql"),
208 down_sql: include_str!("../migrations/mysql/0004_mid_run_output.down.sql"),
209 online_safe: true,
210 },
211 Migration {
212 version: 5,
213 name: "job_progress",
214 up_sql: include_str!("../migrations/mysql/0005_job_progress.up.sql"),
215 down_sql: include_str!("../migrations/mysql/0005_job_progress.down.sql"),
216 online_safe: true,
217 },
218 Migration {
219 version: 6,
220 name: "periodic_origin",
221 up_sql: include_str!("../migrations/mysql/0006_periodic_origin.up.sql"),
222 down_sql: include_str!("../migrations/mysql/0006_periodic_origin.down.sql"),
223 online_safe: true,
224 },
225 Migration {
226 version: 7,
227 name: "scheduler_events",
228 up_sql: include_str!("../migrations/mysql/0007_scheduler_events.up.sql"),
229 down_sql: include_str!("../migrations/mysql/0007_scheduler_events.down.sql"),
230 online_safe: true,
231 },
232 Migration {
233 version: 8,
234 name: "pending_state_barrier",
235 up_sql: include_str!("../migrations/mysql/0008_pending_tags_metrics.up.sql"),
236 down_sql: include_str!("../migrations/mysql/0008_pending_tags_metrics.down.sql"),
237 online_safe: false,
238 },
239 Migration {
240 version: 9,
241 name: "pending_tags_metrics",
242 up_sql: include_str!("../migrations/mysql/0009_pending_tags_metrics.up.sql"),
243 down_sql: include_str!("../migrations/mysql/0009_pending_tags_metrics.down.sql"),
244 online_safe: false,
245 },
246 Migration {
247 version: 10,
248 name: "sticky_routing",
249 up_sql: include_str!("../migrations/mysql/0010_sticky_routing.up.sql"),
250 down_sql: include_str!("../migrations/mysql/0010_sticky_routing.down.sql"),
251 online_safe: false,
252 },
253 Migration {
254 version: 11,
255 name: "partitioned_archive",
256 up_sql: include_str!("../migrations/mysql/0011_partitioned_archive.up.sql"),
257 down_sql: include_str!("../migrations/mysql/0011_partitioned_archive.down.sql"),
258 online_safe: false,
259 },
260 Migration {
261 version: 12,
262 name: "worker_control_state",
263 up_sql: include_str!("../migrations/mysql/0012_worker_control_state.up.sql"),
264 down_sql: include_str!("../migrations/mysql/0012_worker_control_state.down.sql"),
265 online_safe: false,
266 },
267];
268
269pub const fn migrations(backend: Backend) -> &'static [Migration] {
270 match backend {
271 Backend::Postgres => POSTGRES_MIGRATIONS,
272 Backend::Mysql => MYSQL_MIGRATIONS,
273 }
274}
275
276pub fn migration(backend: Backend, version: u32) -> Option<&'static Migration> {
277 migrations(backend).iter().find(|m| m.version == version)
278}
279
280pub fn latest_version(backend: Backend) -> u32 {
281 migrations(backend).last().map_or(0, |m| m.version)
282}
283
284pub fn checksum(migration: &Migration) -> String {
288 let digest = Sha256::digest(migration.up_sql.as_bytes());
289 format!("{digest:x}")
290}
291
292#[derive(Clone, Debug, Eq, PartialEq)]
293pub struct AppliedMigration {
294 pub version: u32,
295 pub name: String,
296 pub checksum: String,
297 pub applied_at_ms: i64,
298}
299
300#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
301pub struct MigrateOptions {
302 pub target_version: Option<u32>,
305 pub max_steps: Option<usize>,
307 pub dry_run: bool,
309}
310
311#[derive(Clone, Debug)]
312pub struct MigrationStep {
313 pub direction: Direction,
314 pub migration: &'static Migration,
315}
316
317impl PartialEq for MigrationStep {
318 fn eq(&self, other: &Self) -> bool {
319 self.direction == other.direction && self.migration.version == other.migration.version
320 }
321}
322
323impl Eq for MigrationStep {}
324
325#[derive(Clone, Debug, Default)]
326pub struct MigrateResult {
327 pub dry_run: bool,
328 pub steps: Vec<MigrationStep>,
329}
330
331#[derive(Clone, Copy, Debug, Eq, PartialEq)]
332pub enum InstallationState {
333 Empty,
334 Unversioned,
335 Versioned,
336}
337
338#[derive(Debug)]
339pub enum MigrationError {
340 Invalid(String),
341 UnversionedSchema,
342 History(String),
343 Schema(Vec<String>),
344 Postgres(tokio_postgres::Error),
345 Mysql(mysql_async::Error),
346}
347
348impl MigrationError {
349 pub(crate) fn schema(messages: Vec<String>) -> Self {
350 Self::Schema(messages)
351 }
352}
353
354impl fmt::Display for MigrationError {
355 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
356 match self {
357 Self::Invalid(message) => write!(f, "invalid migration request: {message}"),
358 Self::UnversionedSchema => write!(
359 f,
360 "headgate tables exist without migration history; validate and adopt the current schema before migrating"
361 ),
362 Self::History(message) => write!(f, "invalid migration history: {message}"),
363 Self::Schema(messages) => {
364 write!(f, "schema validation failed: {}", messages.join("; "))
365 }
366 Self::Postgres(error) => write!(f, "postgres migration failed: {error}"),
367 Self::Mysql(error) => write!(f, "mysql migration failed: {error}"),
368 }
369 }
370}
371
372impl std::error::Error for MigrationError {
373 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
374 match self {
375 Self::Postgres(error) => Some(error),
376 Self::Mysql(error) => Some(error),
377 _ => None,
378 }
379 }
380}
381
382impl From<tokio_postgres::Error> for MigrationError {
383 fn from(value: tokio_postgres::Error) -> Self {
384 Self::Postgres(value)
385 }
386}
387
388impl From<mysql_async::Error> for MigrationError {
389 fn from(value: mysql_async::Error) -> Self {
390 Self::Mysql(value)
391 }
392}
393
394pub fn validate_history(
397 backend: Backend,
398 applied: &[AppliedMigration],
399) -> Result<(), MigrationError> {
400 let known = migrations(backend);
401 for (index, row) in applied.iter().enumerate() {
402 let expected_version = (index + 1) as u32;
403 if row.version != expected_version {
404 return Err(MigrationError::History(format!(
405 "versions must be contiguous from 1; found {} where {} was expected",
406 row.version, expected_version
407 )));
408 }
409 let Some(migration) = known.iter().find(|m| m.version == row.version) else {
410 return Err(MigrationError::History(format!(
411 "database has unknown future version {}",
412 row.version
413 )));
414 };
415 if row.name != migration.name {
416 return Err(MigrationError::History(format!(
417 "version {} is named {:?}, expected {:?}",
418 row.version, row.name, migration.name
419 )));
420 }
421 let expected_checksum = checksum(migration);
422 if row.checksum != expected_checksum {
423 return Err(MigrationError::History(format!(
424 "version {} checksum is {}, expected {}",
425 row.version, row.checksum, expected_checksum
426 )));
427 }
428 }
429 Ok(())
430}
431
432pub fn plan(
433 backend: Backend,
434 applied: &[AppliedMigration],
435 direction: Direction,
436 options: MigrateOptions,
437) -> Result<Vec<MigrationStep>, MigrationError> {
438 validate_history(backend, applied)?;
439 let all = migrations(backend);
440 let current = applied.last().map_or(0, |m| m.version);
441 let latest = latest_version(backend);
442 let target = options.target_version.unwrap_or(match direction {
443 Direction::Up => latest,
444 Direction::Down => 0,
445 });
446 if target > latest {
447 return Err(MigrationError::Invalid(format!(
448 "target version {target} is newer than embedded latest version {latest}"
449 )));
450 }
451 match direction {
452 Direction::Up if target < current => {
453 return Err(MigrationError::Invalid(format!(
454 "target version {target} is below current version {current}; use down"
455 )));
456 }
457 Direction::Down if target > current => {
458 return Err(MigrationError::Invalid(format!(
459 "target version {target} is above current version {current}; use up"
460 )));
461 }
462 _ => {}
463 }
464
465 let mut steps: Vec<_> = match direction {
466 Direction::Up => all
467 .iter()
468 .filter(|m| m.version > current && m.version <= target)
469 .map(|migration| MigrationStep {
470 direction,
471 migration,
472 })
473 .collect(),
474 Direction::Down => all
475 .iter()
476 .rev()
477 .filter(|m| m.version > target && m.version <= current)
478 .map(|migration| MigrationStep {
479 direction,
480 migration,
481 })
482 .collect(),
483 };
484 if let Some(max_steps) = options.max_steps {
485 steps.truncate(max_steps);
486 }
487 Ok(steps)
488}
489
490#[cfg(test)]
491mod tests {
492 use super::*;
493
494 fn applied(version: u32) -> AppliedMigration {
495 let migration = migration(Backend::Postgres, version).unwrap();
496 AppliedMigration {
497 version,
498 name: migration.name.to_owned(),
499 checksum: checksum(migration),
500 applied_at_ms: 1,
501 }
502 }
503
504 #[test]
505 fn plans_up_down_targets_and_idempotent_current() {
506 let up = plan(
507 Backend::Postgres,
508 &[],
509 Direction::Up,
510 MigrateOptions::default(),
511 )
512 .unwrap();
513 assert_eq!(
514 up.iter().map(|s| s.migration.version).collect::<Vec<_>>(),
515 [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
516 );
517
518 let current = [
519 applied(1),
520 applied(2),
521 applied(3),
522 applied(4),
523 applied(5),
524 applied(6),
525 applied(7),
526 applied(8),
527 applied(9),
528 applied(10),
529 applied(11),
530 applied(12),
531 ];
532 assert!(
533 plan(
534 Backend::Postgres,
535 ¤t,
536 Direction::Up,
537 MigrateOptions::default()
538 )
539 .unwrap()
540 .is_empty()
541 );
542 let down = plan(
543 Backend::Postgres,
544 ¤t,
545 Direction::Down,
546 MigrateOptions::default(),
547 )
548 .unwrap();
549 assert_eq!(
550 down.iter().map(|s| s.migration.version).collect::<Vec<_>>(),
551 [12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1]
552 );
553 }
554
555 #[test]
556 fn checksum_or_gap_in_history_is_fatal_even_to_planning() {
557 let mut bad = applied(1);
558 bad.checksum = "tampered".into();
559 assert!(matches!(
560 plan(
561 Backend::Postgres,
562 &[bad],
563 Direction::Up,
564 MigrateOptions {
565 dry_run: true,
566 ..MigrateOptions::default()
567 }
568 ),
569 Err(MigrationError::History(_))
570 ));
571
572 let future = AppliedMigration {
573 version: 11,
574 name: "future".into(),
575 checksum: "x".into(),
576 applied_at_ms: 1,
577 };
578 assert!(matches!(
579 validate_history(Backend::Postgres, &[future]),
580 Err(MigrationError::History(_))
581 ));
582 }
583
584 #[test]
585 fn wrong_direction_and_future_targets_are_rejected() {
586 let current = [applied(1)];
587 assert!(matches!(
588 plan(
589 Backend::Postgres,
590 ¤t,
591 Direction::Up,
592 MigrateOptions {
593 target_version: Some(0),
594 ..MigrateOptions::default()
595 }
596 ),
597 Err(MigrationError::Invalid(_))
598 ));
599 assert!(matches!(
600 plan(
601 Backend::Postgres,
602 &[],
603 Direction::Up,
604 MigrateOptions {
605 target_version: Some(13),
606 ..MigrateOptions::default()
607 }
608 ),
609 Err(MigrationError::Invalid(_))
610 ));
611 }
612}