PostgreSQL change type of primary key using alembic

upgrade.py

from dataclasses import dataclass

from alembic import op

@dataclass
class ForeignKeyConstraint:
    name: str
    table: str
    column: str
    sql: str

    @classmethod
    def from_row(cls, row) -> 'ForeignKeyConstraint':
        return cls(
            name=row['name'],
            table=row['table'],
            column=row['column'],
            sql=row['sql'],
        )


def change_column_to_uuid(table_name: str, column_name: str) -> None:
    # Step 1: Find and remember all foreign key constraints
    constraints_raw = op.get_bind().execute(f'''
        SELECT c.conname as "name",
               (select r.relname from pg_class r where r.oid = c.conrelid) as "table",
               UNNEST((select array_agg(attname) from pg_attribute where attrelid = c.conrelid and array[attnum] <@ c.conkey)) as "column",
               pg_get_constraintdef(c.oid) as "sql"
          FROM pg_constraint c join pg_attribute a on c.confrelid=a.attrelid and a.attnum = ANY(confkey)
         WHERE c.confrelid = (select oid from pg_class where relname = '{table_name}')
           AND c.confrelid != c.conrelid;
    ''')
    constraints_raw = constraints_raw.fetchall()
    constraints = [ForeignKeyConstraint.from_row(row) for row in constraints_raw]

    # Step 2: Drop the foreign key constraints
    for constraint in constraints:
        op.execute(
            f'ALTER TABLE {constraint.table} '
            f'DROP CONSTRAINT {constraint.name}'
        )

    # Step 3: Change the type of the primary key column using "USING"
    op.execute(
        f'ALTER TABLE {table_name} '
        f'ALTER COLUMN {column_name} '
        f'TYPE uuid USING {column_name}::uuid'
    )

    # Step 4: Change the type of all foreign key columns using "USING"
    for constraint in constraints_raw:
        op.execute(
            f'ALTER TABLE {constraint.table} '
            f'ALTER COLUMN {constraint.column} '
            f'TYPE uuid USING {constraint.column}::uuid'
        )

    # Step 5: Create exactly the same foreign key constraints that were dropped
    for constraint in constraints:
        op.execute(
            f'ALTER TABLE {constraint.table} '
            f'ADD CONSTRAINT {constraint.column} '
            f'{constraint.sql}'
        )


def upgrade():
    change_column_to_uuid('documents', 'id')
    
def downgrade():
    pass