Add option to disable notifications during initial backfill

mautrix/__init__.py

@@ -1,3 +1,3 @@
-__version__ = "0.5.2"
+__version__ = "0.5.3"
 __author__ = "Tulir Asokan <[email protected]>"
 __all__ = ["api", "appservice", "bridge", "client", "errors", "util", "types"]

mautrix/bridge/__init__.py

@@ -5,3 +5,4 @@
 from .user import BaseUser
 from .puppet import BasePuppet
 from .bridge import Bridge
+from .notification_disabler import NotificationDisabler

mautrix/bridge/notification_disabler.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2020 Tulir Asokan
+#
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+from typing import Optional, Type
+import logging
+
+from mautrix.types import RoomID
+from mautrix.appservice import IntentAPI
+from mautrix.api import Method, Path, PathBuilder
+from mautrix.util.logging import TraceLogger
+
+from .puppet import BasePuppet
+from .user import BaseUser
+
+
+class NotificationDisabler:
+    puppet_cls: Type[BasePuppet]
+    config_enabled: bool = False
+    log: TraceLogger = logging.getLogger("mau.notification_disabler")
+
+    room_id: RoomID
+    intent: Optional[IntentAPI]
+    enabled: bool
+
+    def __init__(self, room_id: RoomID, user: BaseUser) -> None:
+        self.room_id = room_id
+        self.enabled = False
+        puppet = self.puppet_cls.get_by_custom_mxid(user.mxid)
+        if puppet and puppet.is_real_user:
+            self.intent = puppet.intent
+
+    @property
+    def _path(self) -> PathBuilder:
+        return Path.pushrules["global"].override["net.maunium.silence_while_backfilling"]
+
+    @property
+    def _rule(self) -> dict:
+        return {
+            "actions": ["dont_notify"],
+            "conditions": [{
+                "kind": "event_match",
+                "key": "room_id",
+                "pattern": self.room_id,
+            }]
+        }
+
+    async def __aenter__(self) -> None:
+        if not self.intent or not self.config_enabled:
+            return
+        self.enabled = True
+        try:
+            self.log.debug(f"Disabling notifications in {self.room_id} for {self.intent.mxid}")
+            await self.intent.api.request(Method.PUT, self._path, content=self._rule)
+        except Exception:
+            self.log.warning(f"Failed to disable notifications in {self.room_id} "
+                             f"for {self.intent.mxid} while backfilling", exc_info=True)
+            raise
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        if not self.enabled:
+            return
+        try:
+            self.log.debug(f"Re-enabling notifications in {self.room_id} for {self.intent.mxid}")
+            await self.intent.api.request(Method.DELETE, self._path)
+        except Exception:
+            self.log.warning(f"Failed to re-enable notifications in {self.room_id} "
+                             f"for {self.intent.mxid} after backfilling", exc_info=True)

mautrix/bridge/puppet.py

@@ -36,3 +36,8 @@ def save(self) -> None:
     @abstractmethod
     def get_by_mxid(cls, mxid: UserID) -> 'BasePuppet':
         pass
+
+    @classmethod
+    @abstractmethod
+    def get_by_custom_mxid(cls, mxid: UserID) -> 'BasePuppet':
+        pass

mautrix/util/simple_lock.py

@@ -25,11 +25,17 @@ def __enter__(self) -> None:
         if self._future is None or self._future.done():
             self._future = self._loop.create_future()
 
+    async def __aenter__(self) -> None:
+        self.__enter__()
+
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         if self._future is not None:
             self._future.set_result(None)
             self._future = None
 
+    def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.__exit__(exc_type, exc_val, exc_tb)
+
     @property
     def locked(self) -> bool:
         return self._future is not None and not self._future.done()