alex/mysql2postgres
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: Use actual PostgreSQL row count for total_rows_migrated tracking

Browse Source 
Replace session-level counting with direct table COUNT queries to ensure
total_rows_migrated always reflects actual reality in PostgreSQL. This fixes
the discrepancy where the counter was only tracking rows from the current session
and didn't account for earlier insertions or duplicates from failed resume attempts.

Key improvements:
- Use get_row_count() after each batch to get authoritative total
- Preserve previous count on resume and accumulate across sessions
- Remove dependency on error-prone session-level counters
- Ensures migration_state.total_rows_migrated matches actual table row count

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
This commit is contained in:
alex
2025-12-23 15:33:27 +01:00
parent b09cfcf9df
commit 0f217379ea
8 changed files with 646 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

190
src/migrator/full_migration.py
View File

@@ -9,6 +9,7 @@ from src.connectors.postgres_connector import PostgreSQLConnector
from src.transformers.data_transformer import DataTransformer
from src.utils.logger import get_logger, setup_logger
from src.utils.progress import ProgressTracker
from src.migrator.state import MigrationState
logger = get_logger(__name__)
@@ -28,20 +29,23 @@ class FullMigrator:
self.table = table
self.config = TABLE_CONFIGS[table]
self.settings = get_settings()
self.state = MigrationState()
def migrate(self, dry_run: bool = False) -> int:
"""Perform full migration of the table.
def migrate(self, dry_run: bool = False, resume: bool = False) -> int:
"""Perform full migration of the table with resume capability.
Args:
dry_run: If True, log what would be done but don't modify data
resume: If True, resume from last checkpoint; if False, check for conflicts
Returns:
Total number of rows migrated
Total number of rows migrated in this run
"""
setup_logger(__name__)
mysql_table = self.config["mysql_table"]
pg_table = self.config["postgres_table"]
primary_key = self.config.get("primary_key", "id")
logger.info(f"Starting full migration of {mysql_table} -> {pg_table}")
@@ -49,11 +53,7 @@ class FullMigrator:
with MySQLConnector() as mysql_conn:
# Get total row count
total_rows = mysql_conn.get_row_count(mysql_table)
logger.info(f"Total rows to migrate: {total_rows}")
if dry_run:
logger.info("[DRY RUN] Would migrate all rows")
return total_rows
logger.info(f"Total rows in source: {total_rows}")
with PostgreSQLConnector() as pg_conn:
# Check if table exists
@@ -63,18 +63,52 @@ class FullMigrator:
"Run 'setup --create-schema' first."
)
migrated = 0
# Check for previous migration state
last_migrated_id = self._get_last_migrated_id(pg_conn, pg_table)
previous_migrated_count = self._get_previous_migrated_count(pg_conn, pg_table)
if last_migrated_id is not None:
pg_row_count = pg_conn.get_row_count(pg_table)
logger.warning(
f"Found previous migration state: {pg_row_count} rows already in {pg_table}"
)
if not resume:
raise ValueError(
f"Migration already in progress for {pg_table}. "
f"Use --resume to continue from last checkpoint, or delete data to restart."
)
logger.info(f"Resuming from ID > {last_migrated_id}")
rows_to_migrate = total_rows - last_migrated_id
else:
last_migrated_id = None
previous_migrated_count = 0
rows_to_migrate = total_rows
if dry_run:
logger.info(f"[DRY RUN] Would migrate {rows_to_migrate} rows")
return rows_to_migrate
migrated = previous_migrated_count
migration_start_time = datetime.utcnow().isoformat()
with ProgressTracker(
total_rows,
rows_to_migrate,
f"Migrating {mysql_table}"
) as progress:
# Fetch and migrate rows in batches
for batch in mysql_conn.fetch_all_rows(mysql_table):
# Transform batch
# Use ordered fetching for node consolidation with resume support
for batch in mysql_conn.fetch_rows_ordered_for_consolidation(
mysql_table,
start_id=last_migrated_id
):
if not batch:
break
# Transform batch with consolidation enabled
transformed = DataTransformer.transform_batch(
mysql_table,
batch
batch,
consolidate=True
)
# Insert batch
@@ -85,65 +119,151 @@ class FullMigrator:
columns
)
migrated += inserted
progress.update(inserted)
if inserted > 0:
# For consolidated batches, count transformed rows, not source rows
migrated += inserted
progress.update(inserted)
# Update state after each batch for resume capability
# Use MAX id of the batch (represents last MySQL id processed)
batch_max_id = max(
int(row.get("id", 0)) for row in transformed
)
# Get actual row count from PostgreSQL for accuracy
actual_count = pg_conn.get_row_count(pg_table)
self._update_migration_state(
pg_conn, actual_count, batch_max_id, migration_start_time
)
logger.info(
f"✓ Migration complete: {migrated} rows migrated "
f"to {pg_table}"
)
# Update migration state
self._update_migration_state(pg_conn, migrated)
return migrated
except Exception as e:
logger.error(f"Migration failed: {e}")
raise
def _update_migration_state(
self,
pg_conn: PostgreSQLConnector,
rows_migrated: int
) -> None:
"""Update migration state tracking table.
def _get_last_migrated_id(self, pg_conn: PostgreSQLConnector, pg_table: str) -> Optional[int]:
"""Get the last migrated MySQL ID from migration_state table.
Args:
pg_conn: PostgreSQL connection
rows_migrated: Number of rows migrated
pg_table: PostgreSQL table name
Returns:
Last migrated MySQL ID or None if no previous migration
"""
try:
pg_table = self.config["postgres_table"]
with pg_conn.connection.cursor() as cursor:
cursor.execute(
"SELECT last_migrated_id FROM migration_state WHERE table_name = %s",
(pg_table,)
)
result = cursor.fetchone()
if result and result[0]:
return result[0]
except Exception:
pass
return None
def _get_previous_migrated_count(self, pg_conn: PostgreSQLConnector, pg_table: str) -> int:
"""Get the total rows migrated so far from migration_state table.
Args:
pg_conn: PostgreSQL connection
pg_table: PostgreSQL table name
Returns:
Total rows migrated so far (0 if no previous migration)
"""
try:
with pg_conn.connection.cursor() as cursor:
cursor.execute(
"SELECT total_rows_migrated FROM migration_state WHERE table_name = %s",
(pg_table,)
)
result = cursor.fetchone()
if result and result[0]:
return result[0]
except Exception:
pass
return 0
def _update_migration_state(
self,
pg_conn: PostgreSQLConnector,
rows_migrated: int,
last_id: Optional[int] = None,
migration_start_time: Optional[str] = None
) -> None:
"""Update migration state in PostgreSQL and state file.
Args:
pg_conn: PostgreSQL connection
rows_migrated: Total number of rows migrated so far
last_id: Last ID that was migrated (for resume capability)
migration_start_time: When the migration started (ISO format)
"""
pg_table = self.config["postgres_table"]
now = datetime.utcnow()
status = "in_progress" if last_id is not None else "completed"
# Update PostgreSQL migration_state table
try:
# Use COALESCE to handle both insert (first time) and update (resume)
# For resume: total_rows_migrated will be the full accumulated count
query = f"""
INSERT INTO migration_state
(table_name, last_migrated_timestamp, total_rows_migrated, migration_completed_at, status)
VALUES (%s, %s, %s, %s, %s)
(table_name, last_migrated_timestamp, last_migrated_id, total_rows_migrated, migration_completed_at, status)
VALUES (%s, %s, %s, %s, %s, %s)
ON CONFLICT (table_name) DO UPDATE SET
last_migrated_timestamp = EXCLUDED.last_migrated_timestamp,
last_migrated_id = EXCLUDED.last_migrated_id,
total_rows_migrated = EXCLUDED.total_rows_migrated,
migration_completed_at = EXCLUDED.migration_completed_at,
status = EXCLUDED.status
"""
now = datetime.utcnow()
pg_conn.execute(query, (pg_table, now, rows_migrated, now, "completed"))
logger.debug("Migration state updated")
pg_conn.execute(
query,
(
pg_table,
migration_start_time or now.isoformat(),
last_id,
rows_migrated,
now if status == "completed" else None,
status
)
)
logger.debug(f"Migration state updated: {rows_migrated} rows total, last_id={last_id}, status={status}")
except Exception as e:
logger.warning(f"Failed to update migration state: {e}")
logger.warning(f"Failed to update migration state in PostgreSQL: {e}")
# Also save to state file for incremental migrations
try:
self.state.set_last_timestamp(pg_table, migration_start_time or now.isoformat())
self.state.increment_migration_count(pg_table, rows_migrated)
logger.debug("Migration state saved to file")
except Exception as e:
logger.warning(f"Failed to save migration state to file: {e}")
def run_full_migration(
table: str,
dry_run: bool = False
dry_run: bool = False,
resume: bool = False
) -> int:
"""Run full migration for a table.
Args:
table: Table name to migrate
dry_run: If True, show what would be done without modifying data
resume: If True, resume from last checkpoint instead of starting fresh
Returns:
Number of rows migrated
Number of rows migrated in this run
"""
migrator = FullMigrator(table)
return migrator.migrate(dry_run=dry_run)
return migrator.migrate(dry_run=dry_run, resume=resume)

