Add a module callback to react to account data changes (#12327)

Co-authored-by: Sean Quah <8349537+squahtx@users.noreply.github.com>
Co-authored-by: Andrew Morgan <1342360+anoadragon453@users.noreply.github.com>

2 files changed, 66 insertions, 1 deletions

diff --git a/synapse/handlers/account_data.py b/synapse/handlers/account_data.py
index 177b4f8991..4af9fbc5d1 100644
--- a/synapse/handlers/account_data.py
+++ b/synapse/handlers/account_data.py
@@ -12,8 +12,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import logging
 import random
-from typing import TYPE_CHECKING, Collection, List, Optional, Tuple
+from typing import TYPE_CHECKING, Awaitable, Callable, Collection, List, Optional, Tuple
 
 from synapse.replication.http.account_data import (
     ReplicationAddTagRestServlet,
@@ -27,6 +28,12 @@ from synapse.types import JsonDict, UserID
 if TYPE_CHECKING:
     from synapse.server import HomeServer
 
+logger = logging.getLogger(__name__)
+
+ON_ACCOUNT_DATA_UPDATED_CALLBACK = Callable[
+    [str, Optional[str], str, JsonDict], Awaitable
+]
+
 
 class AccountDataHandler:
     def __init__(self, hs: "HomeServer"):
@@ -40,6 +47,44 @@ class AccountDataHandler:
         self._remove_tag_client = ReplicationRemoveTagRestServlet.make_client(hs)
         self._account_data_writers = hs.config.worker.writers.account_data
 
+        self._on_account_data_updated_callbacks: List[
+            ON_ACCOUNT_DATA_UPDATED_CALLBACK
+        ] = []
+
+    def register_module_callbacks(
+        self, on_account_data_updated: Optional[ON_ACCOUNT_DATA_UPDATED_CALLBACK] = None
+    ) -> None:
+        """Register callbacks from modules."""
+        if on_account_data_updated is not None:
+            self._on_account_data_updated_callbacks.append(on_account_data_updated)
+
+    async def _notify_modules(
+        self,
+        user_id: str,
+        room_id: Optional[str],
+        account_data_type: str,
+        content: JsonDict,
+    ) -> None:
+        """Notifies modules about new account data changes.
+
+        A change can be either a new account data type being added, or the content
+        associated with a type being changed. Account data for a given type is removed by
+        changing the associated content to an empty dictionary.
+
+        Note that this is not called when the tags associated with a room change.
+
+        Args:
+            user_id: The user whose account data is changing.
+            room_id: The ID of the room the account data change concerns, if any.
+            account_data_type: The type of the account data.
+            content: The content that is now associated with this type.
+        """
+        for callback in self._on_account_data_updated_callbacks:
+            try:
+                await callback(user_id, room_id, account_data_type, content)
+            except Exception as e:
+                logger.exception("Failed to run module callback %s: %s", callback, e)
+
     async def add_account_data_to_room(
         self, user_id: str, room_id: str, account_data_type: str, content: JsonDict
     ) -> int:
@@ -63,6 +108,8 @@ class AccountDataHandler:
                 "account_data_key", max_stream_id, users=[user_id]
             )
 
+            await self._notify_modules(user_id, room_id, account_data_type, content)
+
             return max_stream_id
         else:
             response = await self._room_data_client(
@@ -96,6 +143,9 @@ class AccountDataHandler:
             self._notifier.on_new_event(
                 "account_data_key", max_stream_id, users=[user_id]
             )
+
+            await self._notify_modules(user_id, None, account_data_type, content)
+
             return max_stream_id
         else:
             response = await self._user_data_client(
diff --git a/synapse/module_api/__init__.py b/synapse/module_api/__init__.py
index f7f95bae29..9a61593ff5 100644
--- a/synapse/module_api/__init__.py
+++ b/synapse/module_api/__init__.py
@@ -65,6 +65,7 @@ from synapse.events.third_party_rules import (
     ON_THREEPID_BIND_CALLBACK,
     ON_USER_DEACTIVATION_STATUS_CHANGED_CALLBACK,
 )
+from synapse.handlers.account_data import ON_ACCOUNT_DATA_UPDATED_CALLBACK
 from synapse.handlers.account_validity import (
     IS_USER_EXPIRED_CALLBACK,
     ON_LEGACY_ADMIN_REQUEST,
@@ -216,6 +217,7 @@ class ModuleApi:
         self._third_party_event_rules = hs.get_third_party_event_rules()
         self._password_auth_provider = hs.get_password_auth_provider()
         self._presence_router = hs.get_presence_router()
+        self._account_data_handler = hs.get_account_data_handler()
 
     #################################################################################
     # The following methods should only be called during the module's initialisation.
@@ -376,6 +378,19 @@ class ModuleApi:
                 min_batch_size=min_batch_size,
             )
 
+    def register_account_data_callbacks(
+        self,
+        *,
+        on_account_data_updated: Optional[ON_ACCOUNT_DATA_UPDATED_CALLBACK] = None,
+    ) -> None:
+        """Registers account data callbacks.
+
+        Added in Synapse 1.57.0.
+        """
+        return self._account_data_handler.register_module_callbacks(
+            on_account_data_updated=on_account_data_updated,
+        )
+
     def register_web_resource(self, path: str, resource: Resource) -> None:
         """Registers a web resource to be served at the given path.