pagure/alembic/versions/f16ab75e4d32_sshkeys.py

"""Migrate SSH keys to the new format.

Revision ID: f16ab75e4d32
Revises: 0a8f99c161e2
Create Date: 2018-09-24 16:11:21.297620

"""

# revision identifiers, used by Alembic.
revision = "f16ab75e4d32"
down_revision = "0a8f99c161e2"

import datetime

from alembic import op
import sqlalchemy as sa

from pagure.lib.query import is_valid_ssh_key


def upgrade():
    """ Upgrade the database model for the way we store user's public ssh
    keys.

    For this we leverage the existing ``deploykeys`` table.
    It gets renamed to ``sshkeys``, we add the user_id foreign key as now
    ssh keys stored in this table can be linked to an user.
    Then we convert the existing ssh keys to this database model.
    Finally, we drop the ``public_ssh_key`` column from the ``users`` table.
    """
    users_table = sa.sql.table(
        "users",
        sa.sql.column("id", sa.Integer),
        sa.sql.column("public_ssh_key", sa.TEXT()),
    )
    sshkey_table = sa.sql.table(
        "sshkeys",
        sa.sql.column("id", sa.Integer),
        sa.sql.column("user_id", sa.Integer),
        sa.sql.column("public_ssh_key", sa.TEXT()),
        sa.sql.column("ssh_short_key", sa.TEXT()),
        sa.sql.column("ssh_search_key", sa.TEXT()),
        sa.sql.column("creator_user_id", sa.Integer),
        sa.sql.column("pushaccess", sa.Boolean),
        sa.sql.column("date_created", sa.DateTime),
    )

    op.rename_table("deploykeys", "sshkeys")

    op.add_column("sshkeys", sa.Column("user_id", sa.Integer(), nullable=True))
    op.create_index(
        op.f("ix_sshkeys_sshkeys_user_id"),
        "sshkeys",
        ["user_id"],
        unique=False,
    )
    op.create_foreign_key(
        op.f("sshkeys_user_id_fkey"),
        "sshkeys",
        "users",
        ["user_id"],
        ["id"],
        onupdate=u"CASCADE",
    )

    print("Convert existing ssh keys to the new format")

    seen = []
    conn = op.get_bind()
    for key in conn.execute(sshkey_table.select()):
        ssh_short_key = is_valid_ssh_key(key.public_ssh_key).strip()
        ssh_search_key = ssh_short_key.split(" ")[1]
        # There is no chance of dupes in the deploykeys alone
        seen.append(ssh_search_key)
        op.execute(
            sshkey_table.update()
            .where(sshkey_table.c.id == key.id)
            .values({
                "ssh_short_key": ssh_short_key,
                "ssh_search_key": ssh_search_key,
            })
        )

    data = []
    for user in conn.execute(users_table.select()):
        if not user.public_ssh_key:
            continue
        for key in user.public_ssh_key.split("\n"):
            if key in (None, False) or not key.strip():
                print("Skipping one key")
                continue
            ssh_short_key = is_valid_ssh_key(key)
            if not ssh_short_key:
                continue
            ssh_short_key = ssh_short_key.strip()
            ssh_search_key = ssh_short_key.split(" ")[1]
            if ssh_search_key in seen:
                print("Skipping previously seen key")
                continue
            seen.append(ssh_search_key)
            print("Key:    %s" % key)
            print("Short:  %s" % ssh_short_key)
            print("Search: %s" % ssh_search_key)
            tmp = {}
            tmp["user_id"] = user.id
            tmp["creator_user_id"] = user.id
            tmp["public_ssh_key"] = key
            tmp["ssh_search_key"] = ssh_search_key
            tmp["ssh_short_key"] = ssh_short_key
            tmp["pushaccess"] = True
            tmp['date_created'] = datetime.datetime.utcnow()
            data.append(tmp)

    op.bulk_insert(sshkey_table, data)

    op.drop_column("users", "public_ssh_key")


def downgrade():
    """ Downgrade the database model for the way we store user's public ssh
    keys.

    For this we bring back the keys present in the ``sshkeys`` table and
    put them back into the ``public_ssh_key`` column of the ``users`` table.
    """

    users_table = sa.sql.table(
        "users",
        sa.sql.column("id", sa.Integer),
        sa.sql.column("public_ssh_key", sa.TEXT()),
    )
    sshkey_table = sa.sql.table(
        "sshkeys",
        sa.sql.column("user_id", sa.Integer),
        sa.sql.column("public_ssh_key", sa.TEXT()),
        sa.sql.column("ssh_short_key", sa.TEXT()),
        sa.sql.column("ssh_search_key", sa.TEXT()),
        sa.sql.column("creator_user_id", sa.Integer),
        sa.sql.column("pushaccess", sa.Boolean),
        sa.sql.column("date_created", sa.DateTime),
    )

    op.add_column(
        "users", sa.Column("public_ssh_key", sa.TEXT(), nullable=True)
    )

    print("Convert existing ssh keys to the old format")

    conn = op.get_bind()
    data = []
    for key in conn.execute(sshkey_table.select()):
        if not key.user_id:
            continue
        user = [
            u
            for u in conn.execute(
                users_table.select().where(users_table.c.id == key.user_id)
            )
        ]
        user = user[0]
        ssh_key = ""
        if user.public_ssh_key:
            ssh_key = user.public_ssh_key + "\n"
        ssh_key += key.public_ssh_key

        op.execute(
            users_table.update()
            .where(users_table.c.id == key.user_id)
            .values({"public_ssh_key": ssh_key})
        )

    print("Remove the keys associated with users since we moved them")

    op.execute(
        sshkey_table.delete()
        .where(sshkey_table.c.user_id != None)
    )

    op.drop_constraint(
        op.f("sshkeys_user_id_fkey"), "sshkeys", type_="foreignkey"
    )
    op.drop_index(op.f("ix_sshkeys_sshkeys_user_id"), table_name="sshkeys")
    op.drop_column("sshkeys", "user_id")

    op.rename_table("sshkeys", "deploykeys")