H-3599, H-3600: Implement failing on divergent or missing files (#5683)

hashintel · Nov 21, 2024 · 349a618 · 349a618
1 parent 86b2f77
commit 349a618
Show file tree

Hide file tree

Showing 8 changed files with 311 additions and 153 deletions.
diff --git a/libs/@local/graph/migrations-macros/src/lib.rs b/libs/@local/graph/migrations-macros/src/lib.rs
@@ -191,65 +191,7 @@ struct EmbedMigrationsInput {
     location: PathBuf,
 }
 
-fn migration_list() -> TokenStream {
-    quote! {
-        struct Cons<H, T>(H, T);
-        struct Nil;
-
-        macro_rules! tuples {
-            () => { Nil };
-            ($head:expr) => { Cons($head, Nil) };
-            ($head:expr, $($tail:tt)*) => {
-                Cons($head, tuples![$($tail)*])
-            };
-        }
-
-        #[derive(Debug)]
-        pub struct MigrationDefinition<M> {
-            pub migration: M,
-            pub info: MigrationInfo,
-        }
-
-        impl<Head, Tail, C> MigrationList<C> for Cons<MigrationDefinition<Head>, Tail>
-        where
-            Head: Migration,
-            Tail: MigrationList<C>,
-            C: ContextProvider<Head::Context>,
-        {
-            async fn traverse(
-                self,
-                runner: &impl MigrationRunner,
-                context: &mut C,
-                direction: MigrationDirection,
-            ) -> Result<(), Report<MigrationError>> {
-                if direction == MigrationDirection::Down {
-                    self.1.traverse(runner, context, direction).await?;
-                    runner
-                        .run_migration(self.0.migration, &self.0.info, context.provide())
-                        .await
-                } else {
-                    runner
-                        .run_migration(self.0.migration, &self.0.info, context.provide())
-                        .await?;
-                    self.1.traverse(runner, context, direction).await
-                }
-            }
-        }
-
-        // Empty-list implementation.
-        impl<C> MigrationList<C> for Nil {
-            async fn traverse(
-                self,
-                _: &impl MigrationRunner,
-                _: &mut C,
-                _: MigrationDirection,
-            ) -> Result<(), Report<MigrationError>> {
-                Ok(())
-            }
-        }
-    }
-}
-
+#[expect(clippy::too_many_lines)]
 fn embed_migrations_impl(input: &EmbedMigrationsInput) -> Result<TokenStream, Box<dyn Error>> {
     let included_dir = input.location.to_string_lossy();
     let mut migration_files = BTreeMap::new();
@@ -269,9 +211,11 @@ fn embed_migrations_impl(input: &EmbedMigrationsInput) -> Result<TokenStream, Bo
         }
     }
 
-    let mut mod_definitions = Vec::with_capacity(migration_files.len());
-    let mut context_bounds = Vec::with_capacity(migration_files.len());
-    let mut migration_definitions = Vec::with_capacity(migration_files.len());
+    let num_migrations = migration_files.len();
+    let mut mod_definitions = Vec::with_capacity(num_migrations);
+    let mut context_bounds = Vec::with_capacity(num_migrations);
+    let mut migration_infos = Vec::with_capacity(num_migrations);
+    let mut migration_definitions = Vec::with_capacity(num_migrations);
 
     for (n, (number, migration_file)) in migration_files.into_iter().enumerate() {
         let expected_migration_number = u32::try_from(n + 1)?;
@@ -302,21 +246,23 @@ fn embed_migrations_impl(input: &EmbedMigrationsInput) -> Result<TokenStream, Bo
             });
         }
 
