From 1ebf196ae92a7c4cc0ee1bb01c0d99dedf8e9848 Mon Sep 17 00:00:00 2001
From: Olivier Desenfans <desenfans.olivier@gmail.com>
Date: Mon, 30 May 2022 10:30:40 +0200
Subject: [PATCH] [Forget] Forget aggregates

Added support for the new `aggregates` field in FORGET messages.
Using this field, users can now specify one or more aggregates
to forget. This will lead to the deletion of all the messages
linked to these aggregates.

For now, the implementation simply lists all the messages stored
on the node and deletes them.
---
 setup.cfg                    |  2 +-
 src/aleph/handlers/forget.py | 19 +++++++++++++++----
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/setup.cfg b/setup.cfg
index 21000d007..07ab13179 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -38,7 +38,7 @@ install_requires =
     aiohttp==3.8.1
     aioipfs@git+https://github.com/aleph-im/aioipfs.git@76d5624661e879a13b70f3ea87dc9c9604c7eda7
     aleph-client==0.4.6
-    aleph-message==0.1.18
+    aleph-message==0.1.19
     coincurve==15.0.1
     configmanager==1.35.1
     configparser==5.2.0
diff --git a/src/aleph/handlers/forget.py b/src/aleph/handlers/forget.py
index 5daa9faaa..12d56c9d3 100644
--- a/src/aleph/handlers/forget.py
+++ b/src/aleph/handlers/forget.py
@@ -16,6 +16,9 @@
 from aleph.utils import item_type_from_hash
 
 
+logger = logging.getLogger(__name__)
+
+
 @dataclass
 class TargetMessageInfo:
     item_hash: str
@@ -45,9 +48,6 @@ def from_db_object(cls, message_dict: Dict) -> TargetMessageInfo:
         )
 
 
-logger = logging.getLogger(__name__)
-
-
 async def count_file_references(storage_hash: str) -> int:
     """Count the number of references to a file on Aleph."""
     logger.debug(f"Counting documents for {storage_hash}")
@@ -214,8 +214,19 @@ async def handle_forget_message(message: Dict, content: Dict):
     # Parsing and validation
     forget_message = ForgetMessage(**message, content=content)
     logger.debug(f"Handling forget message {forget_message.item_hash}")
+    hashes_to_forget = forget_message.content.hashes
+
+    for target_aggregate in forget_message.content.aggregates:
+        aggregate_hashes = (
+            message["item_hash"]
+            async for message in Message.collection.find(
+                {"type": MessageType.aggregate.value, "content.key": target_aggregate},
+                projection={"_id": 0, "item_hash": 1},
+            )
+        )
+        hashes_to_forget.extend(aggregate_hashes)
 
-    for target_hash in forget_message.content.hashes:
+    for target_hash in hashes_to_forget:
         target_info = await get_target_message_info(target_hash)
 
         if target_info is None: