fix: Ensure last_completed_partition is saved on final migration state update

Problem: The final migration_state update (when marking migration as complete)
was not passing last_partition parameter, so the last completed partition was
being lost in migration_state table. If migration was interrupted at any point,
resume would lose the partition tracking.

Solution:
1. Track last_processed_partition throughout the migration loop
2. Update it when each partition completes
3. Pass it to final _update_migration_state() call when marking migration as complete

Additional fix:
- Use correct postgres_pk column when querying MAX() ID for final state update
- This ensures we get the correct last ID even for tables with non-standard PK names

🤖 Generated with Claude Code

Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
This commit is contained in:
2025-12-26 18:24:28 +01:00
parent 1430ef206f
commit ca2f7c5756

View File

@@ -100,6 +100,7 @@ class FullMigrator:
migrated = previous_migrated_count migrated = previous_migrated_count
migration_start_time = datetime.utcnow().isoformat() migration_start_time = datetime.utcnow().isoformat()
batch_count = 0 batch_count = 0
last_processed_partition = last_completed_partition # Track last partition we process
with ProgressTracker( with ProgressTracker(
rows_to_migrate, rows_to_migrate,
@@ -227,6 +228,7 @@ class FullMigrator:
# NOW partition is complete - update with completed partition # NOW partition is complete - update with completed partition
logger.info(f"Partition {partition} complete: {partition_group_count} groups consolidated") logger.info(f"Partition {partition} complete: {partition_group_count} groups consolidated")
last_processed_partition = partition # Track this partition as processed
self._update_migration_state( self._update_migration_state(
pg_conn, migrated, None, migration_start_time, pg_conn, migrated, None, migration_start_time,
last_partition=partition last_partition=partition
@@ -237,11 +239,12 @@ class FullMigrator:
logger.info(f"Final count from PostgreSQL: {final_count}") logger.info(f"Final count from PostgreSQL: {final_count}")
# Update migration state with final count and mark as completed # Update migration state with final count and mark as completed
# Get the actual last ID from the table # Get the actual last ID from the table using correct PK column
try: try:
with pg_conn.connection.cursor() as cursor: with pg_conn.connection.cursor() as cursor:
pk_column = self.config.get("postgres_pk", "id")
cursor.execute( cursor.execute(
f"SELECT MAX(id) FROM {pg_table}" f"SELECT MAX({pk_column}) FROM {pg_table}"
) )
result = cursor.fetchone() result = cursor.fetchone()
final_last_id = result[0] if result and result[0] else None final_last_id = result[0] if result and result[0] else None
@@ -250,9 +253,10 @@ class FullMigrator:
logger.warning(f"Failed to get final last ID: {e}") logger.warning(f"Failed to get final last ID: {e}")
final_last_id = None final_last_id = None
logger.info(f"About to update migration_state with count={final_count}, last_id={final_last_id}") logger.info(f"About to update migration_state with count={final_count}, last_id={final_last_id}, last_partition={last_processed_partition}")
self._update_migration_state( self._update_migration_state(
pg_conn, final_count, final_last_id, migration_start_time, is_final=True pg_conn, final_count, final_last_id, migration_start_time, is_final=True,
last_partition=last_processed_partition
) )
logger.info(f"Migration state update complete") logger.info(f"Migration state update complete")