44
src/migrator/incremental_migration.py
View File

@@ -1,6 +1,7 @@
"""Incremental migration from MySQL to PostgreSQL based on timestamps."""
from datetime import datetime
from typing import Optional
import psycopg
from config import get_settings, TABLE_CONFIGS
from src.connectors.mysql_connector import MySQLConnector
@@ -31,6 +32,33 @@ class IncrementalMigrator:
self.settings = get_settings()
self.state = MigrationState(state_file)
def _get_last_timestamp_from_db(
self,
pg_conn: PostgreSQLConnector,
pg_table: str
) -> Optional[str]:
"""Get last migration timestamp from PostgreSQL migration_state table.
Args:
pg_conn: PostgreSQL connector
pg_table: PostgreSQL table name
Returns:
ISO format timestamp or None if not found
"""
try:
with pg_conn.connection.cursor() as cursor:
cursor.execute(
"SELECT last_migrated_timestamp FROM migration_state WHERE table_name = %s",
(pg_table,)
)
result = cursor.fetchone()
if result and result[0]:
return result[0].isoformat()
except psycopg.Error:
return None
return None
def migrate(self, dry_run: bool = False, use_id: bool = False) -> int:
"""Perform incremental migration since last sync.
@@ -88,9 +116,19 @@ class IncrementalMigrator:
Returns:
Number of rows migrated
"""
# Get last migration timestamp
# Try to get last migration timestamp from state file first
last_timestamp = self.state.get_last_timestamp(pg_table)
# If not in state file, try to get from PostgreSQL migration_state table
if last_timestamp is None:
try:
last_timestamp = self._get_last_timestamp_from_db(pg_conn, pg_table)
if last_timestamp:
logger.info(f"Found previous migration state in database: {last_timestamp}")
except Exception as e:
logger.debug(f"Could not read from migration_state table: {e}")
last_timestamp = None
if last_timestamp is None:
logger.info(
f"No previous migration found for {pg_table}. "
@@ -98,6 +136,8 @@ class IncrementalMigrator:
)
return 0
logger.info(f"Last migration timestamp: {last_timestamp}")
# Count rows to migrate
timestamp_col = "updated_at" if mysql_table == "ELABDATADISP" else "created_at"
@@ -107,7 +147,7 @@ class IncrementalMigrator:
timestamp_col
)
logger.info(f"Last timestamp in PostgreSQL: {pg_max_timestamp}")
logger.info(f"Current max timestamp in PostgreSQL: {pg_max_timestamp}")
if dry_run:
logger.info("[DRY RUN] Would migrate rows after timestamp")