Adds support for preferring local replication

Author: justin-russell

fractal_database_matrix/controllers/replicate.py

@@ -21,6 +21,7 @@ def to(
         registration_token: str,
         confirm: bool = False,
         set_as_origin: bool = False,
+        local_url: Optional[str] = None,
         **kwargs,
     ):
         """
@@ -30,6 +31,7 @@ def to(
             registration_token: Registration token for the homeserver. Necessary for registering your devices with this homeserver.
             confirm: Consent to replicating your data to the provided homeserver.
             set_as_origin: Set the homeserver as your current database's origin.
+            local_url: Local URL for the homeserver (prefer using this url for faster replication).
         """
         if not confirm:
             res = input(
@@ -67,7 +69,7 @@ def to(
         try:
             homeserver = MatrixHomeserver.objects.get(url=homeserver_url)
             if not homeserver.replication_enabled:
-                homeserver.update(replication_enabled=True)
+                homeserver.update(replication_enabled=True, local_url=local_url)
         except MatrixHomeserver.DoesNotExist:
             with transaction.atomic():
                 # create the homeserver
@@ -78,6 +80,7 @@ def to(
                     registration_token=registration_token,
                     parent_db=current_database,
                     replication_enabled=True,
+                    local_url=local_url,
                 )
                 current_device.add_membership(homeserver)
         else:

fractal_database_matrix/migrations/0002_matrixhomeserver_local_url.py

@@ -0,0 +1,17 @@
+import fractal_database.fields
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('fractal_database_matrix', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='matrixhomeserver',
+            name='local_url',
+            field=fractal_database.fields.LocalURLField(blank=True, null=True),
+        ),
+    ]

fractal_database_matrix/models.py

@@ -10,7 +10,10 @@
 from asgiref.sync import sync_to_async
 from django.conf import settings
 from django.db import models, transaction
+from docker.errors import NotFound
+from docker.models.networks import Network
 from fractal.cli.controllers.auth import AuthenticatedController
+from fractal_database.fields import LocalURLField
 from fractal_database.models import (
     BaseModel,
     Database,
@@ -46,6 +49,7 @@
 class MatrixHomeserver(Service):
     SYNAPSE_COMPOSE_FILE_PATH = f"{fractal_database_matrix.__path__[0]}/synapse"
     SYNAPSE_LOCAL = "http://localhost:8008"
+    LOCAL_MATRIX_NETWORK = "fractal-matrix-network"
 
     credentials: models.QuerySet["MatrixCredentials"]
     gateways: models.QuerySet["Gateway"]
@@ -54,6 +58,7 @@ class MatrixHomeserver(Service):
     priority = models.PositiveIntegerField(default=0, blank=True, null=True)
     registration_token = models.CharField(max_length=255, blank=True, null=True)
     replication_enabled = models.BooleanField(default=False)
+    local_url = LocalURLField()
 
     def __str__(self) -> str:
         return f"{self.url} (MatrixHomeserver)"
@@ -153,10 +158,32 @@ def create(cls, device: Optional["Device"] = None, **ckwargs) -> "MatrixHomeserv
     def _build_images(self) -> None:
         docker_compose("build", _cwd=self.SYNAPSE_COMPOSE_FILE_PATH)
 
+    @classmethod
+    def get_docker_network(cls) -> Network:
+        client = docker.from_env()
+        return client.networks.get(cls.LOCAL_MATRIX_NETWORK)
+
+    @classmethod
+    def create_docker_network(cls) -> None:
+        """
+        Ensures that the external docker network that is specified
+        in the homeserver compose file exists.
+        """
+        client = docker.from_env()
+        try:
+            cls.get_docker_network()
+        except NotFound:
+            logger.info("Creating local matrix network %s" % cls.LOCAL_MATRIX_NETWORK)
+            client.networks.create(cls.LOCAL_MATRIX_NETWORK)
+
     def _render_compose_file(self) -> str:
         """ """
         self._build_images()
 
+        if self.local_url:
+            # ensure that external matrix network is created
+            self.create_docker_network()
+
         # FIXME: handle multiple links
         if not hasattr(self.config, "links") or not hasattr(self, "gateways"):
             raise Exception("No links or gateways found for MatrixHomeserver %s" % self)

fractal_database_matrix/operations.py

@@ -44,6 +44,15 @@ async def put_state(
 
         access_token, homeserver_url, _ = creds
 
+        if homeserver_url != channel.homeserver.url:
+            raise Exception(
+                f"You are logged into the wrong homeserver ({homeserver_url}). You must be logged into the homeserver {channel.homeserver.url}"
+            )
+        else:
+            # credentials will work with the local URL as well
+            # prefer the local URL if it exists
+            homeserver_url = channel.homeserver.local_url or homeserver_url
+
         async with MatrixClient(homeserver_url, access_token, max_timeouts=15) as client:
             res = await client.room_put_state(
                 room_id,
@@ -81,6 +90,13 @@ async def create_room(
 
         access_token, homeserver_url, _ = creds
 
+        if homeserver_url != channel.homeserver.url:
+            raise Exception("You must be logged into the correct homeserver")
+        else:
+            # credentials will work with the local URL as well
+            # prefer the local URL if it exists
+            homeserver_url = channel.homeserver.local_url or homeserver_url
+
         # verify that matrix IDs passed in invite are all lowercase
         if invite:
             if not any([matrix_id.split("@")[1].islower() for matrix_id in invite]):
@@ -117,6 +133,13 @@ async def add_subspace(
 
         access_token, homeserver_url, _ = creds
 
+        if homeserver_url != channel.homeserver.url:
+            raise Exception("You must be logged into the correct homeserver")
+        else:
+            # credentials will work with the local URL as well
+            # prefer the local URL if it exists
+            homeserver_url = channel.homeserver.local_url or homeserver_url
+
         async with MatrixClient(homeserver_url, access_token, max_timeouts=15) as client:
             res = await client.room_put_state(
                 parent_room_id,
@@ -133,7 +156,10 @@ async def add_subspace(
             )
 
     async def accept_invite_as_user(
-        self, room_id: str, homeserver_url: str, matrix_creds: Optional[tuple[str, str]] = None
+        self,
+        room_id: str,
+        channel: "MatrixReplicationChannel",
+        matrix_creds: Optional[tuple[str, str]] = None,
     ):
         """
         Args:
@@ -150,16 +176,27 @@ async def accept_invite_as_user(
         else:
             user_matrix_id, access_token = matrix_creds
 
+        if homeserver_url != channel.homeserver.url:
+            raise Exception(
+                f"You are currently logged into {homeserver_url} not {channel.homeserver.url}"
+            )
+        else:
+            # credentials will work with the local URL as well
+            # prefer the local URL if it exists
+            homeserver_url = channel.homeserver.local_url or homeserver_url
+
         async with MatrixClient(
             homeserver_url=homeserver_url, access_token=access_token, max_timeouts=15
         ) as client:
             logger.info("Accepting invite for %s as %s" % (room_id, user_matrix_id))
             await client.join_room(room_id)
 
     async def accept_invite_as_device(
-        self, device_creds: "MatrixCredentials", room_id: str, homeserver_url: str
+        self, device_creds: "MatrixCredentials", room_id: str, channel: "MatrixReplicationChannel"
     ):
         device_matrix_id = device_creds.matrix_id
+        homeserver_url = channel.homeserver.local_url or channel.homeserver.url
+
         # accept invite on behalf of device
         async with MatrixClient(
             homeserver_url=homeserver_url,
@@ -169,7 +206,9 @@ async def accept_invite_as_device(
             logger.info("Accepting invite for %s as %s" % (room_id, device_matrix_id))
             await client.join_room(room_id)
 
-    async def invite_user(self, matrix_id: str, room_id: str) -> None:
+    async def invite_user(
+        self, matrix_id: str, channel: "MatrixReplicationChannel", room_id: str
+    ) -> None:
         # FIXME: Once user has accounts on many homeservers, we need to strip the
         # host off of the room id and try to find credentials that match that host
         creds = AuthenticatedController.get_creds()
@@ -178,6 +217,15 @@ async def invite_user(self, matrix_id: str, room_id: str) -> None:
         else:
             raise Exception("You must be logged in to Matrix to invite a device")
 
+        if homeserver_url != channel.homeserver.url:
+            raise Exception(
+                f"You are currently logged into {homeserver_url} not {channel.homeserver.url}"
+            )
+        else:
+            # credentials will work with the local URL as well
+            # prefer the local URL if it exists
+            homeserver_url = channel.homeserver.local_url or homeserver_url
+
         async with MatrixClient(
             homeserver_url=homeserver_url,
             access_token=access_token,
@@ -216,7 +264,7 @@ async def register_device_account(
 
     async def set_display_name(
         self,
-        homeserver_url: str,
+        channel: "MatrixReplicationChannel",
         creds: "MatrixCredentials",
         display_name: str,
         owner_matrix_id: Optional[str] = None,
@@ -226,6 +274,8 @@ async def set_display_name(
             owner_username = owner_matrix_id.split("@")[1].split(":")[0]
             display_name = f"{owner_username}'s {display_name}"
 
+        homeserver_url = channel.homeserver.local_url or channel.homeserver.url
+
         async with MatrixClient(
             homeserver_url=homeserver_url,
             access_token=creds.access_token,
@@ -293,7 +343,7 @@ async def run(self, operation: "DurableOperation") -> dict[str, str]:
             async for account in membership.device.matrixcredentials_set.filter(
                 homeserver=channel.homeserver
             ):
-                await self.accept_invite_as_device(account, room_id, channel.homeserver.url)
+                await self.accept_invite_as_device(account, room_id, channel)
 
         logger.info("Successfully created Matrix Room for %s" % name)
         return {metadata_label: room_id}
@@ -545,7 +595,7 @@ async def run(self, operation: "DurableOperation") -> None:
             raise Exception(f"Failed to find room id in channel metadata for {metadata_label}")
 
         try:
-            await self.invite_user(device_creds.matrix_id, room_id)
+            await self.invite_user(device_creds.matrix_id, channel, room_id)
         except Exception as e:
             # if the device is already in the room, no need to accept the invite
             if "is already in the room" in str(e):
@@ -594,7 +644,7 @@ async def run(self, operation: "DurableOperation") -> None:
             % (metadata_label, room_id, device_creds.matrix_id)
         )
         # accept invite on behalf of device
-        await self.accept_invite_as_device(device_creds, room_id, channel.homeserver.url)
+        await self.accept_invite_as_device(device_creds, room_id, channel)
         logger.info("Device has successfully joined space %s for channel %s" % (room_id, channel))
 
         return None
@@ -698,9 +748,7 @@ async def run(self, operation: "DurableOperation") -> None:
             raise Exception(f"Failed to find device credentials for {membership.device}")
 
         # accept invite on behalf of device
-        await self.accept_invite_as_device(
-            device_creds, channel.device_space, channel.homeserver.url
-        )
+        await self.accept_invite_as_device(device_creds, channel.device_space, channel)
         logger.info(
             "Device has successfully joined the devices subspace for channel %s" % channel
         )
@@ -753,7 +801,7 @@ async def run(self, operation: "DurableOperation") -> None:
             raise Exception(f"Failed to find device credentials for {membership.device}")
 
         try:
-            await self.invite_user(device_creds.matrix_id, channel.device_space)
+            await self.invite_user(device_creds.matrix_id, channel, channel.device_space)
         except Exception as e:
             # if the device is already in the room, no need to accept the invite
             if "is already in the room" in str(e):
@@ -1115,7 +1163,7 @@ async def run(self, operation: DurableOperation) -> None:
             )
 
         await self.set_display_name(
-            channel.homeserver.url, device_creds, display_name, owner_matrix_id=owner_matrix_id
+            channel, device_creds, display_name, owner_matrix_id=owner_matrix_id
         )
 
 
@@ -1213,7 +1261,7 @@ async def run(self, operation: "DurableOperation") -> None:
             % (room_id_label, user_matrix_id, channel)
         )
         try:
-            await self.accept_invite_as_user(room_id, homeserver_url=channel.homeserver.url)
+            await self.accept_invite_as_user(room_id, channel=channel)
         except Exception as e:
             # if the user is already in the room, no need to accept the invite
             if "is already in the room" in str(e):
@@ -1297,10 +1345,10 @@ async def run(self, operation: "DurableOperation") -> None:
         )  # type: ignore
 
         # fetch channel in order to get room_id for the group to invite user to
-        channel: (
-            "MatrixReplicationChannel"
-        ) = await operation.channel_type.model_class().objects.aget(
-            pk=operation.channel_id
+        channel: "MatrixReplicationChannel" = (
+            await operation.channel_type.model_class()
+            .objects.select_related("homeserver")
+            .aget(pk=operation.channel_id)
         )  # type: ignore
 
         user_matrix_id = membership.user.matrix_id
@@ -1317,7 +1365,7 @@ async def run(self, operation: "DurableOperation") -> None:
             % (room_id_label, user_matrix_id, channel)
         )
         try:
-            await self.invite_user(user_matrix_id, room_id)
+            await self.invite_user(user_matrix_id, channel, room_id)
         except Exception as e:
             # if the user is already in the room, no need to accept the invite
             if "is already in the room" in str(e):

fractal_database_matrix/synapse/docker-compose.yml

@@ -43,3 +43,8 @@ services:
 
 volumes:
   synapse-db:
+
+networks:
+  default:
+    name: fractal-matrix-network
+    external: true