db_replicator refactor, splitted into several classes

Author: bakwc

mysql_ch_replicator/db_replicator.py

@@ -19,16 +19,23 @@
 from .table_structure import TableStructure, TableField
 from .binlog_replicator import DataReader, LogEvent, EventType
 from .utils import GracefulKiller, touch_all_files, format_floats
+from .db_replicator_initial import DbReplicatorInitial
+from .common import Status
 
 
 logger = getLogger(__name__)
 
 
-class Status(Enum):
-    NONE = 0
-    CREATING_INITIAL_STRUCTURES = 1
-    PERFORMING_INITIAL_REPLICATION = 2
-    RUNNING_REALTIME_REPLICATION = 3
+@dataclass
+class Statistics:
+    last_transaction: tuple = None
+    events_count: int = 0
+    insert_events_count: int = 0
+    insert_records_count: int = 0
+    erase_events_count: int = 0
+    erase_records_count: int = 0
+    no_events_count: int = 0
+    cpu_load: float = 0.0
 
 
 class State:
@@ -87,21 +94,8 @@ def remove(self):
             os.remove(file_name + '.tmp')
 
 
-@dataclass
-class Statistics:
-    last_transaction: tuple = None
-    events_count: int = 0
-    insert_events_count: int = 0
-    insert_records_count: int = 0
-    erase_events_count: int = 0
-    erase_records_count: int = 0
-    no_events_count: int = 0
-    cpu_load: float = 0.0
-
-
 class DbReplicator:
 
-    INITIAL_REPLICATION_BATCH_SIZE = 50000
     SAVE_STATE_INTERVAL = 10
     STATS_DUMP_INTERVAL = 60
     BINLOG_TOUCH_INTERVAL = 120
@@ -189,6 +183,9 @@ def __init__(self, config: Settings, database: str, target_database: str = None,
         self.last_records_upload_time = 0
         self.last_touch_time = 0
         self.start_time = time.time()
+        
+        # Create the initial replicator instance
+        self.initial_replicator = DbReplicatorInitial(self)
 
     def create_state(self):
         return State(self.state_path)
@@ -204,18 +201,6 @@ def validate_database_settings(self):
                     'Otherwise you will get DUPLICATES in your SELECT queries\n\n\n'
                 )
 
-    def validate_mysql_structure(self, mysql_structure: TableStructure):
-        for key_idx in mysql_structure.primary_key_ids:
-            primary_field: TableField = mysql_structure.fields[key_idx]
-            if 'not null' not in primary_field.parameters.lower():
-                logger.warning('primary key validation failed')
-                logger.warning(
-                    f'\n\n\n    !!!  WARNING - PRIMARY KEY NULLABLE (field "{primary_field.name}", table "{mysql_structure.table_name}") !!!\n\n'
-                    'There could be errors replicating nullable primary key\n'
-                    'Please ensure all tables has NOT NULL parameter for primary key\n'
-                    'Or mark tables as skipped, see "exclude_tables" option\n\n\n'
-                )
-
     def run(self):
         try:
             logger.info('launched db_replicator')
@@ -233,7 +218,7 @@ def run(self):
                 self.run_realtime_replication()
                 return
             if self.state.status == Status.PERFORMING_INITIAL_REPLICATION:
-                self.perform_initial_replication()
+                self.initial_replicator.perform_initial_replication()
                 self.run_realtime_replication()
                 return
 
@@ -249,42 +234,13 @@ def run(self):
             self.state.last_processed_transaction = self.data_reader.get_last_transaction_id()
             self.state.save()
             logger.info(f'last known transaction {self.state.last_processed_transaction}')
-            self.create_initial_structure()
-            self.perform_initial_replication()
+            self.initial_replicator.create_initial_structure()
+            self.initial_replicator.perform_initial_replication()
             self.run_realtime_replication()
         except Exception:
             logger.error(f'unhandled exception', exc_info=True)
             raise
 
-    def create_initial_structure(self):
-        self.state.status = Status.CREATING_INITIAL_STRUCTURES
-        for table in self.state.tables:
-            self.create_initial_structure_table(table)
-        self.state.save()
-
-    def create_initial_structure_table(self, table_name):
-        if not self.config.is_table_matches(table_name):
-            return
-
-        if self.single_table and self.single_table != table_name:
-            return
-
-        mysql_create_statement = self.mysql_api.get_table_create_statement(table_name)
-        mysql_structure = self.converter.parse_mysql_table_structure(
-            mysql_create_statement, required_table_name=table_name,
-        )
-        self.validate_mysql_structure(mysql_structure)
-        clickhouse_structure = self.converter.convert_table_structure(mysql_structure)
-        
-        # Always set if_not_exists to True to prevent errors when tables already exist
-        clickhouse_structure.if_not_exists = True
-        
-        self.state.tables_structure[table_name] = (mysql_structure, clickhouse_structure)
-        indexes = self.config.get_indexes(self.database, table_name)
-
-        if not self.is_parallel_worker:
-            self.clickhouse_api.create_table(clickhouse_structure, additional_indexes=indexes)
-
     def prevent_binlog_removal(self):
         if time.time() - self.last_touch_time < self.BINLOG_TOUCH_INTERVAL:
             return
