huesoporro/migrations/20250228112643_uuids.py

"""
This module contains a Caribou migration.

Migration Name: uuids
Migration Version: 20250228112643
"""

import uuid


def upgrade(connection):
    """Several major upgrades
    - Update all table's id columns to be UUIDs.
    - Update the user's table `user` column to be named `username`
    - Update the chatbot's table `user_id` column to reference the `id` column in the user's table,
    thus changing the foreign key constraint.
    """

    # First, create temporary tables with the new schema
    connection.execute("""
     CREATE TABLE users_new (
         id              TEXT         not null PRIMARY KEY,
         username        varchar(255) not null UNIQUE,
         last_updated_at TIMESTAMP    default CURRENT_TIMESTAMP,
         created_at      TIMESTAMP    default CURRENT_TIMESTAMP,
         external_auth   JSON
     );
    """)

    # Fetch users to generate UUIDs in Python
    users = connection.execute(
        "SELECT user, last_updated_at, created_at, external_auth FROM users"
    ).fetchall()

    # Copy data from old users table to new users table, using Python-generated UUIDs
    for user in users:
        new_uuid = uuid.uuid4().hex
        # Safely handle JSON

        connection.execute(
            """
        INSERT INTO users_new (id, username, last_updated_at, created_at, external_auth)
        VALUES (?, ?, ?, ?, ?)
        """,
            (new_uuid, user[0], user[1], user[2], user[3]),
        )

    # Create temporary quotes table with new schema
    connection.execute("""
     CREATE TABLE quotes_new (
         id              TEXT         not null PRIMARY KEY,
         quote           varchar(255) not null UNIQUE,
         author          varchar(255),
         channel         varchar(255),
         created_at      TIMESTAMP    default CURRENT_TIMESTAMP,
         last_updated_at TIMESTAMP    default CURRENT_TIMESTAMP
     );
    """)

    # Fetch quotes to generate UUIDs in Python
    quotes = connection.execute(
        "SELECT quote, author, channel, created_at, last_updated_at FROM quotes"
    ).fetchall()

    # Copy data from old quotes table to new quotes table, using Python-generated UUIDs
    for quote in quotes:
        new_uuid = uuid.uuid4().hex
        connection.execute(
            """
        INSERT INTO quotes_new (id, quote, author, channel, created_at, last_updated_at)
        VALUES (?, ?, ?, ?, ?, ?)
        """,
            (new_uuid, quote[0], quote[1], quote[2], quote[3], quote[4]),
        )

    # Create mapping table to store the relationship between old user IDs and new UUIDs
    connection.execute("""
     CREATE TEMPORARY TABLE user_id_mapping (
         old_user VARCHAR(255) not null,
         new_id TEXT not null
     );
    """)

    # Populate the mapping table
    connection.execute("""
     INSERT INTO user_id_mapping (old_user, new_id)
     SELECT username, id FROM users_new;
    """)

    # Create temporary chatbot table with new schema
    connection.execute("""
     CREATE TABLE chatbot_new (
         id                         TEXT                  not null PRIMARY KEY,
         user_id                    TEXT                  not null UNIQUE,
         automatic_generation_timer INTEGER   default 300 not null,
         automatic_quote_timer      INTEGER   default 500 not null,
         mods                       VARCHAR(255),
         created_at                 TIMESTAMP default CURRENT_TIMESTAMP,
         last_updated_at            TIMESTAMP default CURRENT_TIMESTAMP,
         FOREIGN KEY (user_id) REFERENCES users_new (id)
     );
    """)

    # Fetch chatbot data to generate UUIDs in Python
    chatbots = connection.execute(
        "SELECT user_id, automatic_generation_timer, automatic_quote_timer, mods, created_at, last_updated_at FROM chatbot"
    ).fetchall()

    # Copy data from old chatbot table to new chatbot table with updated foreign key
    for chatbot in chatbots:
        # Get the new UUID for the user
        user_mapping = connection.execute(
            "SELECT new_id FROM user_id_mapping WHERE old_user = ?", (chatbot[0],)
        ).fetchone()
        new_user_id = user_mapping[0] if user_mapping else None

        if new_user_id:
            new_chatbot_uuid = uuid.uuid4().hex
            connection.execute(
                """
            INSERT INTO chatbot_new (id, user_id, automatic_generation_timer, automatic_quote_timer, mods, created_at, last_updated_at)
            VALUES (?, ?, ?, ?, ?, ?, ?)
            """,
                (
                    new_chatbot_uuid,
                    new_user_id,
                    chatbot[1],
                    chatbot[2],
                    chatbot[3],
                    chatbot[4],
                    chatbot[5],
                ),
            )

    # Drop old tables
    connection.execute("DROP TABLE chatbot;")
    connection.execute("DROP TABLE quotes;")
    connection.execute("DROP TABLE users;")

    # Rename new tables to original names
    connection.execute("ALTER TABLE users_new RENAME TO users;")
    connection.execute("ALTER TABLE quotes_new RENAME TO quotes;")
    connection.execute("ALTER TABLE chatbot_new RENAME TO chatbot;")

    # Drop temporary mapping table
    connection.execute("DROP TABLE user_id_mapping;")

    # Commit all changes
    connection.commit()


def downgrade(connection):
    # add your downgrade step here
    pass