From e1129fb30a458c10f98c190d7bcd6dd7133c0486 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ricardo=20Ca=C3=B1uelo?= <ricardo.canuelo@collabora.com>
Date: Wed, 20 Dec 2023 14:15:21 +0100
Subject: [PATCH] migrations: fix existing migrations scripts
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Remove the "set_timeout" migration script, as it won't be needed.
Make some fixes to the "user" script:

  - Keep "admin" users, converting from the old format (pertenence to
    the "admin" group) to the new one ("is_superuser" field)
  - Adjust the "is_verified" and "is_active" fields accordingly
  - Make the script more robust: run sanity checks on the data to apply
    the migrations only when needed
  - On upgrade, remove the "admin" group
  - Implement a symmetric downgrade function

Signed-off-by: Ricardo Cañuelo <ricardo.canuelo@collabora.com>
---
 migrations/20221014061654_set_timeout.py | 33 ---------
 migrations/20231102101356_user.py        | 94 ++++++++++++++++++++++--
 2 files changed, 89 insertions(+), 38 deletions(-)
 delete mode 100644 migrations/20221014061654_set_timeout.py

diff --git a/migrations/20221014061654_set_timeout.py b/migrations/20221014061654_set_timeout.py
deleted file mode 100644
index abf2bc86..00000000
--- a/migrations/20221014061654_set_timeout.py
+++ /dev/null
@@ -1,33 +0,0 @@
-# SPDX-License-Identifier: LGPL-2.1-or-later
-#
-# Copyright (C) 2022 Collabora Limited
-# Author: Jeny Sadadia <jeny.sadadia@collabora.com>
-
-"""
-Migration to update Node.timeout
-"""
-
-from datetime import timedelta
-
-name = '20221014061654_set_timeout'
-dependencies = []
-
-
-def upgrade(db: "pymongo.database.Database"):
-    nodes = db.node.find()
-    for node in nodes:
-        timeout = node['created'] + timedelta(hours=6)
-        db.node.update_one(
-            {
-                "_id": node['_id']
-            },
-            {
-                "$set": {
-                    "timeout": timeout
-                }
-            },
-        )
-
-
-def downgrade(db: "pymongo.database.Database"):
-    pass
diff --git a/migrations/20231102101356_user.py b/migrations/20231102101356_user.py
index 7257f550..61b2341f 100644
--- a/migrations/20231102101356_user.py
+++ b/migrations/20231102101356_user.py
@@ -4,17 +4,51 @@
 # Author: Jeny Sadadia <jeny.sadadia@collabora.com>
 
 """
-Migration for User schema
+Migration for User schema after the adoption of fastapi-users for
+user management in commit:
+
+    api.main: update/add user management routes
+
 """
 
 name = '20231102101356_user'
-dependencies = ['20221014061654_set_timeout']
+dependencies = []
+
+def user_upgrade_needed(user):
+    """Checks if a DB user passed as a parameter needs to be migrated
+    with this script.
+
+    Parameters:
+      user: a mongodb document (dict) defining a KernelCI user
+
+    Returns:
+      True if the user needs to be migrated, False otherwise
+
+    """
+    # The existence of a 'profile' key seems to be enough to detect a
+    # pre-migration user
+    if 'profile' in user:
+        return True
+    else:
+        return False
 
 
 def upgrade(db: "pymongo.database.Database"):
     users = db.user.find()
     db.user.drop_indexes()
     for user in users:
+        # Skip users that don't need any changes
+        if not user_upgrade_needed(user):
+            continue
+        # Check if the user is an admin (superuser), remove it from the
+        # "admin" user group if it is
+        is_superuser = False
+        new_groups_list = [g for g in user['profile']['groups']
+                           if g['name'] != 'admin']
+        if len(new_groups_list) != len(user['profile']['groups']):
+            is_superuser = True
+            user['profile']['groups'] = new_groups_list
+        # User update
         db.user.replace_one(
             {
                 "_id": user['_id']
@@ -23,14 +57,64 @@ def upgrade(db: "pymongo.database.Database"):
                 "_id": user['_id'],
                 "email": user['profile']['email'],
                 "hashed_password": user['profile']['hashed_password'],
-                "is_active": True,
-                "is_superuser": False,
+                "is_active": user['active'],
+                "is_superuser": is_superuser,
                 "is_verified": False,
                 "username": user['profile']['username'],
                 "groups": user['profile']['groups']
             },
         )
+    # Sanity check: check if there are any old-format users in the
+    # "admin" group. Remove the group if there aren't any
+    remaining_admins = db.user.count(
+        {
+            "groups": {
+                "$elemMatch": {"name": "admin"}
+            }
+        }
+    )
+    if remaining_admins == 0:
+        db.usergroup.delete_one({"name": "admin"})
+    else:
+        print("Some old 'admin' users still remain")
 
 
 def downgrade(db: "pymongo.database.Database"):
-    pass
+    superusers = db.user.find({'is_superuser': True})
+    if superusers:
+        # Create the 'admin' group if it doesn't exist
+        db.usergroup.update_one(
+            {'name': 'admin'},
+            {'$setOnInsert': {'name': 'admin'}},
+            upsert=True
+        )
+        admin_group = db.usergroup.find_one({'name': 'admin'})
+
+    users = db.user.find()
+    db.user.drop_indexes()
+    for user in users:
+        # Skip users that weren't migrated (unlikely corner case)
+        if user_upgrade_needed(user):
+            continue
+        if user.get('is_superuser') == True:
+            # Add user to admin group
+            new_groups_list = [g for g in user['groups']
+                               if g['name'] != 'admin']
+            new_groups_list.append(admin_group)
+            user['groups'] = new_groups_list
+
+        db.user.replace_one(
+            {
+                '_id': user['_id'],
+            },
+            {
+                '_id': user['_id'],
+                'active': user['is_active'],
+                'profile': {
+                    'email': user['email'],
+                    'hashed_password': user['hashed_password'],
+                    'username': user['username'],
+                    'groups': user['groups'],
+                }
+            }
+        )