Update room column when room becomes unpublic

2 files changed, 31 insertions, 2 deletions

diff --git a/synapse/handlers/user_directory.py b/synapse/handlers/user_directory.py
index e9488ce554..0cf403f599 100644
--- a/synapse/handlers/user_directory.py
+++ b/synapse/handlers/user_directory.py
@@ -121,12 +121,13 @@ class UserDirectoyHandler(object):
 
         # TODO: Make this faster?
         rooms = yield self.store.get_rooms_for_user(user_id)
-        for room_id in rooms:
+        for j_room_id in rooms:
             is_public = yield self.store.is_room_world_readable_or_publicly_joinable(
-                room_id
+                j_room_id
             )
 
             if is_public:
+                yield self.store.update_user_in_user_dir(user_id, j_room_id)
                 return
 
         yield self.store.remove_from_user_dir(user_id)
@@ -149,6 +150,15 @@ class UserDirectoyHandler(object):
                 if change is None:
                     continue
 
+                is_public = yield self.store.is_room_world_readable_or_publicly_joinable(
+                    room_id
+                )
+
+                if change and is_public:
+                    continue
+                elif not change and not is_public:
+                    continue
+
                 users_with_profile = yield self.state.get_current_user_in_room(room_id)
                 for user_id, profile in users_with_profile.iteritems():
                     if change:
@@ -164,6 +174,15 @@ class UserDirectoyHandler(object):
                 if change is None:
                     continue
 
+                is_public = yield self.store.is_room_world_readable_or_publicly_joinable(
+                    room_id
+                )
+
+                if change and is_public:
+                    continue
+                elif not change and not is_public:
+                    continue
+
                 users_with_profile = yield self.state.get_current_user_in_room(room_id)
                 for user_id, profile in users_with_profile.iteritems():
                     if change:
diff --git a/synapse/storage/user_directory.py b/synapse/storage/user_directory.py
index 6c7c8c4bee..d72b93b585 100644
--- a/synapse/storage/user_directory.py
+++ b/synapse/storage/user_directory.py
@@ -80,6 +80,16 @@ class UserDirectoryStore(SQLBaseStore):
         )
 
     @defer.inlineCallbacks
+    def update_user_in_user_dir(self, user_id, room_id):
+        yield self._simple_update_one(
+            table="user_directory",
+            keyvalues={"user_id": user_id},
+            updatevalues={"room_id": room_id},
+            desc="update_user_in_user_dir",
+        )
+        self.get_user_in_directory.invalidate((user_id,))
+
+    @defer.inlineCallbacks
     def remove_from_user_dir(self, user_id):
         yield self._simple_delete(
             table="user_directory",