postgres_add_constraints_to_create_table_query.py

import re
import sys

def printf(format, *args):
    sys.stdout.write(format % args)

# Open the SQL file
f = open('input.sql', 'r')
sql = f.read()

obj = {}

o = open("output.sql", "w")

# Split the file into a list of statements
statements = sql.split(';')

# Compile regular expressions for each type of constraint
primary_key_regex = re.compile(r'PRIMARY KEY\s*\(([^)]+)\)')
unique_regex = re.compile(r'UNIQUE\s*\(([^)]+)\)')
foreign_key_regex = re.compile(r'FOREIGN KEY\s*\(([^)]+)\)\s*REFERENCES\s*([^(]*)\(([^)]+)\)')

# Go through each statement
for statement in statements:
    # Ignore empty statements
    if not statement.strip():
        continue

    statement = statement.strip();

    printf("\nProcessing statement: %s", statement)

    # Check if the statement is an ALTER TABLE statement
    if statement.lower().startswith('alter table'):
        
        # Split the statement into words
        words = statement.split()

        # Get the name of the table being altered
        table_name = words[2]

        # Check if the statement is adding a PRIMARY KEY constraint
        m = primary_key_regex.search(statement)
        if m:
            # Get the column name for the PRIMARY KEY
            column_name = m.group(1)

            # Find the CREATE TABLE statement for this table
            for s in statements:
                s = s.strip()
                if s.lower().startswith(f'create table {table_name.lower()}'):
                    obj[table_name.lower()] = []
                    for line in s.splitlines():
                        obj[table_name.lower()].append(line + '\n')

            obj[table_name.lower()].insert(1, '    PRIMARY KEY ({}),'.format(column_name))

            printf("Found PRIMARY KEY constraint: %s", statement)

        # Check if the statement is adding a UNIQUE constraint
        m = unique_regex.search(statement)
        if m:
            # Get the column names for the UNIQUE constraint
            column_names = m.group(1)

            for s in statements:
                s = s.strip()
                if s.lower().startswith(f'create table {table_name.lower()}'):
                    obj[table_name.lower()] = []
                    for line in s.splitlines():
                        obj[table_name.lower()].append(line + '\n')

            obj[table_name.lower()].insert(-1, '    UNIQUE ({})'.format(column_names))
            obj[table_name.lower()][-3] = obj[table_name.lower()][-3] + ',' 

            printf("Found UNIQUE constraint: %s", statement)

        # Check if the statement is adding a FOREIGN KEY constraint
        m = foreign_key_regex.search(statement)
        if m:
            # Get the column name for the FOREIGN KEY and the name of the referenced table and column
            column_name, ref_table_name, ref_column_name = m.groups()

            # Find the CREATE TABLE statement for this table
            for s in statements:
                s = s.strip()
                if s.lower().startswith(f'create table {table_name.lower()}'):
                    obj[table_name.lower()] = []
                    for line in s.splitlines():
                        obj[table_name.lower()].append(line + '\n')

            obj[table_name.lower()].insert(-1, '    FOREIGN KEY ({}) REFERENCE {}({})'.format(column_name, ref_table_name, ref_column_name))
            obj[table_name.lower()][-3] = obj[table_name.lower()][-3] + ',' 

            printf("Found FOREIGN KEY constraint: %s", statement)
        
        # Write statement to file
        

for key, item in obj.items():
    print('Writing for ' + key + '...\n')
    o.write('--\n-- Creation of ' + key + ' table\n--\n\n')
    print(obj[key])
    for item in obj[key]:
        o.write(item)
    o.write('\n\n')

# Close the file
o.close()