From b15b04f8fb9cb1648431dbfc6c9a24323cc4a8aa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bj=C3=B6rn=20Petersen?= <git@bjoernpetersen.net>
Date: Sun, 15 Dec 2024 23:19:02 +0100
Subject: [PATCH] Add group import

---
 chart/values.yaml   |  5 +++++
 src/bot/__main__.py | 13 +++++++++++++
 src/bot/dynamo.py   | 12 ++++++++++++
 3 files changed, 30 insertions(+)

diff --git a/chart/values.yaml b/chart/values.yaml
index 3b518a0..471a5b8 100644
--- a/chart/values.yaml
+++ b/chart/values.yaml
@@ -12,8 +12,13 @@ enabledChats:
 crons:
   - schedule: "0 13 * * *"
     command: send-polls
+    suspend:
   - schedule: "0 0 * * *"
     command: close-polls
+    suspend: true
   - schedule: "0 0 * * *"
     command: import-users
     suspend: true
+  - schedule: "0 0 * * *"
+    command: import-user-groups
+    suspend: true
diff --git a/src/bot/__main__.py b/src/bot/__main__.py
index a32e3c4..5f467ba 100644
--- a/src/bot/__main__.py
+++ b/src/bot/__main__.py
@@ -37,6 +37,16 @@ async def _import_users(dynamo: DynamoClient, database: Database) -> None:
         await database.close()
 
 
+async def _import_user_groups(dynamo: DynamoClient, database: Database) -> None:
+    await database.open()
+    try:
+        async for user_id, group_id in dynamo.list_user_groups():
+            _logger.info("Linking user %d to group %d", user_id, group_id)
+            await database.add_to_group(user_id=user_id, group_id=group_id)
+    finally:
+        await database.close()
+
+
 def main() -> None:
     config, database = initialize()
 
@@ -60,6 +70,9 @@ def main() -> None:
         case "import-users":
             _logger.info("Importing users")
             asyncio.run(_import_users(DynamoClient(config.aws), database))
+        case "import-user-groups":
+            _logger.info("Importing user groups")
+            asyncio.run(_import_user_groups(DynamoClient(config.aws), database))
         case other:
             _logger.error("Unknown operation mode: %s", other)
             sys.exit(1)
diff --git a/src/bot/dynamo.py b/src/bot/dynamo.py
index 093232d..9c6561f 100644
--- a/src/bot/dynamo.py
+++ b/src/bot/dynamo.py
@@ -39,3 +39,15 @@ async def list_users(self) -> AsyncIterable[User]:
                 response = await table.scan(ExclusiveStartKey=last_key)
                 for item in response.get("Items"):
                     yield self._build_user(item)
+
+    async def list_user_groups(self) -> AsyncIterable[tuple[int, int]]:
+        async with self._session.resource("dynamodb") as resource:
+            table = await resource.Table("mischebot-groups")
+            response = await table.scan()
+            for item in response.get("Items"):
+                yield int(item["user_id"]), int(item["group_id"])
+
+            while last_key := response.get("LastEvaluatedKey"):
+                response = await table.scan(ExclusiveStartKey=last_key)
+                for item in response.get("Items"):
+                    yield int(item["user_id"]), int(item["group_id"])