wip - move support user logic into handler from storage

5 files changed, 120 insertions, 140 deletions

diff --git a/synapse/handlers/user_directory.py b/synapse/handlers/user_directory.py
index f11b430126..1e233e044f 100644
--- a/synapse/handlers/user_directory.py
+++ b/synapse/handlers/user_directory.py
@@ -125,9 +125,11 @@ class UserDirectoryHandler(object):
         """
         # FIXME(#3714): We should probably do this in the same worker as all
         # the other changes.
-        yield self.store.update_profile_in_user_dir(
-            user_id, profile.display_name, profile.avatar_url, None,
-        )
+        is_support = yield self.store.is_support_user(user_id)
+        if not is_support:
+            yield self.store.update_profile_in_user_dir(
+                user_id, profile.display_name, profile.avatar_url, None,
+            )
 
     @defer.inlineCallbacks
     def handle_user_deactivated(self, user_id):
@@ -323,18 +325,14 @@ class UserDirectoryHandler(object):
                     room_id, prev_event_id, event_id, typ,
                 )
             elif typ == EventTypes.Member:
+
                 change = yield self._get_key_change(
                     prev_event_id, event_id,
                     key_name="membership",
                     public_value=Membership.JOIN,
                 )
 
-                if change is None:
-                    # Handle any profile changes
-                    yield self._handle_profile_change(
-                        state_key, room_id, prev_event_id, event_id,
-                    )
-                    continue
+
 
                 if not change:
                     # Need to check if the server left the room entirely, if so
@@ -349,12 +347,23 @@ class UserDirectoryHandler(object):
                         # need to remove those users or not
                         user_ids = yield self.store.get_users_in_dir_due_to_room(room_id)
                         for user_id in user_ids:
-                            yield self._handle_remove_user(room_id, user_id)
+                            is_support = yield self.store.is_support_user(state_key)
+                            if not is_support:
+                                yield self._handle_remove_user(room_id, user_id)
                         return
                     else:
                         logger.debug("Server is still in room: %r", room_id)
 
-                if change:  # The user joined
+                is_support = yield self.store.is_support_user(state_key)
+
+                if change is None and not is_support:
+                    # Handle any profile changes
+                    yield self._handle_profile_change(
+                        state_key, room_id, prev_event_id, event_id,
+                    )
+                    continue
+
+                if change and not is_support:  # The user joined
                     event = yield self.store.get_event(event_id, allow_none=True)
                     profile = ProfileInfo(
                         avatar_url=event.content.get("avatar_url"),
diff --git a/synapse/storage/registration.py b/synapse/storage/registration.py
index 9f1dddb5ac..4d5ffe375e 100644
--- a/synapse/storage/registration.py
+++ b/synapse/storage/registration.py
@@ -454,6 +454,18 @@ class RegistrationStore(RegistrationWorkerStore,
 
         defer.returnValue(res if res else False)
 
+    @cachedInlineCallbacks()
+    def is_support_user(self, user_id):
+        res = yield self._simple_select_one_onecol(
+            table="users",
+            keyvalues={"name": user_id},
+            retcol="user_type",
+            allow_none=True,
+            desc="is_support_user",
+        )
+
+        defer.returnValue(True if res == UserTypes.SUPPORT else False)
+
     @defer.inlineCallbacks
     def user_add_threepid(self, user_id, medium, address, validated_at, added_at):
         yield self._simple_upsert("user_threepids", {
diff --git a/synapse/storage/user_directory.py b/synapse/storage/user_directory.py
index a77c9a0579..b8f3bfdc97 100644
--- a/synapse/storage/user_directory.py
+++ b/synapse/storage/user_directory.py
@@ -68,10 +68,6 @@ class UserDirectoryStore(SQLBaseStore):
                 or publically joinable
             user_ids (list(str)): Users to add
         """
