fix migrations

Author: Puyodead1

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

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)