+        migration_infos.push(quote! {
+            MigrationInfo {
+                number: #number,
+                name: Cow::Borrowed(#name),
+                size: #size,
+                digest: Digest::new([#(#digest,)*]),
+            }
+        });
+
         migration_definitions.push(quote! {
             MigrationDefinition {
                 migration: self::#module_name::#struct_name,
-                info: ::hash_graph_migrations::MigrationInfo {
-                    number: #number,
-                    name: #name.into(),
-                    size: #size,
-                    digest: [#(#digest,)*].into(),
-                }
+                index: #n,
             }
         });
     }
 
-    let migration_list_impl = migration_list();
-
     Ok(quote! {
         // This enforces rebuilding the crate when the migrations change.
         const _: ::hash_graph_migrations::__export::Dir = ::hash_graph_migrations::__export::include_dir!(#included_dir);
@@ -329,11 +275,74 @@ fn embed_migrations_impl(input: &EmbedMigrationsInput) -> Result<TokenStream, Bo
         where #(#context_bounds)*
         {
             use ::hash_graph_migrations::{
-                __export::Report, ContextProvider, Migration, MigrationDirection,
-                MigrationError, MigrationInfo, MigrationList, MigrationRunner,
+                __export::{Report, alloc::borrow::Cow},
+                ContextProvider, Digest, Migration, MigrationDirection, MigrationError, MigrationInfo,
+                MigrationList, MigrationRunner,
             };
 
-            #migration_list_impl
+            const MIGRATIONS: [MigrationInfo; #num_migrations] = [
+                #(#migration_infos,)*
+            ];
+
+            struct Cons<H, T>(H, T);
+            struct Nil;
+
+            macro_rules! tuples {
+                () => { Nil };
+                ($head:expr) => { Cons($head, Nil) };
+                ($head:expr, $($tail:tt)*) => {
+                    Cons($head, tuples![$($tail)*])
+                };
+            }
+
+            #[derive(Debug)]
+            pub struct MigrationDefinition<M> {
+                pub migration: M,
+                pub index: usize,
+            }
+
+            impl<Head, Tail, C> MigrationList<C> for Cons<MigrationDefinition<Head>, Tail>
+            where
+                Head: Migration,
+                Tail: MigrationList<C>,
+                C: ContextProvider<Head::Context>,
+            {
+                fn infos(&self) -> impl Iterator<Item=&MigrationInfo> {
+                    MIGRATIONS.iter()
+                }
+
+                async fn traverse(
+                    self,
+                    runner: &impl MigrationRunner,
+                    context: &mut C,
+                    direction: MigrationDirection,
+                ) -> Result<(), Report<MigrationError>> {
+                    if direction == MigrationDirection::Down {
+                        self.1.traverse(runner, context, direction).await?;
+                        runner.run_migration(self.0.migration, &MIGRATIONS[self.0.index], context.provide()).await
+                    } else {
+                        runner.run_migration(self.0.migration, &MIGRATIONS[self.0.index], context.provide()).await?;
+                        self.1.traverse(runner, context, direction).await
+                    }
+                }
+            }
+
+            // Empty-list implementation.
+            impl<C> MigrationList<C> for Nil {
+                fn infos(&self) -> impl Iterator<Item=&MigrationInfo> {
+                    MIGRATIONS.iter()
+                }
+
+                async fn traverse(
+                    self,
+                    _: &impl MigrationRunner,
+                    _: &mut C,
+                    _: MigrationDirection,
+                ) -> Result<(), Report<MigrationError>> {
+                    Ok(())
+                }
+            }
+
             tuples![#(#migration_definitions,)*]
         }
     })

diff --git a/libs/@local/graph/migrations/src/bin/cli/subcommand/run.rs b/libs/@local/graph/migrations/src/bin/cli/subcommand/run.rs
@@ -7,12 +7,33 @@ use crate::{Command, subcommand::DatabaseConnectionInfo};
 
 #[derive(Debug, Parser)]
 #[clap(version, author, about, long_about = None)]
+#[expect(clippy::struct_excessive_bools, reason = "This is a CLI command")]
 pub struct RunCommand {
     #[clap(flatten)]
     pub db_info: DatabaseConnectionInfo,
 
+    /// Run the migration up to the specified target.
+    ///
+    /// If the target is not specified, all migrations will be run. If the target is lower than the
+    /// current migration, the migrations will be run in the down direction.
     #[clap(long)]
     pub target: Option<u32>,
+
+    /// Allows the migration to run even if the files are divergent from the database.
+    #[clap(long)]
+    pub allow_divergent: bool,
+
+    /// Divergent migrations will be updated in the database.
+    #[clap(long, requires("allow_divergent"))]
+    pub update_divergent: bool,
+
+    /// Allows the migration to run even if the files are missing from the file system.
+    #[clap(long)]
+    pub allow_missing: bool,
+
+    /// Missing migrations will be removed from the database.
+    #[clap(long, requires("allow_missing"))]
+    pub remove_missing: bool,
 }
 
 hash_graph_migrations::embed_migrations!("graph-migrations");
@@ -29,7 +50,11 @@ impl Command for RunCommand {
 
         let mut builder = MigrationPlanBuilder::new()
             .state(client)
-            .migrations(self::migrations());
+            .migrations(self::migrations())
+            .allow_divergent(self.allow_divergent)
+            .update_divergent(self.update_divergent)
+            .allow_missing(self.allow_missing)
+            .remove_missing(self.remove_missing);
 
         if let Some(target) = self.target {
             builder = builder.target(target);

diff --git a/libs/@local/graph/migrations/src/info.rs b/libs/@local/graph/migrations/src/info.rs
@@ -18,8 +18,9 @@ pub enum InvalidMigrationFile {
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 pub struct Digest([u8; 32]);
 
-impl From<[u8; 32]> for Digest {
-    fn from(bytes: [u8; 32]) -> Self {
+impl Digest {
+    #[must_use]
+    pub const fn new(bytes: [u8; 32]) -> Self {
         Self(bytes)
     }
 }

diff --git a/libs/@local/graph/migrations/src/lib.rs b/libs/@local/graph/migrations/src/lib.rs
@@ -26,6 +26,7 @@ mod state;
 #[cfg(feature = "macros")]
 #[doc(hidden)]
 pub mod __export {
+    pub extern crate alloc;
     // `Error-stack` is required for the return value of the `up` and `down` methods in the
     // `Migration` trait.
     pub use error_stack::Report;

diff --git a/libs/@local/graph/migrations/src/list.rs b/libs/@local/graph/migrations/src/list.rs
@@ -13,6 +13,9 @@ impl MigrationError {
 }
 
 pub trait MigrationList<C> {
+    /// Returns an iterator over the migration infos in this list.
+    fn infos(&self) -> impl Iterator<Item = &MigrationInfo>;
+
     /// Runs this list in the provided plan and context.
     ///
     /// This will only forward the [`Migration`] to the [`MigrationRunner`].