From b15b04f8fb9cb1648431dbfc6c9a24323cc4a8aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Petersen?= 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"])