From 1f0f5ffa1e22dcbe2b0bb605ccaf12bf571dc624 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 5 May 2016 09:51:03 +0100
Subject: [PATCH] Add bulk fetch storage API

---
 synapse/handlers/_base.py       | 10 +++-------
 synapse/handlers/sync.py        |  2 +-
 synapse/storage/account_data.py | 25 ++++++++++++++++++++++---
 3 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/synapse/handlers/_base.py b/synapse/handlers/_base.py
index 0912274e1a..81f7929b50 100644
--- a/synapse/handlers/_base.py
+++ b/synapse/handlers/_base.py
@@ -95,13 +95,9 @@ class BaseHandler(object):
             row["event_id"] for rows in forgotten for row in rows
         )
 
-        # Maps user_id -> account data content
-        ignore_dict_content = yield defer.gatherResults([
-            preserve_fn(self.store.get_global_account_data_by_type_for_user)(
-                user_id, "m.ignored_user_list"
-            ).addCallback(lambda d, u: (u, d), user_id)
-            for user_id, is_peeking in user_tuples
-        ]).addCallback(dict)
+        ignore_dict_content = yield self.store.get_global_account_data_by_type_for_users(
+            "m.ignored_user_list", user_ids=[user_id for user_id, _ in user_tuples]
+        )
 
         # FIXME: This will explode if people upload something incorrect.
         ignore_dict = {
diff --git a/synapse/handlers/sync.py b/synapse/handlers/sync.py
index 27c3c1b525..0bb1913285 100644
--- a/synapse/handlers/sync.py
+++ b/synapse/handlers/sync.py
@@ -522,7 +522,7 @@ class SyncHandler(BaseHandler):
             )
 
         ignored_account_data = yield self.store.get_global_account_data_by_type_for_user(
-            user_id, "m.ignored_user_list"
+            "m.ignored_user_list", user_id=user_id,
         )
 
         if ignored_account_data:
diff --git a/synapse/storage/account_data.py b/synapse/storage/account_data.py
index cc0b92bc89..ec7e8d40d2 100644
--- a/synapse/storage/account_data.py
+++ b/synapse/storage/account_data.py
@@ -16,7 +16,7 @@
 from ._base import SQLBaseStore
 from twisted.internet import defer
 
-from synapse.util.caches.descriptors import cached, cachedInlineCallbacks
+from synapse.util.caches.descriptors import cached, cachedList, cachedInlineCallbacks
 
 import ujson as json
 import logging
@@ -64,7 +64,7 @@ class AccountDataStore(SQLBaseStore):
         )
 
     @cachedInlineCallbacks(num_args=2)
-    def get_global_account_data_by_type_for_user(self, user_id, data_type):
+    def get_global_account_data_by_type_for_user(self, data_type, user_id):
         """
         Returns:
             Deferred: A dict
@@ -85,6 +85,25 @@ class AccountDataStore(SQLBaseStore):
         else:
             defer.returnValue(None)
 
+    @cachedList(cached_method_name="get_global_account_data_by_type_for_user",
+                num_args=2, list_name="user_ids", inlineCallbacks=True)
+    def get_global_account_data_by_type_for_users(self, data_type, user_ids):
+        rows = yield self._simple_select_many_batch(
+            table="account_data",
+            column="user_id",
+            iterable=user_ids,
+            keyvalues={
+                "account_data_type": data_type,
+            },
+            retcols=("user_id", "content",),
+            desc="get_global_account_data_by_type_for_users",
+        )
+
+        defer.returnValue({
+            row["user_id"]: json.loads(row["content"]) if row["content"] else None
+            for row in rows
+        })
+
     def get_account_data_for_room(self, user_id, room_id):
         """Get all the client account_data for a user for a room.
 
@@ -261,7 +280,7 @@ class AccountDataStore(SQLBaseStore):
             txn.call_after(self.get_account_data_for_user.invalidate, (user_id,))
             txn.call_after(
                 self.get_global_account_data_by_type_for_user.invalidate,
-                (user_id, account_data_type,)
+                (account_data_type, user_id,)
             )
             self._update_max_stream_id(txn, next_id)