Move _add_messages_to_device_inbox_txn into a separate method
parent
940d4fad24
commit
f4164edb70
|
@ -37,9 +37,21 @@ class DeviceInboxStore(SQLBaseStore):
|
||||||
inserted.
|
inserted.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def select_devices_txn(txn, user_id, devices):
|
with self._device_inbox_id_gen.get_next() as stream_id:
|
||||||
if not devices:
|
yield self.runInteraction(
|
||||||
return []
|
"add_messages_to_device_inbox",
|
||||||
|
self._add_messages_to_device_inbox_txn,
|
||||||
|
stream_id,
|
||||||
|
messages_by_user_then_device,
|
||||||
|
)
|
||||||
|
|
||||||
|
defer.returnValue(self._device_inbox_id_gen.get_current_token())
|
||||||
|
|
||||||
|
def _add_messages_to_device_inbox_txn(self, txn, stream_id,
|
||||||
|
messages_by_user_then_device):
|
||||||
|
local_users_and_devices = set()
|
||||||
|
for user_id, messages_by_device in messages_by_user_then_device.items():
|
||||||
|
devices = messages_by_device.keys()
|
||||||
sql = (
|
sql = (
|
||||||
"SELECT user_id, device_id FROM devices"
|
"SELECT user_id, device_id FROM devices"
|
||||||
" WHERE user_id = ? AND device_id IN ("
|
" WHERE user_id = ? AND device_id IN ("
|
||||||
|
@ -48,41 +60,24 @@ class DeviceInboxStore(SQLBaseStore):
|
||||||
)
|
)
|
||||||
# TODO: Maybe this needs to be done in batches if there are
|
# TODO: Maybe this needs to be done in batches if there are
|
||||||
# too many local devices for a given user.
|
# too many local devices for a given user.
|
||||||
args = [user_id] + devices
|
txn.execute(sql, [user_id] + devices)
|
||||||
txn.execute(sql, args)
|
local_users_and_devices.update(map(tuple, txn.fetchall()))
|
||||||
return [tuple(row) for row in txn.fetchall()]
|
|
||||||
|
|
||||||
def add_messages_to_device_inbox_txn(txn, stream_id):
|
sql = (
|
||||||
local_users_and_devices = set()
|
"INSERT INTO device_inbox"
|
||||||
for user_id, messages_by_device in messages_by_user_then_device.items():
|
" (user_id, device_id, stream_id, message_json)"
|
||||||
local_users_and_devices.update(
|
" VALUES (?,?,?,?)"
|
||||||
select_devices_txn(txn, user_id, messages_by_device.keys())
|
)
|
||||||
)
|
rows = []
|
||||||
|
for user_id, messages_by_device in messages_by_user_then_device.items():
|
||||||
|
for device_id, message in messages_by_device.items():
|
||||||
|
message_json = ujson.dumps(message)
|
||||||
|
# Only insert into the local inbox if the device exists on
|
||||||
|
# this server
|
||||||
|
if (user_id, device_id) in local_users_and_devices:
|
||||||
|
rows.append((user_id, device_id, stream_id, message_json))
|
||||||
|
|
||||||
sql = (
|
txn.executemany(sql, rows)
|
||||||
"INSERT INTO device_inbox"
|
|
||||||
" (user_id, device_id, stream_id, message_json)"
|
|
||||||
" VALUES (?,?,?,?)"
|
|
||||||
)
|
|
||||||
rows = []
|
|
||||||
for user_id, messages_by_device in messages_by_user_then_device.items():
|
|
||||||
for device_id, message in messages_by_device.items():
|
|
||||||
message_json = ujson.dumps(message)
|
|
||||||
# Only insert into the local inbox if the device exists on
|
|
||||||
# this server
|
|
||||||
if (user_id, device_id) in local_users_and_devices:
|
|
||||||
rows.append((user_id, device_id, stream_id, message_json))
|
|
||||||
|
|
||||||
txn.executemany(sql, rows)
|
|
||||||
|
|
||||||
with self._device_inbox_id_gen.get_next() as stream_id:
|
|
||||||
yield self.runInteraction(
|
|
||||||
"add_messages_to_device_inbox",
|
|
||||||
add_messages_to_device_inbox_txn,
|
|
||||||
stream_id
|
|
||||||
)
|
|
||||||
|
|
||||||
defer.returnValue(self._device_inbox_id_gen.get_current_token())
|
|
||||||
|
|
||||||
def get_new_messages_for_device(
|
def get_new_messages_for_device(
|
||||||
self, user_id, device_id, last_stream_id, current_stream_id, limit=100
|
self, user_id, device_id, last_stream_id, current_stream_id, limit=100
|
||||||
|
|
Loading…
Reference in New Issue