make count_monthly_users async synapse/handlers/auth.py

4 files changed, 46 insertions, 38 deletions

diff --git a/synapse/handlers/register.py b/synapse/handlers/register.py
index f46b8355c0..cc935a5e84 100644
--- a/synapse/handlers/register.py
+++ b/synapse/handlers/register.py
@@ -144,7 +144,7 @@ class RegistrationHandler(BaseHandler):
         Raises:
             RegistrationError if there was a problem registering.
         """
-        self._check_mau_limits()
+        yield self._check_mau_limits()
         password_hash = None
         if password:
             password_hash = yield self.auth_handler().hash(password)
@@ -289,7 +289,7 @@ class RegistrationHandler(BaseHandler):
                 400,
                 "User ID can only contain characters a-z, 0-9, or '=_-./'",
             )
-        self._check_mau_limits()
+        yield self._check_mau_limits()
         user = UserID(localpart, self.hs.hostname)
         user_id = user.to_string()
 
@@ -439,7 +439,7 @@ class RegistrationHandler(BaseHandler):
         """
         if localpart is None:
             raise SynapseError(400, "Request must include user id")
-        self._check_mau_limits()
+        yield self._check_mau_limits()
         need_register = True
 
         try:
@@ -534,13 +534,14 @@ class RegistrationHandler(BaseHandler):
             action="join",
         )
 
+    @defer.inlineCallbacks
     def _check_mau_limits(self):
         """
         Do not accept registrations if monthly active user limits exceeded
          and limiting is enabled
         """
         if self.hs.config.limit_usage_by_mau is True:
