insertion into monthly_active_users

5 files changed, 99 insertions, 11 deletions

diff --git a/synapse/api/auth.py b/synapse/api/auth.py
index 5bbbe8e2e7..d8022bcf8e 100644
--- a/synapse/api/auth.py
+++ b/synapse/api/auth.py
@@ -213,7 +213,7 @@ class Auth(object):
                 default=[b""]
             )[0]
             if user and access_token and ip_addr:
-                self.store.insert_client_ip(
+                yield self.store.insert_client_ip(
                     user_id=user.to_string(),
                     access_token=access_token,
                     ip=ip_addr,
diff --git a/synapse/storage/__init__.py b/synapse/storage/__init__.py
index 04ff1f8006..2120f46ed5 100644
--- a/synapse/storage/__init__.py
+++ b/synapse/storage/__init__.py
@@ -39,6 +39,7 @@ from .filtering import FilteringStore
 from .group_server import GroupServerStore
 from .keys import KeyStore
 from .media_repository import MediaRepositoryStore
+from .monthly_active_users import MonthlyActiveUsersStore
 from .openid import OpenIdStore
 from .presence import PresenceStore, UserPresenceState
 from .profile import ProfileStore
@@ -87,6 +88,7 @@ class DataStore(RoomMemberStore, RoomStore,
                 UserDirectoryStore,
                 GroupServerStore,
                 UserErasureStore,
+                MonthlyActiveUsersStore,
                 ):
 
     def __init__(self, db_conn, hs):
diff --git a/synapse/storage/client_ips.py b/synapse/storage/client_ips.py
index b8cefd43d6..506915a1ef 100644
--- a/synapse/storage/client_ips.py
+++ b/synapse/storage/client_ips.py
@@ -35,6 +35,7 @@ LAST_SEEN_GRANULARITY = 120 * 1000
 
 class ClientIpStore(background_updates.BackgroundUpdateStore):
     def __init__(self, db_conn, hs):
+
         self.client_ip_last_seen = Cache(
             name="client_ip_last_seen",
             keylen=4,
@@ -74,6 +75,7 @@ class ClientIpStore(background_updates.BackgroundUpdateStore):
             "before", "shutdown", self._update_client_ips_batch
         )
 
+    @defer.inlineCallbacks
     def insert_client_ip(self, user_id, access_token, ip, user_agent, device_id,
                          now=None):
         if not now:
@@ -84,7 +86,7 @@ class ClientIpStore(background_updates.BackgroundUpdateStore):
             last_seen = self.client_ip_last_seen.get(key)
         except KeyError:
             last_seen = None
-
+        yield self._populate_monthly_active_users(user_id)
         # Rate-limited inserts
         if last_seen is not None and (now - last_seen) < LAST_SEEN_GRANULARITY:
             return
@@ -93,6 +95,24 @@ class ClientIpStore(background_updates.BackgroundUpdateStore):
 
         self._batch_row_update[key] = (user_agent, device_id, now)
 
+    @defer.inlineCallbacks
+    def _populate_monthly_active_users(self, user_id):
+        store = self.hs.get_datastore()
+        print "entering _populate_monthly_active_users"
+        if self.hs.config.limit_usage_by_mau:
+            print "self.hs.config.limit_usage_by_mau is TRUE"
+            is_user_monthly_active = yield store.is_user_monthly_active(user_id)
+            print "is_user_monthly_active is %r" % is_user_monthly_active
+            if is_user_monthly_active:
+                yield store.upsert_monthly_active_user(user_id)
+            else:
+                count = yield store.get_monthly_active_count()
+                print "count is %d" % count
+                if count < self.hs.config.max_mau_value:
+                    print "count is less than self.hs.config.max_mau_value "
+                    res = yield store.upsert_monthly_active_user(user_id)
+                    print "upsert response is %r" % res
+
     def _update_client_ips_batch(self):
         def update():
             to_update = self._batch_row_update
diff --git a/synapse/storage/monthly_active_users.py b/synapse/storage/monthly_active_users.py
index 2337438c58..aada9bd2b9 100644
--- a/synapse/storage/monthly_active_users.py
+++ b/synapse/storage/monthly_active_users.py
@@ -4,7 +4,7 @@ from ._base import SQLBaseStore
 
 
 class MonthlyActiveUsersStore(SQLBaseStore):
-    def __init__(self, hs):
+    def __init__(self, dbconn, hs):
         super(MonthlyActiveUsersStore, self).__init__(None, hs)
         self._clock = hs.get_clock()
         self.max_mau_value = hs.config.max_mau_value
@@ -14,24 +14,28 @@ class MonthlyActiveUsersStore(SQLBaseStore):
         Cleans out monthly active user table to ensure that no stale
         entries exist.
         Return:
-            defered, no return type
+            Defered()
         """
         def _reap_users(txn):
             thirty_days_ago = (
                 int(self._clock.time_msec()) - (1000 * 60 * 60 * 24 * 30)
             )
-
             sql = "DELETE FROM monthly_active_users WHERE timestamp < ?"
-
             txn.execute(sql, (thirty_days_ago,))
+            sql = """
+                DELETE FROM monthly_active_users
+                ORDER BY timestamp desc
+                LIMIT -1 OFFSET ?
+                """
+            txn.execute(sql, (self.max_mau_value,))
 
         return self.runInteraction("reap_monthly_active_users", _reap_users)
 
     def get_monthly_active_count(self):
         """
             Generates current count of monthly active users.abs
-            return:
-                defered resolves to int
+            Return:
+                Defered(int): Number of current monthly active users
         """
         def _count_users(txn):
             sql = "SELECT COALESCE(count(*), 0) FROM monthly_active_users"
@@ -46,6 +50,8 @@ class MonthlyActiveUsersStore(SQLBaseStore):
             Updates or inserts monthly active user member
             Arguments:
                 user_id (str): user to add/update
+            Deferred(bool): True if a new entry was created, False if an
+                existing one was updated.
         """
         return self._simple_upsert(
             desc="upsert_monthly_active_user",
diff --git a/tests/storage/test_client_ips.py b/tests/storage/test_client_ips.py
index bd6fda6cb1..e1552510cc 100644
--- a/tests/storage/test_client_ips.py
+++ b/tests/storage/test_client_ips.py
@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from mock import Mock
 
 from twisted.internet import defer
 
@@ -27,9 +28,9 @@ class ClientIpStoreTestCase(tests.unittest.TestCase):
 
     @defer.inlineCallbacks
     def setUp(self):
-        hs = yield tests.utils.setup_test_homeserver()
-        self.store = hs.get_datastore()
-        self.clock = hs.get_clock()
+        self.hs = yield tests.utils.setup_test_homeserver()
+        self.store = self.hs.get_datastore()
+        self.clock = self.hs.get_clock()
 
     @defer.inlineCallbacks
     def test_insert_new_client_ip(self):
@@ -54,3 +55,62 @@ class ClientIpStoreTestCase(tests.unittest.TestCase):
             },
             r
         )
+
+    @defer.inlineCallbacks
+    def test_disabled_monthly_active_user(self):
+        self.hs.config.limit_usage_by_mau = False
+        self.hs.config.max_mau_value = 50
+        user_id = "@user:server"
+        yield self.store.insert_client_ip(
+            user_id, "access_token", "ip", "user_agent", "device_id",
+        )
+        active = yield self.store.is_user_monthly_active(user_id)
+        self.assertFalse(active)
+
+    @defer.inlineCallbacks
+    def test_adding_monthly_active_user_when_full(self):
+        self.hs.config.limit_usage_by_mau = True
+        self.hs.config.max_mau_value = 50
+        lots_of_users = 100
+        user_id = "@user:server"
+
+        self.store.get_monthly_active_count = Mock(
+            return_value=defer.succeed(lots_of_users)
+        )
+        yield self.store.insert_client_ip(
+            user_id, "access_token", "ip", "user_agent", "device_id",
+        )
+        active = yield self.store.is_user_monthly_active(user_id)
+        self.assertFalse(active)
+
+    @defer.inlineCallbacks
+    def test_adding_monthly_active_user_when_space(self):
+        self.hs.config.limit_usage_by_mau = True
+        self.hs.config.max_mau_value = 50
+        user_id = "@user:server"
+        active = yield self.store.is_user_monthly_active(user_id)
+        self.assertFalse(active)
+
+        yield self.store.insert_client_ip(
+            user_id, "access_token", "ip", "user_agent", "device_id",
+        )
+        active = yield self.store.is_user_monthly_active(user_id)
+        self.assertTrue(active)
+
+    @defer.inlineCallbacks
+    def test_updating_monthly_active_user_when_space(self):
+        self.hs.config.limit_usage_by_mau = True
+        self.hs.config.max_mau_value = 50
+        user_id = "@user:server"
+
+        active = yield self.store.is_user_monthly_active(user_id)
+        self.assertFalse(active)
+
+        yield self.store.insert_client_ip(
+            user_id, "access_token", "ip", "user_agent", "device_id",
+        )
+        yield self.store.insert_client_ip(
+            user_id, "access_token", "ip", "user_agent", "device_id",
+        )
+        active = yield self.store.is_user_monthly_active(user_id)
+        self.assertTrue(active)