Skip to main content

headgate_migrate/
lib.rs

1//! Versioned, embedded schema migrations for headgate's SQL backends.
2//!
3//! The migration history is data, not a guess made from whichever columns happen to
4//! exist. Every applied version records the SHA-256 of its immutable up SQL. A changed
5//! historical migration therefore fails validation instead of silently turning two
6//! installations at "version 1" into different schemas.
7//!
8//! Postgres applies each version and its history row in one transaction. MySQL DDL
9//! commits implicitly, so its migrations must be resumable: a connection-scoped lock
10//! serializes migrators, every statement in an up migration is idempotent, and the
11//! version row is written only after the resulting schema passes the current manifest.
12
13use 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/// The two stores with durable schemas. Redis key layouts are versioned by code and Lua,
35/// not by a DDL migrator, so claiming a Redis migration backend would be dishonest.
36#[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/// A checked-in migration. Existing versions are immutable: add a new version instead
79/// of editing an applied one, even when the edit appears additive.
80#[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    /// Whether the UP direction may run while workers and clients continue using the
87    /// schema. DOWN is always operator-confirmed and offline.
88    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        // This is a fresh install containing non-concurrent index creation. There is no
98        // prior application traffic it needs to preserve, so call it offline explicitly.
99        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        // Baseline + trigger installation is one cut-over; stop producers first.
107        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        // Forward-only enum extension: schedule an operator-reviewed maintenance window.
150        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
284/// The checksum stored in `headgate_schema_migration`. It covers the UP SQL because that
285/// is the schema an applied version claims was installed; changing DOWN SQL is caught by
286/// source parity tests and review, while it cannot make an existing schema differ.
287pub 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    /// Desired schema version after the call. `None` means latest for UP and zero for
303    /// DOWN. Migrating to the current version is an idempotent no-op.
304    pub target_version: Option<u32>,
305    /// Bound work in one invocation. `None` means all required versions.
306    pub max_steps: Option<usize>,
307    /// Plan and return exact SQL without changing either schema or history.
308    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
394/// Validate history independently of a database. This is also the planner's first step,
395/// so a checksum mismatch cannot be bypassed with `--dry-run` or a target version.
396pub 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                &current,
536                Direction::Up,
537                MigrateOptions::default()
538            )
539            .unwrap()
540            .is_empty()
541        );
542        let down = plan(
543            Backend::Postgres,
544            &current,
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                &current,
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}