@@ -295,195 +251,6 @@ def prevent_binlog_removal(self):
         self.last_touch_time = time.time()
         touch_all_files(binlog_directory)
 
-    def perform_initial_replication(self):
-        self.clickhouse_api.database = self.target_database_tmp
-        logger.info('running initial replication')
-        self.state.status = Status.PERFORMING_INITIAL_REPLICATION
-        self.state.save()
-        start_table = self.state.initial_replication_table
-        for table in self.state.tables:
-            if start_table and table != start_table:
-                continue
-            if self.single_table and self.single_table != table:
-                continue
-            self.perform_initial_replication_table(table)
-            start_table = None
-
-        if not self.is_parallel_worker:
-            logger.info(f'initial replication - swapping database')
-            if self.target_database in self.clickhouse_api.get_databases():
-                self.clickhouse_api.execute_command(
-                    f'RENAME DATABASE `{self.target_database}` TO `{self.target_database}_old`',
-                )
-                self.clickhouse_api.execute_command(
-                    f'RENAME DATABASE `{self.target_database_tmp}` TO `{self.target_database}`',
-                )
-                self.clickhouse_api.drop_database(f'{self.target_database}_old')
-            else:
-                self.clickhouse_api.execute_command(
-                    f'RENAME DATABASE `{self.target_database_tmp}` TO `{self.target_database}`',
-                )
-            self.clickhouse_api.database = self.target_database
-        logger.info(f'initial replication - done')
-
-    def perform_initial_replication_table(self, table_name):
-        logger.info(f'running initial replication for table {table_name}')
-
-        if not self.config.is_table_matches(table_name):
-            logger.info(f'skip table {table_name} - not matching any allowed table')
-            return
-
-        if not self.is_parallel_worker and self.config.initial_replication_threads > 1:
-            self.state.initial_replication_table = table_name
-            self.state.initial_replication_max_primary_key = None
-            self.state.save()
-            self.perform_initial_replication_table_parallel(table_name)
-            return
-
-        max_primary_key = None
-        if self.state.initial_replication_table == table_name:
-            # continue replication from saved position
-            max_primary_key = self.state.initial_replication_max_primary_key
-            logger.info(f'continue from primary key {max_primary_key}')
-        else:
-            # starting replication from zero
-            logger.info(f'replicating from scratch')
-            self.state.initial_replication_table = table_name
-            self.state.initial_replication_max_primary_key = None
-            self.state.save()
-
-        mysql_table_structure, clickhouse_table_structure = self.state.tables_structure[table_name]
-
-        logger.debug(f'mysql table structure: {mysql_table_structure}')
-        logger.debug(f'clickhouse table structure: {clickhouse_table_structure}')
-
-        field_types = [field.field_type for field in clickhouse_table_structure.fields]
-
-        primary_keys = clickhouse_table_structure.primary_keys
-        primary_key_ids = clickhouse_table_structure.primary_key_ids
-        primary_key_types = [field_types[key_idx] for key_idx in primary_key_ids]
-
-        #logger.debug(f'primary key name: {primary_key}, type: {primary_key_type}')
-
-        stats_number_of_records = 0
-        last_stats_dump_time = time.time()
-
-        while True:
-
-            query_start_values = max_primary_key
-            if query_start_values is not None:
-                for i in range(len(query_start_values)):
-                    key_type = primary_key_types[i]
-                    value = query_start_values[i]
-                    if 'int' not in key_type.lower():
-                        value = f"'{value}'"
-                        query_start_values[i] = value
-
-            records = self.mysql_api.get_records(
-                table_name=table_name,
-                order_by=primary_keys,
-                limit=DbReplicator.INITIAL_REPLICATION_BATCH_SIZE,
-                start_value=query_start_values,
-                worker_id=self.worker_id,
-                total_workers=self.total_workers,
-            )
-            logger.debug(f'extracted {len(records)} records from mysql')
-
-            records = self.converter.convert_records(records, mysql_table_structure, clickhouse_table_structure)
-
-            if self.config.debug_log_level:
-                logger.debug(f'records: {records}')
-
-            if not records:
-                break
-            self.clickhouse_api.insert(table_name, records, table_structure=clickhouse_table_structure)
-            for record in records:
-                record_primary_key = [record[key_idx] for key_idx in primary_key_ids]
-                if max_primary_key is None:
-                    max_primary_key = record_primary_key
-                else:
-                    max_primary_key = max(max_primary_key, record_primary_key)
-
-            self.state.initial_replication_max_primary_key = max_primary_key
-            self.save_state_if_required()
-            self.prevent_binlog_removal()
-
-            stats_number_of_records += len(records)
-            curr_time = time.time()
-            if curr_time - last_stats_dump_time >= 60.0:
-                last_stats_dump_time = curr_time
-                logger.info(
-                    f'replicating {table_name}, '
-                    f'replicated {stats_number_of_records} records, '
-                    f'primary key: {max_primary_key}',
-                )
-
-        logger.info(
-            f'finish replicating {table_name}, '
-            f'replicated {stats_number_of_records} records, '
-            f'primary key: {max_primary_key}',
-        )
-
-    def perform_initial_replication_table_parallel(self, table_name):
-        """
-        Execute initial replication for a table using multiple parallel worker processes.
-        Each worker will handle a portion of the table based on its worker_id and total_workers.
-        """
-        logger.info(f"Starting parallel replication for table {table_name} with {self.config.initial_replication_threads} workers")
-
-        # Create and launch worker processes
-        processes = []
-        for worker_id in range(self.config.initial_replication_threads):
-            # Prepare command to launch a worker process
-            cmd = [
-                sys.executable, "-m", "mysql_ch_replicator.main",
-                "db_replicator",  # Required positional mode argument
-                "--config", self.settings_file,
-                "--db", self.database,
-                "--worker_id", str(worker_id),
-                "--total_workers", str(self.config.initial_replication_threads),
-                "--table", table_name,
-                "--target_db", self.target_database_tmp,
-                "--initial_only=True",
-            ]
-                
-            logger.info(f"Launching worker {worker_id}: {' '.join(cmd)}")
-            process = subprocess.Popen(cmd)
-            processes.append(process)
-        
-        # Wait for all worker processes to complete
-        logger.info(f"Waiting for {len(processes)} workers to complete replication of {table_name}")
-        
-        try:
-            while processes:
-                for i, process in enumerate(processes[:]):
-                    # Check if process is still running
-                    if process.poll() is not None:
-                        exit_code = process.returncode
-                        if exit_code == 0:
-                            logger.info(f"Worker process {i} completed successfully")
-                        else:
-                            logger.error(f"Worker process {i} failed with exit code {exit_code}")
-                            # Optional: can raise an exception here to abort the entire operation
-                            raise Exception(f"Worker process failed with exit code {exit_code}")
-                        
-                        processes.remove(process)
-                
-                if processes:
-                    # Wait a bit before checking again
-                    time.sleep(0.1)
-                    
-                    # Every 30 seconds, log progress
-                    if int(time.time()) % 30 == 0:
-                        logger.info(f"Still waiting for {len(processes)} workers to complete")
-        except KeyboardInterrupt:
-            logger.warning("Received interrupt, terminating worker processes")
-            for process in processes:
-                process.terminate()
-            raise
-        
-        logger.info(f"All workers completed replication of table {table_name}")
-
     def run_realtime_replication(self):
         if self.initial_only:
             logger.info('skip running realtime replication, only initial replication was requested')

