"""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")