Fixed tests

Author: bakwc

mysql_ch_replicator/clickhouse_api.py

@@ -16,16 +16,15 @@
 (
 {fields},
     `_version` UInt64,
-    INDEX _version _version TYPE minmax GRANULARITY 1,
-    INDEX idx_id {primary_key} TYPE bloom_filter GRANULARITY 1
+    {indexes}
 )
 ENGINE = ReplacingMergeTree(_version)
 {partition_by}ORDER BY {primary_key}
 SETTINGS index_granularity = 8192
 '''
 
 DELETE_QUERY = '''
-DELETE FROM {db_name}.{table_name} WHERE {field_name} IN ({field_values})
+DELETE FROM {db_name}.{table_name} WHERE ({field_name}) IN ({field_values})
 '''
 
 
@@ -63,8 +62,6 @@ def get_databases(self):
         return database_list
 
     def execute_command(self, query):
-        #print(' === executing ch query', query)
-
         for attempt in range(ClickhouseApi.MAX_RETRIES):
             try:
                 self.client.command(query)
@@ -76,7 +73,6 @@ def execute_command(self, query):
                 time.sleep(ClickhouseApi.RETRY_INTERVAL)
 
     def recreate_database(self):
-        #print(' === creating database', self.database)
         self.execute_command(f'DROP DATABASE IF EXISTS {self.database}')
         self.execute_command(f'CREATE DATABASE {self.database}')
 
@@ -90,28 +86,36 @@ def create_table(self, structure: TableStructure):
         if not structure.primary_keys:
             raise Exception(f'missing primary key for {structure.table_name}')
 
-        primary_key_type = ''
-        for field in structure.fields:
-            if field.name == structure.primary_key:
-                primary_key_type = field.field_type
-        if not primary_key_type:
-            raise Exception(f'failed to get type of primary key {structure.table_name} {structure.primary_key}')
-
         fields = [
             f'    `{field.name}` {field.field_type}' for field in structure.fields
         ]
         fields = ',\n'.join(fields)
         partition_by = ''
 
-        if 'int' in primary_key_type.lower():
-            partition_by = f'PARTITION BY intDiv({structure.primary_key}, 4294967)\n'
+        if len(structure.primary_keys) == 1:
+            if 'int' in structure.fields[structure.primary_key_ids[0]].field_type.lower():
+                partition_by = f'PARTITION BY intDiv({structure.primary_keys[0]}, 4294967)\n'
+
+        indexes = [
+            'INDEX _version _version TYPE minmax GRANULARITY 1',
+        ]
+        if len(structure.primary_keys) == 1:
+            indexes.append(
+                f'INDEX idx_id {structure.primary_keys[0]} TYPE bloom_filter GRANULARITY 1',
+            )
+
+        indexes = ',\n'.join(indexes)
+        primary_key = ','.join(structure.primary_keys)
+        if len(structure.primary_keys) > 1:
+            primary_key = f'({primary_key})'
 
         query = CREATE_TABLE_QUERY.format(**{
             'db_name': self.database,
             'table_name': structure.table_name,
             'fields': fields,
-            'primary_key': structure.primary_key,
+            'primary_key': primary_key,
             'partition_by': partition_by,
+            'indexes': indexes,
         })
         self.execute_command(query)
 
@@ -161,6 +165,7 @@ def insert(self, table_name, records, table_structure: TableStructure = None):
         self.set_last_used_version(table_name, current_version)
 
     def erase(self, table_name, field_name, field_values):
+        field_name = ','.join(field_name)
         field_values = ', '.join(list(map(str, field_values)))
         query = DELETE_QUERY.format(**{
             'db_name': self.database,

mysql_ch_replicator/db_replicator.py

@@ -277,13 +277,11 @@ def perform_initial_replication_table(self, table_name):
         logger.debug(f'mysql table structure: {mysql_table_structure}')
         logger.debug(f'clickhouse table structure: {clickhouse_table_structure}')
 
-        field_names = [field.name for field in clickhouse_table_structure.fields]
         field_types = [field.field_type for field in clickhouse_table_structure.fields]
 
-        primary_key = clickhouse_table_structure.primary_keys
+        primary_keys = clickhouse_table_structure.primary_keys
         primary_key_ids = clickhouse_table_structure.primary_key_ids
-
-        #primary_key_type = field_types[primary_key_index]
+        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}')
 
@@ -292,15 +290,20 @@ def perform_initial_replication_table(self, table_name):
 
         while True:
 
-            query_start_value = max_primary_key
-            if 'int' not in primary_key_type.lower() and query_start_value is not None:
-                query_start_value = f"'{query_start_value}'"
+            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_key,
+                order_by=primary_keys,
                 limit=DbReplicator.INITIAL_REPLICATION_BATCH_SIZE,
-                start_value=query_start_value,
+                start_value=query_start_values,
             )
             logger.debug(f'extracted {len(records)} records from mysql')
 
@@ -313,7 +316,7 @@ def perform_initial_replication_table(self, table_name):
                 break
             self.clickhouse_api.insert(table_name, records, table_structure=clickhouse_table_structure)
             for record in records:
-                record_primary_key = record[primary_key_index]
+                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:
@@ -406,6 +409,16 @@ def save_state_if_required(self, force=False):
         self.state.tables_last_record_version = self.clickhouse_api.tables_last_record_version
         self.state.save()
 
+    def _get_record_id(self, ch_table_structure, record: list):
+        result = []
+        for idx in ch_table_structure.primary_key_ids:
+            field_type = ch_table_structure.fields[idx].field_type
+            if field_type == 'String':
+                result.append(f"'{record[idx]}'")
+            else:
+                result.append(record[idx])
+        return ','.join(map(str, result))
+
     def handle_insert_event(self, event: LogEvent):
         if self.config.debug_log_level:
             logger.debug(
@@ -420,12 +433,10 @@ def handle_insert_event(self, event: LogEvent):
         clickhouse_table_structure = self.state.tables_structure[event.table_name][1]
         records = self.converter.convert_records(event.records, mysql_table_structure, clickhouse_table_structure)
 
-        primary_key_ids = mysql_table_structure.primary_key_idx
-
         current_table_records_to_insert = self.records_to_insert[event.table_name]
         current_table_records_to_delete = self.records_to_delete[event.table_name]
         for record in records:
-            record_id = record[primary_key_ids]
+            record_id = self._get_record_id(clickhouse_table_structure, record)
             current_table_records_to_insert[record_id] = record
             current_table_records_to_delete.discard(record_id)
 
@@ -439,16 +450,9 @@ def handle_erase_event(self, event: LogEvent):
         self.stats.erase_events_count += 1
         self.stats.erase_records_count += len(event.records)
 
-        table_structure: TableStructure = self.state.tables_structure[event.table_name][0]
         table_structure_ch: TableStructure = self.state.tables_structure[event.table_name][1]
 
-        primary_key_name_idx = table_structure.primary_key_idx
-        field_type_ch = table_structure_ch.fields[primary_key_name_idx].field_type
-
-        if field_type_ch == 'String':
-            keys_to_remove = [f"'{record[primary_key_name_idx]}'" for record in event.records]
-        else:
-            keys_to_remove = [record[primary_key_name_idx] for record in event.records]
+        keys_to_remove = [self._get_record_id(table_structure_ch, record) for record in event.records]
 
         current_table_records_to_insert = self.records_to_insert[event.table_name]
         current_table_records_to_delete = self.records_to_delete[event.table_name]
@@ -548,12 +552,12 @@ def upload_records(self):
             if not keys_to_remove:
                 continue
             table_structure: TableStructure = self.state.tables_structure[table_name][0]
-            primary_key_name = table_structure.primary_key
+            primary_key_names = table_structure.primary_keys
             if self.config.debug_log_level:
-                logger.debug(f'erasing from {table_name}, primary key: {primary_key_name}, values: {keys_to_remove}')
+                logger.debug(f'erasing from {table_name}, primary key: {primary_key_names}, values: {keys_to_remove}')
             self.clickhouse_api.erase(
                 table_name=table_name,
-                field_name=primary_key_name,
+                field_name=primary_key_names,
                 field_values=keys_to_remove,
             )
 

mysql_ch_replicator/mysql_api.py

@@ -48,7 +48,6 @@ def create_database(self, db_name):
         self.cursor.execute(f'CREATE DATABASE {db_name}')
 
     def execute(self, command, commit=False):
-        #print(f'Executing: <{command}>')
         self.cursor.execute(command)
         if commit:
             self.db.commit()
@@ -88,9 +87,11 @@ def get_table_create_statement(self, table_name) -> str:
 
     def get_records(self, table_name, order_by, limit, start_value=None):
         self.reconnect_if_required()
+        order_by = ','.join(order_by)
         where = ''
         if start_value is not None:
-            where = f'WHERE {order_by} > {start_value} '
+            start_value = ','.join(map(str, start_value))
+            where = f'WHERE ({order_by}) > ({start_value}) '
         query = f'SELECT * FROM {table_name} {where}ORDER BY {order_by} LIMIT {limit}'
         self.cursor.execute(query)
         res = self.cursor.fetchall()