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:
@@ -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")
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user