-            current_mau = self.store.count_monthly_users()
+            current_mau = yield self.store.count_monthly_users()
             if current_mau >= self.hs.config.max_mau_value:
                 raise RegistrationError(
                     403, "MAU Limit Exceeded", Codes.MAU_LIMIT_EXCEEDED
diff --git a/synapse/storage/__init__.py b/synapse/storage/__init__.py
index 4747118ed7..f9682832ca 100644
--- a/synapse/storage/__init__.py
+++ b/synapse/storage/__init__.py
@@ -273,24 +273,24 @@ class DataStore(RoomMemberStore, RoomStore,
         This method should be refactored with count_daily_users - the only
         reason not to is waiting on definition of mau
         returns:
-            int: count of current monthly active users
+            defered: resolves to int
         """
+        def _count_monthly_users(txn):
+            thirty_days_ago = int(self._clock.time_msec()) - (1000 * 60 * 60 * 24 * 30)
+            sql = """
+                SELECT COALESCE(count(*), 0) FROM (
+                    SELECT user_id FROM user_ips
+                    WHERE last_seen > ?
+                    GROUP BY user_id
+                ) u
+            """
 
-        thirty_days_ago = int(self._clock.time_msec()) - (1000 * 60 * 60 * 24 * 30)
-        sql = """
-            SELECT COALESCE(count(*), 0) FROM (
-                SELECT user_id FROM user_ips
-                WHERE last_seen > ?
-                GROUP BY user_id
-            ) u
-        """
-        try:
-            txn = self.db_conn.cursor()
             txn.execute(sql, (thirty_days_ago,))
             count, = txn.fetchone()
+            print "Count is %d" % (count,)
             return count
-        finally:
-            txn.close()
+
+        return self.runInteraction("count_monthly_users", _count_monthly_users)
 
     def count_r30_users(self):
         """
diff --git a/tests/handlers/test_auth.py b/tests/handlers/test_auth.py
index e01f14a10a..440a453082 100644
--- a/tests/handlers/test_auth.py
+++ b/tests/handlers/test_auth.py
@@ -77,38 +77,37 @@ class AuthTestCase(unittest.TestCase):
         v.satisfy_general(verify_nonce)
         v.verify(macaroon, self.hs.config.macaroon_secret_key)
 
+    @defer.inlineCallbacks
     def test_short_term_login_token_gives_user_id(self):
         self.hs.clock.now = 1000
 
         token = self.macaroon_generator.generate_short_term_login_token(
             "a_user", 5000
         )
-
-        self.assertEqual(
-            "a_user",
-            self.auth_handler.validate_short_term_login_token_and_get_user_id(
-                token
-            )
+        user_id = yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
+            token
         )
+        self.assertEqual("a_user", user_id)
 
         # when we advance the clock, the token should be rejected
         self.hs.clock.now = 6000
         with self.assertRaises(synapse.api.errors.AuthError):
-            self.auth_handler.validate_short_term_login_token_and_get_user_id(
+            yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
                 token
             )
 
+    @defer.inlineCallbacks
     def test_short_term_login_token_cannot_replace_user_id(self):
         token = self.macaroon_generator.generate_short_term_login_token(
             "a_user", 5000
         )
         macaroon = pymacaroons.Macaroon.deserialize(token)
 
+        user_id = yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
+            macaroon.serialize()
+        )
         self.assertEqual(
-            "a_user",
-            self.auth_handler.validate_short_term_login_token_and_get_user_id(
-                macaroon.serialize()
-            )
+            "a_user", user_id
         )
 
         # add another "user_id" caveat, which might allow us to override the
@@ -116,7 +115,7 @@ class AuthTestCase(unittest.TestCase):
         macaroon.add_first_party_caveat("user_id = b_user")
 
         with self.assertRaises(synapse.api.errors.AuthError):
-            self.auth_handler.validate_short_term_login_token_and_get_user_id(
+            yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
                 macaroon.serialize()
             )
 
@@ -126,7 +125,7 @@ class AuthTestCase(unittest.TestCase):
         # Ensure does not throw exception
         yield self.auth_handler.get_access_token_for_user_id('user_a')
 
-        self.auth_handler.validate_short_term_login_token_and_get_user_id(
+        yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
             self._get_macaroon().serialize()
         )
 
@@ -134,24 +133,30 @@ class AuthTestCase(unittest.TestCase):
     def test_mau_limits_exceeded(self):
         self.hs.config.limit_usage_by_mau = True
         self.hs.get_datastore().count_monthly_users = Mock(
-            return_value=self.large_number_of_users
+            return_value=defer.succeed(self.large_number_of_users)
         )
+
         with self.assertRaises(AuthError):
             yield self.auth_handler.get_access_token_for_user_id('user_a')
+
+        self.hs.get_datastore().count_monthly_users = Mock(
+            return_value=defer.succeed(self.large_number_of_users)
+        )
         with self.assertRaises(AuthError):
-            self.auth_handler.validate_short_term_login_token_and_get_user_id(
+            yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
                 self._get_macaroon().serialize()
             )
 
     @defer.inlineCallbacks
     def test_mau_limits_not_exceeded(self):
         self.hs.config.limit_usage_by_mau = True
+
         self.hs.get_datastore().count_monthly_users = Mock(
-            return_value=self.small_number_of_users
+            return_value=defer.succeed(self.small_number_of_users)
         )
         # Ensure does not raise exception
         yield self.auth_handler.get_access_token_for_user_id('user_a')
-        self.auth_handler.validate_short_term_login_token_and_get_user_id(
+        yield self.auth_handler.validate_short_term_login_token_and_get_user_id(
             self._get_macaroon().serialize()
         )
 
diff --git a/tests/handlers/test_register.py b/tests/handlers/test_register.py
index a5a8e7c954..0937d71cf6 100644
--- a/tests/handlers/test_register.py
+++ b/tests/handlers/test_register.py
@@ -90,7 +90,7 @@ class RegistrationTestCase(unittest.TestCase):
         lots_of_users = 100
         small_number_users = 1
 
-        store.count_monthly_users = Mock(return_value=lots_of_users)
+        store.count_monthly_users = Mock(return_value=defer.succeed(lots_of_users))
 
         # Ensure does not throw exception
         yield self.handler.get_or_create_user(requester, 'a', display_name)
@@ -100,7 +100,7 @@ class RegistrationTestCase(unittest.TestCase):
         with self.assertRaises(RegistrationError):
             yield self.handler.get_or_create_user(requester, 'b', display_name)
 
-        store.count_monthly_users = Mock(return_value=small_number_users)
+        store.count_monthly_users = Mock(return_value=defer.succeed(small_number_users))
 
         self._macaroon_mock_generator("another_secret")
 
@@ -108,12 +108,14 @@ class RegistrationTestCase(unittest.TestCase):
         yield self.handler.get_or_create_user("@neil:matrix.org", 'c', "Neil")
 
         self._macaroon_mock_generator("another another secret")
-        store.count_monthly_users = Mock(return_value=lots_of_users)
+        store.count_monthly_users = Mock(return_value=defer.succeed(lots_of_users))
+
         with self.assertRaises(RegistrationError):
             yield self.handler.register(localpart=local_part)
 
         self._macaroon_mock_generator("another another secret")
-        store.count_monthly_users = Mock(return_value=lots_of_users)
+        store.count_monthly_users = Mock(return_value=defer.succeed(lots_of_users))
+
         with self.assertRaises(RegistrationError):
             yield self.handler.register_saml2(local_part)