started implementing multi-primary-key

Author: bakwc

mysql_ch_replicator/clickhouse_api.py

@@ -86,7 +86,7 @@ def set_last_used_version(self, table_name, last_used_version):
         self.tables_last_record_version[table_name] = last_used_version
 
     def create_table(self, structure: TableStructure):
-        if not structure.primary_key:
+        if not structure.primary_keys:
             raise Exception(f'missing primary key for {structure.table_name}')
 
         primary_key_type = ''

mysql_ch_replicator/converter.py

@@ -2,7 +2,7 @@
 import json
 import sqlparse
 import re
-from pyparsing import Word, alphas, alphanums
+from pyparsing import Suppress, CaselessKeyword, Word, alphas, alphanums, delimitedList
 
 from .table_structure import TableStructure, TableField
 
@@ -216,7 +216,7 @@ def convert_table_structure(self, mysql_structure: TableStructure) -> TableStruc
                 name=field.name,
                 field_type=clickhouse_field_type,
             ))
-        clickhouse_structure.primary_key = mysql_structure.primary_key
+        clickhouse_structure.primary_keys = mysql_structure.primary_keys
         clickhouse_structure.preprocess()
         return clickhouse_structure
 
@@ -519,9 +519,26 @@ def parse_mysql_table_structure(self, create_statement, required_table_name=None
             if line.lower().startswith('constraint'):
                 continue
             if line.lower().startswith('primary key'):
-                pattern = 'PRIMARY KEY (' + Word(alphanums + '_`') + ')'
+                # pattern = 'PRIMARY KEY (' + Word(alphanums + '_`') + ')'
+                # result = pattern.parseString(line)
+                # structure.primary_key = strip_sql_name(result[1])
+
+                # Define identifier to match column names, handling backticks and unquoted names
+                identifier = (Suppress('`') + Word(alphas + alphanums + '_') + Suppress('`')) | Word(
+                    alphas + alphanums + '_')
+
+                # Build the parsing pattern
+                pattern = CaselessKeyword('PRIMARY') + CaselessKeyword('KEY') + Suppress('(') + delimitedList(
+                    identifier)('column_names') + Suppress(')')
+
+                # Parse the line
                 result = pattern.parseString(line)
-                structure.primary_key = strip_sql_name(result[1])
+
+                # Extract and process the primary key column names
+                primary_keys = [strip_sql_name(name) for name in result['column_names']]
+
+                structure.primary_keys = primary_keys
+
                 continue
 
             #print(" === processing line", line)
@@ -541,16 +558,16 @@ def parse_mysql_table_structure(self, create_statement, required_table_name=None
             #print(' ---- params:', field_parameters)
 
 
-        if not structure.primary_key:
+        if not structure.primary_keys:
             for field in structure.fields:
                 if 'primary key' in field.parameters.lower():
-                    structure.primary_key = field.name
+                    structure.primary_keys.append(field.name)
 
-        if not structure.primary_key:
+        if not structure.primary_keys:
             if structure.has_field('id'):
-                structure.primary_key = 'id'
+                structure.primary_keys = ['id']
 
-        if not structure.primary_key:
+        if not structure.primary_keys:
             raise Exception(f'No primary key for table {structure.table_name}, {create_statement}')
 
         structure.preprocess()

mysql_ch_replicator/db_replicator.py

@@ -148,15 +148,16 @@ def validate_database_settings(self):
                 )
 
     def validate_mysql_structure(self, mysql_structure: TableStructure):
-        primary_field: TableField = mysql_structure.fields[mysql_structure.primary_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'
-            )
+        pass
+        # primary_field: TableField = mysql_structure.fields[mysql_structure.primary_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:
@@ -279,11 +280,12 @@ def perform_initial_replication_table(self, table_name):
         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_key
-        primary_key_index = field_names.index(primary_key)
-        primary_key_type = field_types[primary_key_index]
+        primary_key = clickhouse_table_structure.primary_keys
+        primary_key_ids = clickhouse_table_structure.primary_key_ids
+
+        #primary_key_type = field_types[primary_key_index]
 
-        logger.debug(f'primary key name: {primary_key}, type: {primary_key_type}')
+        #logger.debug(f'primary key name: {primary_key}, type: {primary_key_type}')
 
         stats_number_of_records = 0
         last_stats_dump_time = time.time()

mysql_ch_replicator/table_structure.py

@@ -9,13 +9,15 @@ class TableField:
 @dataclass
 class TableStructure:
     fields: list = field(default_factory=list)
-    primary_key: str = ''
-    primary_key_idx: int = 0
+    primary_keys: str = ''
+    primary_key_ids: int = 0
     table_name: str = ''
 
     def preprocess(self):
         field_names = [f.name for f in self.fields]
-        self.primary_key_idx = field_names.index(self.primary_key)
+        self.primary_key_ids = [
+            field_names.index(key) for key in self.primary_keys
+        ]
 
     def add_field_after(self, new_field: TableField, after: str):
 

test_mysql_ch_replicator.py

@@ -226,7 +226,7 @@ def test_e2e_multistatement():
     id int NOT NULL AUTO_INCREMENT,
     name varchar(255),
     age int,
-    PRIMARY KEY (id)
+    PRIMARY KEY (id, `name`)
 ); 
     ''')