-        for u in user_ids:
-            is_support = yield self.store.is_support_user(u)
-            if is_support:
-                user_ids.remove(u)
 
         yield self._simple_insert_many(
             table="users_in_public_rooms",
@@ -87,7 +83,7 @@ class UserDirectoryStore(SQLBaseStore):
         for user_id in user_ids:
             self.get_user_in_public_room.invalidate((user_id,))
 
-    @defer.inlineCallbacks
+
     def add_profiles_to_user_dir(self, room_id, users_with_profile):
         """Add profiles to the user directory
 
@@ -96,13 +92,6 @@ class UserDirectoryStore(SQLBaseStore):
             users_with_profile (dict): Users to add to directory in the form of
                 mapping of user_id -> ProfileInfo
         """
-        user_ids_to_pop = []
-        for user_id in iterkeys(users_with_profile):
-            is_support = yield self.store.is_support_user(user_id)
-            if is_support:
-                user_ids_to_pop.append(user_id)
-        for u in user_ids_to_pop:
-            users_with_profile.pop(u, None)
 
         if isinstance(self.database_engine, PostgresEngine):
             # We weight the loclpart most highly, then display name and finally
@@ -158,27 +147,22 @@ class UserDirectoryStore(SQLBaseStore):
                     self.get_user_in_directory.invalidate, (user_id,)
                 )
 
-        self.runInteraction(
+        return self.runInteraction(
             "add_profiles_to_user_dir", _add_profiles_to_user_dir_txn
         )
 
     @defer.inlineCallbacks
     def update_user_in_user_dir(self, user_id, room_id):
-        is_support = yield self.store.is_support_user(user_id)
-        if not is_support:
-            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,))
+        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,))
 
     def update_profile_in_user_dir(self, user_id, display_name, avatar_url, room_id):
         def _update_profile_in_user_dir_txn(txn):
-            is_support = yield self.store.is_support_user(user_id)
-            if is_support:
-                return
 
             new_entry = self._simple_upsert_txn(
                 txn,
@@ -243,15 +227,13 @@ class UserDirectoryStore(SQLBaseStore):
 
     @defer.inlineCallbacks
     def update_user_in_public_user_list(self, user_id, room_id):
-        is_support = yield self.store.is_support_user(user_id)
-        if not is_support:
-            yield self._simple_update_one(
-                table="users_in_public_rooms",
-                keyvalues={"user_id": user_id},
-                updatevalues={"room_id": room_id},
-                desc="update_user_in_public_user_list",
-            )
-            self.get_user_in_public_room.invalidate((user_id,))
+        yield self._simple_update_one(
+            table="users_in_public_rooms",
+            keyvalues={"user_id": user_id},
+            updatevalues={"room_id": room_id},
+            desc="update_user_in_public_user_list",
+        )
+        self.get_user_in_public_room.invalidate((user_id,))
 
     def remove_from_user_dir(self, user_id):
         def _remove_from_user_dir_txn(txn):
@@ -354,7 +336,6 @@ class UserDirectoryStore(SQLBaseStore):
         rows = yield self._execute("get_all_local_users", None, sql)
         defer.returnValue([name for name, in rows])
 
-    @defer.inlineCallbacks
     def add_users_who_share_room(self, room_id, share_private, user_id_tuples):
         """Insert entries into the users_who_share_rooms table. The first
         user should be a local user.
@@ -364,19 +345,8 @@ class UserDirectoryStore(SQLBaseStore):
             share_private (bool): Is the room private
             user_id_tuples([(str, str)]): iterable of 2-tuple of user IDs.
         """
-        # user_id_tuples =
-        #   [u for u in user_id_tuples if not self.store.is_support_user(u)]
-        user_id_tuples_filtered = []
-        for ut in user_id_tuples:
-            safe = True
-            for u in ut:
-                is_support = yield self.store.is_support_user(u)
-                if is_support:
-                    safe = False
-            if safe:
-                user_id_tuples_filtered.append(ut)
-
-        def _add_users_who_share_room_txn(txn, user_id_tuples_filtered):
+
+        def _add_users_who_share_room_txn(txn, user_id_tuples):
 
             self._simple_insert_many_txn(
                 txn,
@@ -388,10 +358,10 @@ class UserDirectoryStore(SQLBaseStore):
                         "room_id": room_id,
                         "share_private": share_private,
                     }
-                    for user_id, other_user_id in user_id_tuples_filtered
+                    for user_id, other_user_id in user_id_tuples
                 ],
             )
-            for user_id, other_user_id in user_id_tuples_filtered:
+            for user_id, other_user_id in user_id_tuples:
                 txn.call_after(
                     self.get_users_who_share_room_from_dir.invalidate,
                     (user_id,),
@@ -400,10 +370,10 @@ class UserDirectoryStore(SQLBaseStore):
                     self.get_if_users_share_a_room.invalidate,
                     (user_id, other_user_id),
                 )
