diff --git a/src/migrator/full_migration.py b/src/migrator/full_migration.py
index 36e529a..200b888 100644
--- a/src/migrator/full_migration.py
+++ b/src/migrator/full_migration.py
@@ -100,6 +100,7 @@ class FullMigrator:
                     migrated = previous_migrated_count
                     migration_start_time = datetime.utcnow().isoformat()
                     batch_count = 0
+                    last_processed_partition = last_completed_partition  # Track last partition we process
 
                     with ProgressTracker(
                         rows_to_migrate,
@@ -227,6 +228,7 @@ class FullMigrator:
 
                             # NOW partition is complete - update with completed partition
                             logger.info(f"Partition {partition} complete: {partition_group_count} groups consolidated")
+                            last_processed_partition = partition  # Track this partition as processed
                             self._update_migration_state(
                                 pg_conn, migrated, None, migration_start_time,
                                 last_partition=partition
@@ -237,11 +239,12 @@ class FullMigrator:
                     logger.info(f"Final count from PostgreSQL: {final_count}")
 
                     # 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:
                         with pg_conn.connection.cursor() as cursor:
+                            pk_column = self.config.get("postgres_pk", "id")
                             cursor.execute(
-                                f"SELECT MAX(id) FROM {pg_table}"
+                                f"SELECT MAX({pk_column}) FROM {pg_table}"
                             )
                             result = cursor.fetchone()
                             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}")
                         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(
-                        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")