Revamp merging.py module (#604)

ekassos · web-flow · commit fc6dc56cbb09 · 2024-11-08T14:25:56.000-05:00
diff --git a/pingpong/__main__.py b/pingpong/__main__.py
@@ -11,6 +11,17 @@
 from sqlalchemy import text
 from sqlalchemy.ext.asyncio import create_async_engine
 
+from pingpong.merge import (
+    get_merged_user_tuples,
+    list_all_permissions,
+    merge_missing_assistant_permissions,
+    merge_missing_class_file_permissions,
+    merge_missing_thread_permissions,
+    merge_missing_user_file_permissions,
+    merge_permissions,
+    merge,
+)
+
 from .auth import encode_auth_token
 from .bg import get_server
 from .canvas import canvas_sync_all
@@ -118,6 +129,79 @@ async def _get_or_create(email) -> int:
     webbrowser.open(url)
 
 
+# This command lists all explicitly granted permissions for a user
+@auth.command("list_permissions")
+@click.argument("user_id", type=int)
+def list_permissions(user_id: int) -> None:
+    async def _list_permissions() -> None:
+        await config.authz.driver.init()
+        async with config.authz.driver.get_client() as c:
+            perms = await list_all_permissions(c, user_id)
+            logging.info(f"Permissions for user {user_id}: {perms}")
+
+    asyncio.run(_list_permissions())
+
+
+# This command attempts to merge any outstanding permissions
+# from one user to another based on the users_merged_users table
+@auth.command("redo_permission_merges")
+def users_merge_permissions() -> None:
+    async def _users_merge_permissions() -> None:
+        await config.authz.driver.init()
+        async with config.db.driver.async_session() as session:
+            async with config.authz.driver.get_client() as c:
+                logger.info("Merging permissions for all users...")
+                async for row in get_merged_user_tuples(session):
+                    logging.info(
+                        f"Merging permissions for {row.merged_user_id} into {row.current_user_id}"
+                    )
+                    await merge_permissions(c, row.current_user_id, row.merged_user_id)
+
+    asyncio.run(_users_merge_permissions())
+
+
+# This command attempts to recover any missing permissions for a user
+# after a user(s) has/have been merged into said user. This command uses
+# fields in the database to infer which permissions the user should have
+@auth.command("add_missing_permissions")
+@click.argument("new_user_id", type=int)
+def add_missing_permissions(new_user_id: int) -> None:
+    async def _add_missing_permissions() -> None:
+        await config.authz.driver.init()
+        async with config.db.driver.async_session() as session:
+            async with config.authz.driver.get_client() as c:
+                logger.info(f"Adding missing permissions for user {new_user_id}...")
+                logger.info("Merging assistant permissions...")
+                await merge_missing_assistant_permissions(c, session, new_user_id)
+                logger.info("Merging thread permissions...")
+                await merge_missing_thread_permissions(c, session, new_user_id)
+                logger.info("Merging user file permissions...")
+                await merge_missing_user_file_permissions(c, session, new_user_id)
+                logger.info("Merging class file permissions...")
+                await merge_missing_class_file_permissions(c, session, new_user_id)
+                logger.info("Done!")
+
+    asyncio.run(_add_missing_permissions())
+
+
+# This command attempts to merge all permissions from old_user_id to new_user_id.
+# This command can be used if a user has been merged into another user
+# and some permissions were not transferred over, or the tuple was not added in users_merged_users.
+# In other words, it can be used with `old_user_id`s of users who have already been deleted.
+@auth.command("merge_users")
+@click.argument("new_user_id", type=int)
+@click.argument("old_user_id", type=int)
+def merge_users(new_user_id: int, old_user_id: int) -> None:
+    async def _merge_users() -> None:
+        await config.authz.driver.init()
+        async with config.db.driver.async_session() as session:
+            async with config.authz.driver.get_client() as c:
+                await merge(session, c, new_user_id, old_user_id)
+            await session.commit()
+
+    asyncio.run(_merge_users())
+
+
 def _load_alembic(alembic_config="alembic.ini") -> alembic.config.Config:
     """Load the Alembic config."""
     al_cfg = alembic.config.Config(alembic_config)
diff --git a/pingpong/authz/openfga.py b/pingpong/authz/openfga.py
@@ -78,16 +78,29 @@ async def list_entities(self, target: str, relation: str, type_: str) -> List[in
         return [int(user.object.id) for user in response.users]
 
     async def read(self, key: ReadRequestTupleKey) -> List[ClientTuple]:
-        response: ReadResponse = await self._cli.read(key)
+        client_tuples = []
+        continuation_token = None
 
-        client_tuples = [
-            ClientTuple(
-                user=tuple.key.user,
-                relation=tuple.key.relation,
-                object=tuple.key.object,
+        while True:
+            response: ReadResponse = await self._cli.read(
+                key, {"continuation_token": continuation_token}
             )
-            for tuple in response.tuples
-        ]
+
+            client_tuples.extend(
+                [
+                    ClientTuple(
+                        user=tuple.key.user,
+                        relation=tuple.key.relation,
+                        object=tuple.key.object,
+                    )
+                    for tuple in response.tuples
+                ]
+            )
+
+            if not response.continuation_token:
+                break
+
+            continuation_token = response.continuation_token
 
         return client_tuples
 
diff --git a/pingpong/merge.py b/pingpong/merge.py
@@ -1,6 +1,7 @@
 import asyncio
 import json
-from sqlalchemy import delete, text, update
+from typing import AsyncGenerator
+from sqlalchemy import Select, and_, delete, select, text, update
 from sqlalchemy.ext.asyncio import AsyncSession
 import logging
 
@@ -14,10 +15,12 @@
     User,
     UserClassRole,
     UserInstitutionRole,
+    _get_upsert_stmt,
     user_thread_association,
     user_merge_association,
     File,
 )
+from pingpong.schemas import MergedUserTuple
 
 logger = logging.getLogger(__name__)
 
@@ -28,19 +31,53 @@ async def merge(
     new_user_id: int,
     old_user_id: int,
 ) -> "User":
-    await asyncio.gather(
-        merge_classes(session, new_user_id, old_user_id),
-        merge_institutions(session, new_user_id, old_user_id),
-        merge_assistants(session, new_user_id, old_user_id),
-        merge_threads(session, new_user_id, old_user_id),
-        merge_lms_users(session, new_user_id, old_user_id),
-        merge_external_logins(session, new_user_id, old_user_id),
-        merge_user_files(session, new_user_id, old_user_id),
-        merge_permissions(client, new_user_id, old_user_id),
-    )
+    await merge_db_operations(session, new_user_id, old_user_id)
+    await merge_permissions(client, new_user_id, old_user_id)
     return await merge_users(session, new_user_id, old_user_id)
 
 
+async def merge_db_operations(
+    session: AsyncSession,
+    new_user_id: int,
+    old_user_id: int,
+):
+    await merge_classes(session, new_user_id, old_user_id)
+    await merge_institutions(session, new_user_id, old_user_id)
+    await merge_assistants(session, new_user_id, old_user_id)
+    await merge_threads(session, new_user_id, old_user_id)
+    await merge_lms_users(session, new_user_id, old_user_id)
+    await merge_external_logins(session, new_user_id, old_user_id)
+    await merge_user_files(session, new_user_id, old_user_id)
+
+
+async def merge_missing_permissions(
+    session: AsyncSession,
+    client: OpenFgaAuthzClient,
+    stmt: Select,
+    obj_type: str,
+    rel: str,
+    new_user_id: int,
+) -> None:
+    result = await session.execute(stmt)
+    grants: list[Relation] = []
+    revokes: list[Relation] = []
+
+    async def process_row(row):
+        old_user_ids = await client.list_entities(f"{obj_type}:{row[0]}", rel, "user")
+
+        grants.extend([(f"user:{new_user_id}", rel, f"{obj_type}:{row[0]}")])
+        revokes.extend(
+            [
+                (f"user:{old_id}", rel, f"{obj_type}:{row[0]}")
+                for old_id in old_user_ids
+                if old_id != new_user_id
+            ]
+        )
+
+    await asyncio.gather(*(process_row(row) for row in result))
+    await client.write_safe(grant=grants, revoke=revokes)
+
+
 async def merge_classes(
     session: AsyncSession, new_user_id: int, old_user_id: int
 ) -> None:
@@ -95,6 +132,15 @@ async def merge_assistants(
     await session.execute(stmt)
 
 
+async def merge_missing_assistant_permissions(
+    client: OpenFgaAuthzClient, session: AsyncSession, new_user_id: int
+) -> None:
+    stmt = select(Assistant.id).where(Assistant.creator_id == new_user_id)
+    await merge_missing_permissions(
+        session, client, stmt, "assistant", "owner", new_user_id
+    )
+
+
 async def merge_threads(
     session: AsyncSession, new_user_id: int, old_user_id: int
 ) -> None:
@@ -106,6 +152,18 @@ async def merge_threads(
     await session.execute(stmt)
 
 
+async def merge_missing_thread_permissions(
+    client: OpenFgaAuthzClient, session: AsyncSession, new_user_id: int
+) -> None:
+    # Working assumption is that only one persion is associated with a thread, as we currently don't have multiparty threads
+    stmt = select(user_thread_association.c.thread_id).where(
+        user_thread_association.c.user_id == new_user_id
+    )
+    await merge_missing_permissions(
+        session, client, stmt, "thread", "party", new_user_id
+    )
+
+
 async def merge_lms_users(
     session: AsyncSession, new_user_id: int, old_user_id: int
 ) -> None:
@@ -148,6 +206,28 @@ async def merge_user_files(
     await session.execute(stmt)
 
 
+async def merge_missing_user_file_permissions(
+    client: OpenFgaAuthzClient, session: AsyncSession, new_user_id: int
+) -> None:
+    stmt = select(File.id).where(
+        and_(File.uploader_id == new_user_id, File.private.is_(True))
+    )
+    await merge_missing_permissions(
+        session, client, stmt, "user_file", "owner", new_user_id
+    )
+
+
+async def merge_missing_class_file_permissions(
+    client: OpenFgaAuthzClient, session: AsyncSession, new_user_id: int
+) -> None:
+    stmt = select(File.id).where(
+        and_(File.uploader_id == new_user_id, File.private.is_(False))
+    )
+    await merge_missing_permissions(
+        session, client, stmt, "class_file", "owner", new_user_id
+    )
+
+
 def get_types() -> list[str]:
     """Get a list of object types used in the authz model."""
     with open(config.authz.driver.model_config) as f:
@@ -180,14 +260,23 @@ async def list_all_permissions(
     return all_relations
 
 
+async def get_merged_user_tuples(
+    session: AsyncSession,
+) -> AsyncGenerator[MergedUserTuple, None]:
+    stmt = select(
+        user_merge_association.c.user_id, user_merge_association.c.merged_user_id
+    )
+    result = await session.execute(stmt)
+    for row in result:
+        yield MergedUserTuple(current_user_id=row[0], merged_user_id=row[1])
+
+
 async def merge_permissions(
     client: OpenFgaAuthzClient, new_user_id: int, old_user_id: int
 ) -> None:
     logging.info(f"Merging permissions for {old_user_id} into {new_user_id}")
     old_permissions = await list_all_permissions(client, old_user_id)
     new_permissions = [(f"user:{new_user_id}", r, o) for _, r, o in old_permissions]
-    logging.info(f"Revoking {old_permissions}")
-    logging.info(f"Granting {new_permissions}")
     await client.write_safe(grant=new_permissions, revoke=old_permissions)
 
 
@@ -196,33 +285,42 @@ async def merge_users(
 ) -> "User":
     old_user = await User.get_by_id(session, old_user_id)
     new_user = await User.get_by_id(session, new_user_id)
+    if not new_user:
+        raise ValueError(f"New user {new_user_id} not found.")
     logging.info(f"Merging user {old_user_id} into {new_user_id}")
-    logging.info(
-        f"Old user: {old_user.id}, {old_user.email}, {old_user.state}, {'Super admin' if old_user.super_admin else 'Not super admin'}"
-    )
-    match old_user.state:
-        case "verified":
-            new_user.state = (
-                "verified" if new_user.state != "banned" else new_user.state
-            )
-        case "banned":
-            new_user.state = "banned"
-        case _:
-            pass
-
-    new_user.super_admin = new_user.super_admin or old_user.super_admin
-    stmt = (
-        update(user_merge_association)
-        .where(user_merge_association.c.user_id == old_user_id)
-        .values(user_id=new_user_id)
+    if not old_user:
+        logging.warning(
+            f"Old user {old_user_id} not found, continuing with adding the merge tuple only."
+        )
+    if old_user:
+        match old_user.state:
+            case "verified":
+                new_user.state = (
+                    "verified" if new_user.state != "banned" else new_user.state
+                )
+            case "banned":
+                new_user.state = "banned"
+            case _:
+                pass
+
+        new_user.super_admin = new_user.super_admin or old_user.super_admin
+        update_merged_account_tuple_stmt = (
+            update(user_merge_association)
+            .where(user_merge_association.c.user_id == old_user_id)
+            .values(user_id=new_user_id)
+        )
+        await session.execute(update_merged_account_tuple_stmt)
+        delete_old_user_stmt = delete(User).where(User.id == old_user_id)
+        await session.execute(delete_old_user_stmt)
+    add_new_merge_tuple_stmt = (
+        _get_upsert_stmt(session)(user_merge_association)
+        .values(user_id=new_user_id, merged_user_id=old_user_id)
+        .on_conflict_do_nothing(
+            index_elements=["user_id", "merged_user_id"],
+        )
     )
-    await session.execute(stmt)
-    stmt_ = delete(User).where(User.id == old_user_id)
-    await session.execute(stmt_)
+    await session.execute(add_new_merge_tuple_stmt)
     session.add(new_user)
     await session.flush()
     await session.refresh(new_user)
-    logging.info(
-        f"New user: {new_user.id}, {new_user.email}, {new_user.state}, {'Super admin' if new_user.super_admin else 'Not super admin'}"
-    )
     return new_user
diff --git a/pingpong/schemas.py b/pingpong/schemas.py
@@ -136,6 +136,11 @@ def has_real_name(self) -> bool:
         return bool(self.display_name or self.first_name or self.last_name)
 
 
+class MergedUserTuple(BaseModel):
+    current_user_id: int
+    merged_user_id: int
+
+
 class User(BaseModel, UserNameMixin):
     id: int
     state: UserState