-        self.runInteraction(
+        return self.runInteraction(
             "add_users_who_share_room",
             _add_users_who_share_room_txn,
-            user_id_tuples_filtered,
+            user_id_tuples,
         )
 
     def update_users_who_share_room(self, room_id, share_private, user_id_sets):
@@ -775,18 +745,6 @@ class UserDirectoryStore(SQLBaseStore):
             "results": results,
         })
 
-    @cachedInlineCallbacks()
-    def is_support_user(self, user_id):
-        res = yield self._simple_select_one_onecol(
-            table="users",
-            keyvalues={"name": user_id},
-            retcol="user_type",
-            allow_none=True,
-            desc="is_support_user",
-        )
-
-        defer.returnValue(True if res == UserTypes.SUPPORT else False)
-
 
 def _parse_query_sqlite(search_term):
     """Takes a plain unicode string from the user and converts it into a form
diff --git a/tests/storage/test_registration.py b/tests/storage/test_registration.py
index 3dfb7b903a..ec2032ea63 100644
--- a/tests/storage/test_registration.py
+++ b/tests/storage/test_registration.py
@@ -99,6 +99,26 @@ class RegistrationStoreTestCase(unittest.TestCase):
         user = yield self.store.get_user_by_access_token(self.tokens[0])
         self.assertIsNone(user, "access token was not deleted without device_id")
 
+    @defer.inlineCallbacks
+    def test_is_support_user(self):
+        TEST_USER = "@test:test"
+        SUPPORT_USER = "@support:test"
+
+        res = yield self.store.is_support_user(None)
+        self.assertFalse(res)
+        yield self.datastore.register(user_id=TEST_USER, token="123", password_hash=None)
+        res = yield self.store.is_support_user(TEST_USER)
+        self.assertFalse(res)
+
+        self.datastore.register(
+            user_id=SUPPORT_USER,
+            token="456",
+            password_hash=None,
+            user_type=UserTypes.SUPPORT
+        )
+        res = yield self.store.is_support_user(SUPPORT_USER)
+        self.assertTrue(res)
+
 
 class TokenGenerator:
     def __init__(self):
diff --git a/tests/storage/test_user_directory.py b/tests/storage/test_user_directory.py
index 6cf6d4a48a..04cfa8c691 100644
--- a/tests/storage/test_user_directory.py
+++ b/tests/storage/test_user_directory.py
@@ -32,8 +32,8 @@ class UserDirectoryStoreTestCase(unittest.TestCase):
     @defer.inlineCallbacks
     def setUp(self):
         self.hs = yield setup_test_homeserver(self.addCleanup)
-        self.datastore = self.hs.get_datastore()
-        self.store = UserDirectoryStore(self.hs.get_db_conn(), self.hs)
+        self.store = self.hs.get_datastore()
+        # self.store = UserDirectoryStore(self.hs.get_db_conn(), self.hs)
 
         # alice and bob are both in !room_id. bobby is not but shares
         # a homeserver with alice.
@@ -79,66 +79,47 @@ class UserDirectoryStoreTestCase(unittest.TestCase):
         finally:
             self.hs.config.user_directory_search_all_users = False
 
-    @defer.inlineCallbacks
-    def test_cannot_add_support_user_to_directory(self):
-        self.hs.config.user_directory_search_all_users = True
-        SUPPORT_USER = "@support:test"
-        SUPPOER_USER_SCREEN_NAME = "Support"
-
-        yield self.datastore.register(user_id=SUPPORT_USER, token="123",
-                                      password_hash=None, user_type=UserTypes.SUPPORT)
-        yield self.datastore.register(user_id=ALICE, token="456", password_hash=None)
-
-        yield self.store.add_profiles_to_user_dir(
-            ROOM,
-            {SUPPORT_USER: ProfileInfo(None, SUPPOER_USER_SCREEN_NAME)},
-        )
-        yield self.store.add_users_to_public_room(ROOM, [SUPPORT_USER])
-        yield self.store.add_users_who_share_room(
-            ROOM, False, ((ALICE, SUPPORT_USER),)
-        )
-
-        r = yield self.store.search_user_dir(ALICE, SUPPOER_USER_SCREEN_NAME, 10)
-        self.assertFalse(r["limited"])
-        self.assertEqual(0, len(r["results"]))
-
-        # Check that enabled support user does not prevent all users being added
-        r = yield self.store.search_user_dir(ALICE, ALICE, 10)
-        self.assertFalse(r["limited"])
-        self.assertEqual(1, len(r["results"]))
-
-        yield self.store.update_user_in_user_dir(SUPPORT_USER, ROOM)
-        yield self.store.update_profile_in_user_dir(
-            SUPPORT_USER, SUPPOER_USER_SCREEN_NAME, None, ROOM
-        )
-        yield self.store.update_user_in_public_user_list(SUPPORT_USER, ROOM)
-
-        r = yield self.store.search_user_dir(ALICE, SUPPOER_USER_SCREEN_NAME, 10)
-        self.assertFalse(r["limited"])
-        self.assertEqual(0, len(r["results"]))
-
-        r = yield self.store.get_user_in_directory(SUPPORT_USER)
-        self.assertEqual(r, None)
-
-        r = yield self.store.get_user_in_public_room(SUPPORT_USER)
-        self.assertEqual(r, None)
-
-    @defer.inlineCallbacks
-    def test_is_support_user(self):
-        TEST_USER = "@test:test"
-        SUPPORT_USER = "@support:test"
-
-        res = yield self.store.is_support_user(None)
-        self.assertFalse(res)
-        yield self.datastore.register(user_id=TEST_USER, token="123", password_hash=None)
-        res = yield self.store.is_support_user(TEST_USER)
-        self.assertFalse(res)
-
-        self.datastore.register(
-            user_id=SUPPORT_USER,
-            token="456",
-            password_hash=None,
-            user_type=UserTypes.SUPPORT
-        )
-        res = yield self.store.is_support_user(SUPPORT_USER)
-        self.assertTrue(res)
+    # @defer.inlineCallbacks
+    # def test_cannot_add_support_user_to_directory(self):
+    #     self.hs.config.user_directory_search_all_users = True
+    #     SUPPORT_USER = "@support:test"
+    #     SUPPOER_USER_SCREEN_NAME = "Support"
+    #
+    #     yield self.store.register(user_id=SUPPORT_USER, token="123",
+    #                               password_hash=None,
+    #                               user_type=UserTypes.SUPPORT)
+    #     yield self.store.register(user_id=ALICE, token="456", password_hash=None)
+    #
+    #     yield self.store.add_profiles_to_user_dir(
+    #         ROOM,
+    #         {SUPPORT_USER: ProfileInfo(None, SUPPOER_USER_SCREEN_NAME)},
+    #     )
+    #     yield self.store.add_users_to_public_room(ROOM, [SUPPORT_USER])
+    #     yield self.store.add_users_who_share_room(
+    #         ROOM, False, ((ALICE, SUPPORT_USER),)
+    #     )
+    #
+    #     r = yield self.store.search_user_dir(ALICE, SUPPOER_USER_SCREEN_NAME, 10)
+    #     self.assertFalse(r["limited"])
+    #     self.assertEqual(0, len(r["results"]))
+    #
+    #     # Check that enabled support user does not prevent all users being added
+    #     r = yield self.store.search_user_dir(ALICE, ALICE, 10)
+    #     self.assertFalse(r["limited"])
+    #     self.assertEqual(1, len(r["results"]))
+    #
+    #     yield self.store.update_user_in_user_dir(SUPPORT_USER, ROOM)
+    #     yield self.store.update_profile_in_user_dir(
+    #         SUPPORT_USER, SUPPOER_USER_SCREEN_NAME, None, ROOM
+    #     )
+    #     yield self.store.update_user_in_public_user_list(SUPPORT_USER, ROOM)
+    #
+    #     r = yield self.store.search_user_dir(ALICE, SUPPOER_USER_SCREEN_NAME, 10)
+    #     self.assertFalse(r["limited"])
+    #     self.assertEqual(0, len(r["results"]))
+    #
+    #     r = yield self.store.get_user_in_directory(SUPPORT_USER)
+    #     self.assertEqual(r, None)
+    #
+    #     r = yield self.store.get_user_in_public_room(SUPPORT_USER)
+    #     self.assertEqual(r, None)