Improve AuthN DB migration, address API key scopes

Author: nmaytan

tiled/authn_database/migrations/versions/d829476bc173_split_apikeys_scope_into_create_and_.py

@@ -6,9 +6,10 @@
 
 """
 from alembic import op
+from sqlalchemy.orm import lazyload
 from sqlalchemy.orm.session import Session
 
-from tiled.authn_database.orm import Role
+from tiled.authn_database.orm import APIKey, Role
 
 # revision identifiers, used by Alembic.
 revision = "d829476bc173"
@@ -18,57 +19,78 @@
 
 
 ROLES = ["admin", "user"]
-NEW_SCOPES_USER = ["create:apikeys", "revoke:apikeys", "create:node"]
-OLD_SCOPES_USER = ["apikeys", "create"]
-NEW_SCOPES_ADMIN = ["create:node"]
-OLD_SCOPES_ADMIN = ["create"]
+OLD_TO_NEW_SCOPES = {
+    "create": ["create:node"],
+    "apikeys": ["revoke:apikeys", "create:apikeys"],
+}
+NEW_TO_OLD_SCOPES = {
+    "create:node": ["create"],
+    "revoke:apikeys": ["apikeys"],
+    "create:apikeys": ["apikeys"],
+}
 
 
 def upgrade():
     """
     Add new scopes to Roles.
     Remove old scopes from Roles, if present.
+
+    Also adjust scopes for server API keys.
     """
     connection = op.get_bind()
     with Session(bind=connection) as db:
         for role_name in ROLES:
             role = db.query(Role).filter(Role.name == role_name).first()
-            scopes = role.scopes.copy()
-            if role_name == "admin":
-                NEW_SCOPES = NEW_SCOPES_ADMIN
-                OLD_SCOPES = OLD_SCOPES_ADMIN
-            else:
-                NEW_SCOPES = NEW_SCOPES_USER
-                OLD_SCOPES = OLD_SCOPES_USER
-            for scope in OLD_SCOPES:
-                if scope in scopes:
-                    scopes.remove(scope)
-            scopes.extend(NEW_SCOPES)
-            role.scopes = scopes
-            db.commit()
+            if not role:
+                raise RuntimeError(f"Expected role '{role_name}' not found in db!")
+            scopes = set(role.scopes or [])
+            for old_scope, new_scopes in OLD_TO_NEW_SCOPES.items():
+                if old_scope in scopes:
+                    scopes.discard(old_scope)
+                    scopes |= set(new_scopes)
+            role.scopes = list(scopes)
+
+        for apikey in (
+            db.query(APIKey).options(lazyload(APIKey.principal)).yield_per(500)
+        ):
+            scopes = set(apikey.scopes)
+            for old_scope, new_scopes in OLD_TO_NEW_SCOPES.items():
+                if old_scope in scopes:
+                    scopes.discard(old_scope)
+                    scopes |= set(new_scopes)
+            apikey.scopes = list(scopes)
+
+        db.commit()
 
 
 def downgrade():
     """
     Remove new scopes from Roles, if present.
     Add old scopes to Roles, if not preesent.
+
+    Also adjust scopes for server API keys.
     """
     connection = op.get_bind()
     with Session(bind=connection) as db:
         for role_name in ROLES:
             role = db.query(Role).filter(Role.name == role_name).first()
-            scopes = role.scopes.copy()
-            if role_name == "admin":
-                NEW_SCOPES = NEW_SCOPES_ADMIN
-                OLD_SCOPES = OLD_SCOPES_ADMIN
-            else:
-                NEW_SCOPES = NEW_SCOPES_USER
-                OLD_SCOPES = OLD_SCOPES_USER
-            for scope in NEW_SCOPES:
-                if scope in scopes:
-                    scopes.remove(scope)
-            for scope in OLD_SCOPES:
-                if scope not in scopes:
-                    scopes.append(scope)
-            role.scopes = scopes
-            db.commit()
+            if not role:
+                raise RuntimeError(f"Expected role '{role_name}' not found in db!")
+            scopes = set(role.scopes or [])
+            for new_scope, old_scopes in NEW_TO_OLD_SCOPES.items():
+                if new_scope in scopes:
+                    scopes.discard(new_scope)
+                    scopes |= set(old_scopes)
+            role.scopes = list(scopes)
+
+        for apikey in (
+            db.query(APIKey).options(lazyload(APIKey.principal)).yield_per(500)
+        ):
+            scopes = set(apikey.scopes)
+            for new_scope, old_scopes in NEW_TO_OLD_SCOPES.items():
+                if new_scope in scopes:
+                    scopes.discard(new_scope)
+                    scopes |= set(old_scopes)
+            apikey.scopes = list(scopes)
+
+        db.commit()