Fix: Automation script of merging ALTER PK to CREATE DDL in tests

migtests/scripts/add-pk-from-alter-to-create

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+
+
+import os
+import re
+import shutil
+
+table_file_path = os.getenv('TEST_DIR')+'/export-dir/schema/tables/table.sql'
+
+#copy the table_file_path to table_file_path+'.old'
+src = table_file_path
+dest = table_file_path + '.old'
+
+shutil.copy(src,dest)
+
+with open(table_file_path, 'r') as file:
+    lines = file.readlines()
+
+create_tables = []
+alter_tables = []
+create_table_map = {}
+alter_table_map = {}
+for i, line in enumerate(lines):
+    if line.startswith('CREATE TABLE'):
+        match = re.search(r'CREATE TABLE (.*) ', line.split('(')[0]) 
+        if match:
+            table_name = match.group(1)
+            create_table_map[table_name] = line
+        create_tables.append(line)
+    elif line.strip(' ').startswith('ADD CONSTRAINT') and line.strip(' ').find('PRIMARY KEY') != -1 and lines[i-1].startswith('ALTER TABLE'):
+        match = re.search(r'ALTER TABLE ONLY (.*)', lines[i-1])
+        if match:
+            table_name = match.group(1)
+            alter_table_map[table_name] = lines[i]
+        alter_tables.append(lines[i])
+
+
+output = []
+
+for k,v in create_table_map.items():
+    # Extract the primary key constraint from the ALTER TABLE statement
+    if k not in alter_table_map:
+        output.append(v)
+        continue
+    match = re.search(r'ADD CONSTRAINT (.*) PRIMARY KEY \((.*)\);', alter_table_map[k])
+    if match:
+        constraint = match.group(1)
+        columns = match.group(2)
+
+        # Add the primary key constraint to the CREATE TABLE statement
+        if "WITH" in v:
+            create_table = v.replace(') WITH', ',' + 'PRIMARY KEY (' +  columns + ')) WITH')
+        else:
+            create_table = v + '    PRIMARY KEY (' + columns + '),\n'
+        output.append(create_table)
+
+with open(table_file_path, 'w') as file:
+    # modify create_table lines to output ones
+    for i, line in enumerate(lines):
+        if line in create_tables:
+            file.write(output.pop(0))
+        elif line in alter_tables and line.find('PRIMARY KEY') != -1: # remove the ADD CONTRAINT lines with PK only
+            pass
+        elif line.startswith('ALTER TABLE') and (i < len(lines)-1 and lines[i+1].find('PRIMARY KEY') != -1): # remove the ALTER TABLE lines for PK only
+            pass
+        else:
+            file.write(line)

migtests/scripts/live-migration-fallb-run-test.sh

@@ -113,9 +113,9 @@ main() {
 		run_ysql yugabyte "CREATE DATABASE ${TARGET_DB_NAME}"
 	fi
 
-	if [ -x "${TEST_DIR}/add-pk-from-alter-to-create" ]
+	if [ -x "${SCRIPTS}/add-pk-from-alter-to-create" ]
 	then
-		"${TEST_DIR}/add-pk-from-alter-to-create"
+		"${SCRIPTS}/add-pk-from-alter-to-create"
 	fi
 
 	step "Import schema."

migtests/scripts/live-migration-fallf-run-test.sh

@@ -121,9 +121,9 @@ main() {
 		run_ysql yugabyte "CREATE DATABASE ${TARGET_DB_NAME}"
 	fi
 
-	if [ -x "${TEST_DIR}/add-pk-from-alter-to-create" ]
+	if [ -x "${SCRIPTS}/add-pk-from-alter-to-create" ]
 	then
-		"${TEST_DIR}/add-pk-from-alter-to-create"
+		"${SCRIPTS}/add-pk-from-alter-to-create"
 	fi
 
 	step "Import schema."

migtests/scripts/live-migration-run-test.sh

@@ -112,9 +112,9 @@ main() {
 		run_ysql yugabyte "CREATE DATABASE ${TARGET_DB_NAME}"
 	fi
 
-	if [ -x "${TEST_DIR}/add-pk-from-alter-to-create" ]
+	if [ -x "${SCRIPTS}/add-pk-from-alter-to-create" ]
 	then
-		"${TEST_DIR}/add-pk-from-alter-to-create"
+		"${SCRIPTS}/add-pk-from-alter-to-create"
 	fi
 
 	step "Import schema."

migtests/scripts/run-test.sh

@@ -127,9 +127,9 @@ main() {
 		run_ysql yugabyte "CREATE DATABASE ${TARGET_DB_NAME}"
 	fi
 
-	if [ -x "${TEST_DIR}/add-pk-from-alter-to-create" ]
+	if [ -x "${SCRIPTS}/add-pk-from-alter-to-create" ]
 	then
-		"${TEST_DIR}/add-pk-from-alter-to-create"
+		"${SCRIPTS}/add-pk-from-alter-to-create"
 	fi
 
 	step "Import schema."