diff --git a/alembic/versions/684292138a0a_many_devices_to_many_users.py b/alembic/versions/684292138a0a_many_devices_to_many_users.py index 9bdfe45..d2de730 100644 --- a/alembic/versions/684292138a0a_many_devices_to_many_users.py +++ b/alembic/versions/684292138a0a_many_devices_to_many_users.py @@ -9,7 +9,6 @@ from typing import Sequence, Union import sqlalchemy as sa -from sqlalchemy.dialects import mysql from alembic import op @@ -21,22 +20,12 @@ def upgrade() -> None: - # create a user <-> device association table - op.create_table( - "user_device", - sa.Column("user_id", sa.VARCHAR(255), nullable=False), - sa.Column("device_code", sa.VARCHAR(255), nullable=False), - sa.ForeignKeyConstraint(["device_code"], ["devices.code"], ondelete="CASCADE"), - sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"), - ) - - # remove useless columns, these never served any purpose + # Remove columns from devices table that are no longer needed op.drop_column("devices", "security_level") op.drop_column("devices", "session_id_type") # rename code to info op.alter_column("devices", "code", new_column_name="info", existing_type=sa.String(255)) - # add new code column with a more unique value op.add_column( "devices", sa.Column( @@ -60,8 +49,6 @@ def upgrade() -> None: # remove all rows where uploaded_by is null, as they are not associated with a user op.execute("DELETE FROM devices WHERE uploaded_by IS NULL") - - # set uploaded_by to non nullable op.alter_column("devices", "uploaded_by", nullable=False, existing_type=sa.String(255)) # generate code for rows as a sha256 in the format of "client_id_blob_filename sha256:device_private_key sha265:uploaded_by" @@ -80,30 +67,27 @@ def upgrade() -> None: ) # for mapping, cause we need to dedupe and create a unique constraint - old_devices = op.get_bind().execute(sa.text("SELECT * FROM devices")).fetchall() + op.execute("DELETE FROM devices WHERE id NOT IN (SELECT MIN(id) FROM devices GROUP BY code)") + op.create_unique_constraint(None, "devices", ["code"]) - user_device_insert = [] - for device in old_devices: - user_device_insert.append(f"('{device[6]}', '{device[7]}')") + # create a user <-> device association table + op.create_table( + "user_device", + sa.Column("user_id", sa.VARCHAR(255), nullable=False), + sa.Column("device_code", sa.VARCHAR(255), nullable=False), + sa.ForeignKeyConstraint(["device_code"], ["devices.code"], ondelete="CASCADE"), + sa.ForeignKeyConstraint(["user_id"], ["users.id"], ondelete="CASCADE"), + ) - # Join the values and create the full INSERT statement + # Insert existing data into user_device table + old_devices = op.get_bind().execute(sa.text("SELECT * FROM devices")).fetchall() + user_device_insert = [f"('{device[4]}', '{device[5]}')" for device in old_devices] if user_device_insert: user_device_insert_sql = "INSERT INTO user_device (user_id, device_code) VALUES " + ",".join(user_device_insert) - - # insert data into user_device op.execute(sa.text(user_device_insert_sql)) else: print("No devices found to insert into user_device table.") - # deduplicate devices by code - op.execute("DELETE FROM devices WHERE id NOT IN (SELECT MIN(id) FROM devices GROUP BY code)") - - # add a unique constraint to the code column - op.create_unique_constraint(None, "devices", ["code"]) - - # reset ids - # op.execute("SET @id=0; UPDATE devices SET id=@id:=@id+1") - def downgrade() -> None: raise NotImplementedError("Downgrade is not supported for this migration.") diff --git a/alembic/versions/8fbc59f03986_convert_device_to_wvd.py b/alembic/versions/8fbc59f03986_convert_device_to_wvd.py index 1136b66..09cb2d1 100644 --- a/alembic/versions/8fbc59f03986_convert_device_to_wvd.py +++ b/alembic/versions/8fbc59f03986_convert_device_to_wvd.py @@ -37,19 +37,24 @@ def upgrade() -> None: op.get_bind().execute(sa.text("SELECT code,client_id_blob_filename,device_private_key FROM devices")).fetchall() ) for code, client_id, private_key in devices: - logger.info(f"Converting device {code} to WVD") - wvd = Device( - type_=DeviceTypes.ANDROID, - security_level=3, - flags=None, - private_key=base64.b64decode(private_key), - client_id=base64.b64decode(client_id), - ) - - wvd_b64 = base64.b64encode(wvd.dumps()).decode() - op.get_bind().execute( - sa.text("UPDATE devices SET wvd = :wvd WHERE code = :code"), {"wvd": wvd_b64, "code": code} - ) + try: + logger.info(f"Converting device {code} to WVD") + wvd = Device( + type_=DeviceTypes.ANDROID, + security_level=3, + flags=None, + private_key=base64.b64decode(private_key), + client_id=base64.b64decode(client_id), + ) + + wvd_b64 = base64.b64encode(wvd.dumps()).decode() + op.get_bind().execute( + sa.text("UPDATE devices SET wvd = :wvd WHERE code = :code"), {"wvd": wvd_b64, "code": code} + ) + except Exception as e: + logger.error(f"Failed to convert device {code} to WVD: {e}\nPK: {private_key}\nCID: {client_id}") + # remove the device from the database + op.get_bind().execute(sa.text("DELETE FROM devices WHERE code = :code"), {"code": code}) # make the wvd column non-nullable op.alter_column("devices", "wvd", existing_type=sa.Text(), nullable=False)