Skip to content

Commit

Permalink
migrations: fix existing migrations scripts
Browse files Browse the repository at this point in the history
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 <[email protected]>
  • Loading branch information
r-c-n committed Jan 5, 2024
1 parent cbae5f6 commit e1129fb
Show file tree
Hide file tree
Showing 2 changed files with 89 additions and 38 deletions.
33 changes: 0 additions & 33 deletions migrations/20221014061654_set_timeout.py

This file was deleted.

94 changes: 89 additions & 5 deletions migrations/20231102101356_user.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,51 @@
# Author: Jeny Sadadia <[email protected]>

"""
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']
Expand All @@ -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'],
}
}
)

0 comments on commit e1129fb

Please sign in to comment.