models: Move query_for_ids to zerver.lib.query_helpers.

Signed-off-by: Anders Kaseorg <anders@zulip.com>
This commit is contained in:
Anders Kaseorg 2023-12-15 12:10:20 -08:00 committed by Tim Abbott
parent 33d140c8dc
commit c343d7c30e
5 changed files with 35 additions and 32 deletions

View File

@ -68,6 +68,7 @@ from zerver.lib.notification_data import (
get_user_group_mentions_data, get_user_group_mentions_data,
user_allows_notifications_in_StreamTopic, user_allows_notifications_in_StreamTopic,
) )
from zerver.lib.query_helpers import query_for_ids
from zerver.lib.queue import queue_json_publish from zerver.lib.queue import queue_json_publish
from zerver.lib.recipient_users import recipient_for_user_profiles from zerver.lib.recipient_users import recipient_for_user_profiles
from zerver.lib.stream_subscription import ( from zerver.lib.stream_subscription import (
@ -101,7 +102,6 @@ from zerver.models import (
UserPresence, UserPresence,
UserProfile, UserProfile,
UserTopic, UserTopic,
query_for_ids,
) )
from zerver.models.clients import get_client from zerver.models.clients import get_client
from zerver.models.groups import SystemGroups from zerver.models.groups import SystemGroups

View File

@ -42,6 +42,7 @@ from zerver.lib.exceptions import JsonableError, MissingAuthenticationError
from zerver.lib.markdown import MessageRenderingResult, markdown_convert, topic_links from zerver.lib.markdown import MessageRenderingResult, markdown_convert, topic_links
from zerver.lib.markdown import version as markdown_version from zerver.lib.markdown import version as markdown_version
from zerver.lib.mention import MentionData from zerver.lib.mention import MentionData
from zerver.lib.query_helpers import query_for_ids
from zerver.lib.request import RequestVariableConversionError from zerver.lib.request import RequestVariableConversionError
from zerver.lib.stream_subscription import ( from zerver.lib.stream_subscription import (
get_stream_subscriptions_for_user, get_stream_subscriptions_for_user,
@ -74,7 +75,6 @@ from zerver.models import (
UserMessage, UserMessage,
UserProfile, UserProfile,
UserTopic, UserTopic,
query_for_ids,
) )
from zerver.models.constants import MAX_TOPIC_NAME_LENGTH from zerver.models.constants import MAX_TOPIC_NAME_LENGTH
from zerver.models.messages import get_usermessage_by_message_id from zerver.models.messages import get_usermessage_by_message_id

View File

@ -6,9 +6,10 @@ from typing import Any, Dict, Mapping, Optional, Sequence, Set
from django.conf import settings from django.conf import settings
from django.utils.timezone import now as timezone_now from django.utils.timezone import now as timezone_now
from zerver.lib.query_helpers import query_for_ids
from zerver.lib.timestamp import datetime_to_timestamp from zerver.lib.timestamp import datetime_to_timestamp
from zerver.lib.users import check_user_can_access_all_users, get_accessible_user_ids from zerver.lib.users import check_user_can_access_all_users, get_accessible_user_ids
from zerver.models import PushDeviceToken, Realm, UserPresence, UserProfile, query_for_ids from zerver.models import PushDeviceToken, Realm, UserPresence, UserProfile
def get_presence_dicts_for_rows( def get_presence_dicts_for_rows(

View File

@ -0,0 +1,30 @@
from typing import List, TypeVar
from django.db import models
from django_stubs_ext import ValuesQuerySet
ModelT = TypeVar("ModelT", bound=models.Model)
RowT = TypeVar("RowT")
def query_for_ids(
query: ValuesQuerySet[ModelT, RowT],
user_ids: List[int],
field: str,
) -> ValuesQuerySet[ModelT, RowT]:
"""
This function optimizes searches of the form
`user_profile_id in (1, 2, 3, 4)` by quickly
building the where clauses. Profiling shows significant
speedups over the normal Django-based approach.
Use this very carefully! Also, the caller should
guard against empty lists of user_ids.
"""
assert user_ids
clause = f"{field} IN %s"
query = query.extra(
where=[clause],
params=(tuple(user_ids),),
)
return query

View File

@ -1,9 +1,8 @@
from typing import List, Tuple, TypeVar, Union from typing import List, Tuple, Union
from django.db import models from django.db import models
from django.db.backends.base.base import BaseDatabaseWrapper from django.db.backends.base.base import BaseDatabaseWrapper
from django.db.models.sql.compiler import SQLCompiler from django.db.models.sql.compiler import SQLCompiler
from django_stubs_ext import ValuesQuerySet
from typing_extensions import override from typing_extensions import override
from zerver.models.alert_words import AlertWord as AlertWord from zerver.models.alert_words import AlertWord as AlertWord
@ -98,30 +97,3 @@ class AndNonZero(models.Lookup[int]):
lhs, lhs_params = self.process_lhs(compiler, connection) lhs, lhs_params = self.process_lhs(compiler, connection)
rhs, rhs_params = self.process_rhs(compiler, connection) rhs, rhs_params = self.process_rhs(compiler, connection)
return f"{lhs} & {rhs} != 0", lhs_params + rhs_params return f"{lhs} & {rhs} != 0", lhs_params + rhs_params
ModelT = TypeVar("ModelT", bound=models.Model)
RowT = TypeVar("RowT")
def query_for_ids(
query: ValuesQuerySet[ModelT, RowT],
user_ids: List[int],
field: str,
) -> ValuesQuerySet[ModelT, RowT]:
"""
This function optimizes searches of the form
`user_profile_id in (1, 2, 3, 4)` by quickly
building the where clauses. Profiling shows significant
speedups over the normal Django-based approach.
Use this very carefully! Also, the caller should
guard against empty lists of user_ids.
"""
assert user_ids
clause = f"{field} IN %s"
query = query.extra(
where=[clause],
params=(tuple(user_ids),),
)
return query