fix migration primary key generator state bug

Author: Handy-caT

codegen/src/generators/persist/table/impls.rs

@@ -20,6 +20,7 @@ impl PersistGenerator {
         let reinsert_fn = self.gen_table_reinsert_fn();
         let upsert_fn = self.gen_table_upsert_fn();
         let get_next_fn = self.gen_table_get_next_fn();
+        let pk_gen_state_fn = self.gen_table_pk_gen_state_fn();
         let iter_with_fn = self.gen_table_iter_with_fn();
         let iter_with_async_fn = self.gen_table_iter_with_async_fn();
         let count_fn = self.gen_table_count_fn();
@@ -38,6 +39,7 @@ impl PersistGenerator {
                 #upsert_fn
                 #count_fn
                 #get_next_fn
+                #pk_gen_state_fn
                 #iter_with_fn
                 #iter_with_async_fn
                 #system_info_fn
@@ -279,6 +281,17 @@ impl PersistGenerator {
         }
     }
 
+    fn gen_table_pk_gen_state_fn(&self) -> TokenStream {
+        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
+        let primary_key_type = name_generator.get_primary_key_type_ident();
+
+        quote! {
+            pub fn pk_gen_state(&self) -> <<#primary_key_type as TablePrimaryKey>::Generator as PrimaryKeyGeneratorState>::State {
+                self.0.pk_gen.get_state()
+            }
+        }
+    }
+
     fn gen_table_iter_with_fn(&self) -> TokenStream {
         let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
         let row_type = name_generator.get_row_type_ident();

codegen/src/generators/read_only/table/impls.rs

@@ -19,6 +19,7 @@ impl ReadOnlyGenerator {
         let reinsert_fn = self.gen_table_reinsert_fn();
         let upsert_fn = self.gen_table_upsert_fn();
         let get_next_fn = self.gen_table_get_next_fn();
+        let pk_gen_state_fn = self.gen_table_pk_gen_state_fn();
         let iter_with_fn = self.gen_table_iter_with_fn();
         let iter_with_async_fn = self.gen_table_iter_with_async_fn();
         let count_fn = self.gen_table_count_fn();
@@ -37,6 +38,7 @@ impl ReadOnlyGenerator {
                 #upsert_fn
                 #count_fn
                 #get_next_fn
+                #pk_gen_state_fn
                 #iter_with_fn
                 #iter_with_async_fn
                 #system_info_fn
@@ -229,6 +231,17 @@ impl ReadOnlyGenerator {
         }
     }
 
+    fn gen_table_pk_gen_state_fn(&self) -> TokenStream {
+        let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
+        let primary_key_type = name_generator.get_primary_key_type_ident();
+
+        quote! {
+            pub fn pk_gen_state(&self) -> <<#primary_key_type as TablePrimaryKey>::Generator as PrimaryKeyGeneratorState>::State {
+                self.0.pk_gen.get_state()
+            }
+        }
+    }
+
     fn gen_table_iter_with_fn(&self) -> TokenStream {
         let name_generator = WorktableNameGenerator::from_table_name(self.name.to_string());
         let row_type = name_generator.get_row_type_ident();

codegen/src/migration_engine/generator.rs

@@ -36,12 +36,13 @@ pub fn generate(input: MigrationEngineInput) -> TokenStream {
         quote! {
             async fn #fn_name(
                 source_path: &str,
-                target: &#current_table,
+                target: &mut #current_table,
                 ctx: &#ctx_type,
             ) -> eyre::Result<()> {
                 let config = DiskConfig::new_with_table_name(source_path, #table_name_lit, #version);
                 let engine = ReadOnlyPersistenceEngine::create(config).await?;
                 let source = #table_path::load(engine).await?;
+                target.0.pk_gen = worktable::prelude::PrimaryKeyGeneratorState::from_state(source.pk_gen_state());
 
                 let rows = source.select_all().execute()?;
                 for row in rows {
@@ -57,7 +58,7 @@ pub fn generate(input: MigrationEngineInput) -> TokenStream {
     let match_arms = input.version_tables.keys().map(|version| {
         let fn_name = Ident::new(&format!("migrate_v{}", version), current_table.span());
         quote! {
-            #version => Self::#fn_name(source_path, &target, ctx).await?,
+            #version => Self::#fn_name(source_path, &mut target, ctx).await?,
         }
     });
 
@@ -77,7 +78,7 @@ pub fn generate(input: MigrationEngineInput) -> TokenStream {
 
                 let target_config = DiskConfig::new_with_table_name(target_path, #table_name_lit, #current_version);
                 let target_engine = #persistence_engine::new(target_config).await?;
-                let target = #current_table::new(target_engine).await?;
+                let mut target = #current_table::new(target_engine).await?;
 
                 match version {
                     #( #match_arms )*

tests/migration/mod.rs

@@ -242,6 +242,108 @@ fn test_migrate_v2_to_current() {
     });
 }
 
+#[test]
+fn test_next_pk_and_indexes_after_migration() {
+    let runtime = tokio::runtime::Builder::new_multi_thread()
+        .worker_threads(2)
+        .enable_io()
+        .enable_time()
+        .build()
+        .unwrap();
+
+    runtime.block_on(async {
+        let src = "tests/data/migration/next_pk_after_migration";
+        let dst = "tests/data/migration/next_pk_after_migration/dst";
+        remove_dir_if_exists(src.to_string()).await;
+
+        {
+            let config = DiskConfig::new_with_table_name(
+                src,
+                v2::UserWorkTable::name_snake_case(),
+                v2::UserWorkTable::version(),
+            );
+            let engine = v2::UserPersistenceEngine::new(config).await.unwrap();
+            let table = v2::UserWorkTable::load(engine).await.unwrap();
+
+            table
+                .insert(v2::UserRow {
+                    id: table.get_next_pk().into(),
+                    name: "Eve".to_string(),
+                    email: "eve@test.com".to_string(),
+                })
+                .unwrap();
+            table
+                .insert(v2::UserRow {
+                    id: table.get_next_pk().into(),
+                    name: "Frank".to_string(),
+                    email: "frank@test.com".to_string(),
+                })
+                .unwrap();
+
+            table.wait_for_ops().await;
+        }
+
+        let ctx = UserMigrationContext {
+            default_email: "unknown@example.com".to_string(),
+            default_created_at: chrono::Utc::now().timestamp() as u64,
+        };
+
+        let report = UserMigrationEngine::migrate(src, dst, &ctx).await.unwrap();
+        assert_eq!(report.source_version, v2::UserWorkTable::version());
+
+        let inserted = {
+            let config =
+                DiskConfig::new_with_table_name(dst, UserWorkTable::name_snake_case(), UserWorkTable::version());
+            let engine = UserPersistenceEngine::new(config).await.unwrap();
+            let table = UserWorkTable::load(engine).await.unwrap();
+
+            let next_pk = table.get_next_pk();
+            assert_eq!(next_pk, 2_u64.into());
+
+            let inserted = UserRow {
+                id: next_pk.into(),
+                name: "Grace".to_string(),
+                email: "grace@test.com".to_string(),
+                created_at: ctx.default_created_at + 1,
+            };
+
+            table.insert(inserted.clone()).unwrap();
+            table.wait_for_ops().await;
+
+            assert_eq!(table.count(), 3);
+            assert_eq!(table.select(inserted.id), Some(inserted.clone()));
+
+            let by_name = table.select_by_name(inserted.name.clone()).execute().unwrap();
+            assert_eq!(by_name, vec![inserted.clone()]);
+
+            let by_name_range = table
+                .select_by_name_range(inserted.name.clone()..=inserted.name.clone())
+                .execute()
+                .unwrap();
+            assert_eq!(by_name_range, vec![inserted.clone()]);
+
+            let by_pk_range = table.select_by_pk_range(inserted.id..=inserted.id).execute().unwrap();
+            assert_eq!(by_pk_range, vec![inserted.clone()]);
+
+            inserted
+        };
+
+        {
+            let config =
+                DiskConfig::new_with_table_name(dst, UserWorkTable::name_snake_case(), UserWorkTable::version());
+            let engine = UserPersistenceEngine::new(config).await.unwrap();
+            let table = UserWorkTable::load(engine).await.unwrap();
+
+            assert_eq!(table.count(), 3);
+            assert_eq!(table.select(inserted.id), Some(inserted.clone()));
+            assert_eq!(
+                table.select_by_name(inserted.name.clone()).execute().unwrap(),
+                vec![inserted.clone()]
+            );
+        }
+    });
+}
+
 /// Nonexistent source returns an error
 #[test]
 fn test_nonexistent_source_error() {