From 935af0da380f39ba284b78054270331bdbad7712 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Mon, 20 May 2019 10:13:05 +0100
Subject: [PATCH] Correctly update aggregation counts after redaction

---
 synapse/storage/events.py                    |  3 ++
 synapse/storage/relations.py                 | 17 +++++++++
 tests/rest/client/v2_alpha/test_relations.py | 37 ++++++++++++++++++++
 3 files changed, 57 insertions(+)

diff --git a/synapse/storage/events.py b/synapse/storage/events.py
index b025ebc926..881d6d0126 100644
--- a/synapse/storage/events.py
+++ b/synapse/storage/events.py
@@ -1325,6 +1325,9 @@ class EventsStore(
                     txn, event.room_id, event.redacts
                 )
 
+                # Remove from relations table.
+                self._handle_redaction(txn, event.redacts)
+
         # Update the event_forward_extremities, event_backward_extremities and
         # event_edges tables.
         self._handle_mult_prev_events(
diff --git a/synapse/storage/relations.py b/synapse/storage/relations.py
index 6e216066ab..63e6185ee3 100644
--- a/synapse/storage/relations.py
+++ b/synapse/storage/relations.py
@@ -415,3 +415,20 @@ class RelationsStore(RelationsWorkerStore):
 
         if rel_type == RelationTypes.REPLACES:
             txn.call_after(self.get_applicable_edit.invalidate, (parent_id,))
+
+    def _handle_redaction(self, txn, redacted_event_id):
+        """Handles receiving a redaction and checking whether we need to remove
+        any redacted relations from the database.
+
+        Args:
+            txn
+            redacted_event_id (str): The event that was redacted.
+        """
+
+        self._simple_delete_txn(
+            txn,
+            table="event_relations",
+            keyvalues={
+                "event_id": redacted_event_id,
+            }
+        )
diff --git a/tests/rest/client/v2_alpha/test_relations.py b/tests/rest/client/v2_alpha/test_relations.py
index cd965167f8..3737cdc396 100644
--- a/tests/rest/client/v2_alpha/test_relations.py
+++ b/tests/rest/client/v2_alpha/test_relations.py
@@ -320,6 +320,43 @@ class RelationsTestCase(unittest.HomeserverTestCase):
             },
         )
 
+    def test_aggregation_redactions(self):
+        """Test that annotations get correctly aggregated after a redactions.
+        """
+
+        channel = self._send_relation(RelationTypes.ANNOTATION, "m.reaction", "a")
+        self.assertEquals(200, channel.code, channel.json_body)
+        to_redact_event_id = channel.json_body["event_id"]
+
+        channel = self._send_relation(
+            RelationTypes.ANNOTATION, "m.reaction", "a", access_token=self.user2_token
+        )
+        self.assertEquals(200, channel.code, channel.json_body)
+
+        # Now lets redact the 'a' reaction
+        request, channel = self.make_request(
+            "POST",
+            "/_matrix/client/r0/rooms/%s/redact/%s" % (self.room, to_redact_event_id),
+            access_token=self.user_token,
+            content={},
+        )
+        self.render(request)
+        self.assertEquals(200, channel.code, channel.json_body)
+
+        request, channel = self.make_request(
+            "GET",
+            "/_matrix/client/unstable/rooms/%s/aggregations/%s"
+            % (self.room, self.parent_id),
+            access_token=self.user_token,
+        )
+        self.render(request)
+        self.assertEquals(200, channel.code, channel.json_body)
+
+        self.assertEquals(
+            channel.json_body,
+            {"chunk": [{"type": "m.reaction", "key": "a", "count": 1}]},
+        )
+
     def test_aggregation_must_be_annotation(self):
         """Test that aggregations must be annotations.
         """