mirror of
https://github.com/zulip/zulip.git
synced 2025-11-04 05:53:43 +00:00
365fe0b3d5
Fixes #2665. Regenerated by tabbott with `lint --fix` after a rebase and change in parameters. Note from tabbott: In a few cases, this converts technical debt in the form of unsorted imports into different technical debt in the form of our largest files having very long, ugly import sequences at the start. I expect this change will increase pressure for us to split those files, which isn't a bad thing. Signed-off-by: Anders Kaseorg <anders@zulip.com>
74 lines
2.4 KiB
Python
74 lines
2.4 KiB
Python
from typing import Dict, List
|
|
|
|
from django.db import connection
|
|
from psycopg2.sql import SQL
|
|
|
|
from zerver.models import Recipient
|
|
|
|
|
|
class StreamRecipientMap:
|
|
'''
|
|
This class maps stream_id -> recipient_id and vice versa.
|
|
It is useful for bulk operations. Call the populate_* methods
|
|
to initialize the data structures. You should try to avoid
|
|
excessive queries by finding ids up front, but you can call
|
|
this repeatedly, and it will only look up new ids.
|
|
|
|
You should ONLY use this class for READ operations.
|
|
|
|
Note that this class uses raw SQL, because we want to highly
|
|
optimize page loads.
|
|
'''
|
|
|
|
def __init__(self) -> None:
|
|
self.recip_to_stream: Dict[int, int] = dict()
|
|
self.stream_to_recip: Dict[int, int] = dict()
|
|
|
|
def populate_with(self, *, stream_id: int, recipient_id: int) -> None:
|
|
# We use * to enforce using named arguments when calling this function,
|
|
# to avoid confusion about the ordering of the two integers.
|
|
self.recip_to_stream[recipient_id] = stream_id
|
|
self.stream_to_recip[stream_id] = recipient_id
|
|
|
|
def populate_for_recipient_ids(self, recipient_ids: List[int]) -> None:
|
|
recipient_ids = sorted([
|
|
recip_id for recip_id in recipient_ids
|
|
if recip_id not in self.recip_to_stream
|
|
])
|
|
|
|
if not recipient_ids:
|
|
return
|
|
|
|
# see comment at the top of the class
|
|
query = SQL('''
|
|
SELECT
|
|
zerver_recipient.id as recipient_id,
|
|
zerver_recipient.type_id as stream_id
|
|
FROM
|
|
zerver_recipient
|
|
WHERE
|
|
zerver_recipient.type = %(STREAM)s
|
|
AND
|
|
zerver_recipient.id in %(recipient_ids)s
|
|
''')
|
|
|
|
cursor = connection.cursor()
|
|
cursor.execute(query, {
|
|
"STREAM": Recipient.STREAM,
|
|
"recipient_ids": tuple(recipient_ids),
|
|
})
|
|
rows = cursor.fetchall()
|
|
cursor.close()
|
|
for recip_id, stream_id in rows:
|
|
self.recip_to_stream[recip_id] = stream_id
|
|
self.stream_to_recip[stream_id] = recip_id
|
|
|
|
def recipient_id_for(self, stream_id: int) -> int:
|
|
return self.stream_to_recip[stream_id]
|
|
|
|
def stream_id_for(self, recip_id: int) -> int:
|
|
return self.recip_to_stream[recip_id]
|
|
|
|
def recipient_to_stream_id_dict(self) -> Dict[int, int]:
|
|
return self.recip_to_stream
|