test_mysql_ch_replicator.py

@@ -14,7 +14,7 @@
 from mysql_ch_replicator import mysql_api
 from mysql_ch_replicator import clickhouse_api
 from mysql_ch_replicator.binlog_replicator import State as BinlogState, FileReader, EventType, BinlogReplicator
-from mysql_ch_replicator.db_replicator import State as DbReplicatorState, DbReplicator
+from mysql_ch_replicator.db_replicator import State as DbReplicatorState, DbReplicator, DbReplicatorInitial
 from mysql_ch_replicator.converter import MysqlToClickhouseConverter
 
 from mysql_ch_replicator.runner import ProcessRunner
@@ -1049,7 +1049,7 @@ def test_json():
 
 
 def test_string_primary_key(monkeypatch):
-    monkeypatch.setattr(DbReplicator, 'INITIAL_REPLICATION_BATCH_SIZE', 1)
+    monkeypatch.setattr(DbReplicatorInitial, 'INITIAL_REPLICATION_BATCH_SIZE', 1)
 
     cfg = config.Settings()
     cfg.load(CONFIG_FILE)
@@ -1111,7 +1111,7 @@ def test_string_primary_key(monkeypatch):
 
 
 def test_if_exists_if_not_exists(monkeypatch):
-    monkeypatch.setattr(DbReplicator, 'INITIAL_REPLICATION_BATCH_SIZE', 1)
+    monkeypatch.setattr(DbReplicatorInitial, 'INITIAL_REPLICATION_BATCH_SIZE', 1)
 
     cfg = config.Settings()
     cfg.load(CONFIG_FILE)
@@ -1152,7 +1152,7 @@ def test_if_exists_if_not_exists(monkeypatch):
 
 
 def test_percona_migration(monkeypatch):
-    monkeypatch.setattr(DbReplicator, 'INITIAL_REPLICATION_BATCH_SIZE', 1)
+    monkeypatch.setattr(DbReplicatorInitial, 'INITIAL_REPLICATION_BATCH_SIZE', 1)
 
     cfg = config.Settings()
     cfg.load(CONFIG_FILE)
@@ -1230,7 +1230,7 @@ def test_percona_migration(monkeypatch):
 
 
 def test_add_column_first_after_and_drop_column(monkeypatch):
-    monkeypatch.setattr(DbReplicator, 'INITIAL_REPLICATION_BATCH_SIZE', 1)
+    monkeypatch.setattr(DbReplicatorInitial, 'INITIAL_REPLICATION_BATCH_SIZE', 1)
 
     cfg = config.Settings()
     cfg.load(CONFIG_FILE)