Skip to content

Commit 14887f1

Browse files
bakwcbakwc
committed
Refactor
1 parent b29286a commit 14887f1

File tree

1 file changed

+70
-72
lines changed

1 file changed

+70
-72
lines changed

‎mysql_ch_replicator/converter.py‎

Lines changed: 70 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -736,6 +736,72 @@ def __convert_alter_table_change_column(self, db_name, table_name, tokens):
736736
query = f'ALTER TABLE `{db_name}`.`{table_name}` RENAME COLUMN {column_name} TO {new_column_name}'
737737
self.db_replicator.clickhouse_api.execute_command(query)
738738

739+
def _handle_create_table_like(self, create_statement, source_table_name, target_table_name, is_query_api=True):
740+
"""
741+
Helper method to handle CREATE TABLE LIKE statements.
742+
743+
Args:
744+
create_statement: The original CREATE TABLE LIKE statement
745+
source_table_name: Name of the source table being copied
746+
target_table_name: Name of the new table being created
747+
is_query_api: If True, returns both MySQL and CH structures; if False, returns only MySQL structure
748+
749+
Returns:
750+
Either (mysql_structure, ch_structure) if is_query_api=True, or just mysql_structure otherwise
751+
"""
752+
# Try to get the actual structure from the existing table structures first
753+
if (hasattr(self, 'db_replicator') and
754+
self.db_replicator is not None and
755+
hasattr(self.db_replicator, 'state') and
756+
hasattr(self.db_replicator.state, 'tables_structure')):
757+
758+
# Check if the source table structure is already in our state
759+
if source_table_name in self.db_replicator.state.tables_structure:
760+
# Get the existing structure
761+
source_mysql_structure, source_ch_structure = self.db_replicator.state.tables_structure[source_table_name]
762+
763+
# Create a new structure with the target table name
764+
new_mysql_structure = copy.deepcopy(source_mysql_structure)
765+
new_mysql_structure.table_name = target_table_name
766+
767+
# Convert to ClickHouse structure
768+
new_ch_structure = copy.deepcopy(source_ch_structure)
769+
new_ch_structure.table_name = target_table_name
770+
771+
return (new_mysql_structure, new_ch_structure) if is_query_api else new_mysql_structure
772+
773+
# If we couldn't get it from state, try with MySQL API
774+
if (hasattr(self, 'db_replicator') and
775+
self.db_replicator is not None and
776+
hasattr(self.db_replicator, 'mysql_api') and
777+
self.db_replicator.mysql_api is not None):
778+
779+
try:
780+
# Get the CREATE statement for the source table
781+
source_create_statement = self.db_replicator.mysql_api.get_table_create_statement(source_table_name)
782+
783+
# Parse the source table structure
784+
source_structure = self.parse_mysql_table_structure(source_create_statement)
785+
786+
# Copy the structure but keep the new table name
787+
mysql_structure = copy.deepcopy(source_structure)
788+
mysql_structure.table_name = target_table_name
789+
790+
if is_query_api:
791+
# Convert to ClickHouse structure
792+
ch_structure = self.convert_table_structure(mysql_structure)
793+
return mysql_structure, ch_structure
794+
else:
795+
return mysql_structure
796+
797+
except Exception as e:
798+
error_msg = f"Could not get source table structure for LIKE statement: {str(e)}"
799+
print(f"Error: {error_msg}")
800+
raise Exception(error_msg, create_statement)
801+
802+
# If we got here, we couldn't determine the structure
803+
raise Exception(f"Could not determine structure for source table '{source_table_name}' in LIKE statement", create_statement)
804+
739805
def parse_create_table_query(self, mysql_query) -> tuple[TableStructure, TableStructure]:
740806
# Special handling for CREATE TABLE LIKE statements
741807
if 'LIKE' in mysql_query.upper():
@@ -748,55 +814,8 @@ def parse_create_table_query(self, mysql_query) -> tuple[TableStructure, TableSt
748814
new_table_name = match.group(1).strip('`"')
749815
source_table_name = match.group(2).strip('`"')
750816

751-
# Try to get the actual structure from the existing table structures first
752-
if (hasattr(self, 'db_replicator') and
753-
self.db_replicator is not None and
754-
hasattr(self.db_replicator, 'state') and
755-
hasattr(self.db_replicator.state, 'tables_structure')):
756-
757-
# Check if the source table structure is already in our state
758-
if source_table_name in self.db_replicator.state.tables_structure:
759-
# Get the existing structure
760-
source_mysql_structure, source_ch_structure = self.db_replicator.state.tables_structure[source_table_name]
761-
762-
# Create a new structure with the target table name
763-
new_mysql_structure = copy.deepcopy(source_mysql_structure)
764-
new_mysql_structure.table_name = new_table_name
765-
766-
# Convert to ClickHouse structure
767-
new_ch_structure = copy.deepcopy(source_ch_structure)
768-
new_ch_structure.table_name = new_table_name
769-
770-
return new_mysql_structure, new_ch_structure
771-
772-
# If we couldn't get it from state, try with MySQL API
773-
if (hasattr(self, 'db_replicator') and
774-
self.db_replicator is not None and
775-
hasattr(self.db_replicator, 'mysql_api') and
776-
self.db_replicator.mysql_api is not None):
777-
778-
try:
779-
# Get the CREATE statement for the source table
780-
source_create_statement = self.db_replicator.mysql_api.get_table_create_statement(source_table_name)
781-
782-
# Parse the source table structure
783-
source_structure = self.parse_mysql_table_structure(source_create_statement)
784-
785-
# Create a new structure with the target table name
786-
mysql_table_structure = copy.deepcopy(source_structure)
787-
mysql_table_structure.table_name = new_table_name
788-
789-
# Convert to ClickHouse structure
790-
ch_table_structure = self.convert_table_structure(mysql_table_structure)
791-
792-
return mysql_table_structure, ch_table_structure
793-
except Exception as e:
794-
error_msg = f"Could not get source table structure for LIKE statement: {str(e)}"
795-
print(f"Error: {error_msg}")
796-
raise Exception(error_msg, mysql_query)
797-
798-
# If we got here, we couldn't determine the structure
799-
raise Exception(f"Could not determine structure for source table '{source_table_name}' in LIKE statement", mysql_query)
817+
# Use the common helper method to handle the LIKE statement
818+
return self._handle_create_table_like(mysql_query, source_table_name, new_table_name, True)
800819

801820
# Regular parsing for non-LIKE statements
802821
mysql_table_structure = self.parse_mysql_table_structure(mysql_query)
@@ -851,29 +870,8 @@ def parse_mysql_table_structure(self, create_statement, required_table_name=None
851870
source_table_name = strip_sql_name(tokens[4].get_real_name())
852871
target_table_name = strip_sql_name(tokens[2].get_real_name())
853872

854-
# Try to get the actual structure - we need a valid MySQL API
855-
if not (hasattr(self, 'db_replicator') and
856-
self.db_replicator is not None and
857-
hasattr(self.db_replicator, 'mysql_api') and
858-
self.db_replicator.mysql_api is not None):
859-
raise Exception(f"Cannot parse LIKE statement - no MySQL API available to get source table structure", create_statement)
860-
861-
try:
862-
# Try to get the CREATE statement from the current database
863-
mysql_api = self.db_replicator.mysql_api
864-
source_create_statement = mysql_api.get_table_create_statement(source_table_name)
865-
866-
# Parse the source table structure
867-
source_structure = self.parse_mysql_table_structure(source_create_statement)
868-
869-
# Copy the structure but keep the new table name
870-
structure = copy.deepcopy(source_structure)
871-
structure.table_name = target_table_name
872-
return structure
873-
except Exception as e:
874-
error_msg = f"Could not get source table structure: {str(e)}"
875-
print(f"Error: {error_msg}")
876-
raise Exception(error_msg, create_statement)
873+
# Use the common helper method to handle the LIKE statement
874+
return self._handle_create_table_like(create_statement, source_table_name, target_table_name, False)
877875

878876
if not isinstance(tokens[3], sqlparse.sql.Parenthesis):
879877
raise Exception('wrong create statement', create_statement)

0 commit comments

Comments
 (0)