message_fetch: Extract fetch_messages helper to zerver.lib.narrow.

Signed-off-by: Anders Kaseorg <anders@zulip.com>
This commit is contained in:
Anders Kaseorg 2022-11-09 15:35:52 -08:00 committed by Tim Abbott
parent bd5d200722
commit dae4633745
4 changed files with 134 additions and 97 deletions

View File

@ -23,7 +23,7 @@ from django.core.exceptions import ValidationError
from django.db import connection
from django.utils.translation import gettext as _
from sqlalchemy.dialects import postgresql
from sqlalchemy.engine import Connection
from sqlalchemy.engine import Connection, Row
from sqlalchemy.sql import (
ClauseElement,
ColumnElement,
@ -45,7 +45,9 @@ from sqlalchemy.types import ARRAY, Boolean, Integer, Text
from zerver.lib.addressee import get_user_profiles, get_user_profiles_by_ids
from zerver.lib.exceptions import ErrorCode, JsonableError
from zerver.lib.message import get_first_visible_message_id
from zerver.lib.recipient_users import recipient_for_user_profiles
from zerver.lib.sqlalchemy_utils import get_sqlalchemy_connection
from zerver.lib.streams import (
can_access_stream_history_by_id,
can_access_stream_history_by_name,
@ -1145,3 +1147,120 @@ def post_process_limited_query(
found_oldest=found_oldest,
history_limited=history_limited,
)
@dataclass
class FetchedMessages(LimitedMessages[Row]):
anchor: int
include_history: bool
is_search: bool
def fetch_messages(
*,
narrow: OptionalNarrowListT,
user_profile: Optional[UserProfile],
realm: Realm,
is_web_public_query: bool,
anchor: Optional[int],
num_before: int,
num_after: int,
) -> FetchedMessages:
include_history = ok_to_include_history(narrow, user_profile, is_web_public_query)
if include_history:
# The initial query in this case doesn't use `zerver_usermessage`,
# and isn't yet limited to messages the user is entitled to see!
#
# This is OK only because we've made sure this is a narrow that
# will cause us to limit the query appropriately elsewhere.
# See `ok_to_include_history` for details.
#
# Note that is_web_public_query=True goes here, since
# include_history is semantically correct for is_web_public_query.
need_message = True
need_user_message = False
elif narrow is None:
# We need to limit to messages the user has received, but we don't actually
# need any fields from Message
need_message = False
need_user_message = True
else:
need_message = True
need_user_message = True
query: SelectBase
query, inner_msg_id_col = get_base_query_for_search(
user_profile=user_profile,
need_message=need_message,
need_user_message=need_user_message,
)
query, is_search = add_narrow_conditions(
user_profile=user_profile,
inner_msg_id_col=inner_msg_id_col,
query=query,
narrow=narrow,
realm=realm,
is_web_public_query=is_web_public_query,
)
with get_sqlalchemy_connection() as sa_conn:
if anchor is None:
# `anchor=None` corresponds to the anchor="first_unread" parameter.
anchor = find_first_unread_anchor(
sa_conn,
user_profile,
narrow,
)
anchored_to_left = anchor == 0
# Set value that will be used to short circuit the after_query
# altogether and avoid needless conditions in the before_query.
anchored_to_right = anchor >= LARGER_THAN_MAX_MESSAGE_ID
if anchored_to_right:
num_after = 0
first_visible_message_id = get_first_visible_message_id(realm)
query = limit_query_to_range(
query=query,
num_before=num_before,
num_after=num_after,
anchor=anchor,
anchored_to_left=anchored_to_left,
anchored_to_right=anchored_to_right,
id_col=inner_msg_id_col,
first_visible_message_id=first_visible_message_id,
)
main_query = query.subquery()
query = (
select(*main_query.c)
.select_from(main_query)
.order_by(column("message_id", Integer).asc())
)
# This is a hack to tag the query we use for testing
query = query.prefix_with("/* get_messages */")
rows = list(sa_conn.execute(query).fetchall())
query_info = post_process_limited_query(
rows=rows,
num_before=num_before,
num_after=num_after,
anchor=anchor,
anchored_to_left=anchored_to_left,
anchored_to_right=anchored_to_right,
first_visible_message_id=first_visible_message_id,
)
return FetchedMessages(
rows=query_info.rows,
found_anchor=query_info.found_anchor,
found_newest=query_info.found_newest,
found_oldest=query_info.found_oldest,
history_limited=query_info.history_limited,
anchor=anchor,
include_history=include_history,
is_search=is_search,
)

View File

@ -92,7 +92,7 @@ def mute_stream(realm: Realm, user_profile: UserProfile, stream_name: str) -> No
def first_visible_id_as(message_id: int) -> Any:
return mock.patch(
"zerver.views.message_fetch.get_first_visible_message_id",
"zerver.lib.narrow.get_first_visible_message_id",
return_value=message_id,
)

View File

@ -136,7 +136,7 @@ class FirstUnreadAnchorTests(ZulipTestCase):
self.assertEqual(messages_response["anchor"], new_message_id)
with mock.patch(
"zerver.views.message_fetch.get_first_visible_message_id", return_value=new_message_id
"zerver.lib.narrow.get_first_visible_message_id", return_value=new_message_id
):
messages_response = self.get_messages_response(
anchor="first_unread", num_before=0, num_after=1
@ -145,7 +145,7 @@ class FirstUnreadAnchorTests(ZulipTestCase):
self.assertEqual(messages_response["anchor"], new_message_id)
with mock.patch(
"zerver.views.message_fetch.get_first_visible_message_id",
"zerver.lib.narrow.get_first_visible_message_id",
return_value=new_message_id + 1,
):
messages_reponse = self.get_messages_response(
@ -155,7 +155,7 @@ class FirstUnreadAnchorTests(ZulipTestCase):
self.assertIn("anchor", messages_reponse)
with mock.patch(
"zerver.views.message_fetch.get_first_visible_message_id",
"zerver.lib.narrow.get_first_visible_message_id",
return_value=new_message_id - 1,
):
messages = self.get_messages(anchor="first_unread", num_before=0, num_after=1)

View File

@ -5,26 +5,19 @@ from django.http import HttpRequest, HttpResponse
from django.utils.html import escape as escape_html
from django.utils.translation import gettext as _
from sqlalchemy.sql import and_, column, join, literal, literal_column, select, table
from sqlalchemy.sql.selectable import SelectBase
from sqlalchemy.types import Integer, Text
from zerver.context_processors import get_valid_realm_from_request
from zerver.lib.exceptions import JsonableError, MissingAuthenticationError
from zerver.lib.message import get_first_visible_message_id, messages_for_ids
from zerver.lib.narrow import (
LARGER_THAN_MAX_MESSAGE_ID,
NarrowBuilder,
OptionalNarrowListT,
add_narrow_conditions,
find_first_unread_anchor,
get_base_query_for_search,
fetch_messages,
is_spectator_compatible,
is_web_public_narrow,
limit_query_to_range,
narrow_parameter,
ok_to_include_history,
parse_anchor_value,
post_process_limited_query,
)
from zerver.lib.request import REQ, RequestNotes, has_request_variables
from zerver.lib.response import json_success
@ -149,44 +142,6 @@ def get_messages_backend(
# clients cannot compute gravatars, so we force-set it to false.
client_gravatar = False
include_history = ok_to_include_history(narrow, user_profile, is_web_public_query)
if include_history:
# The initial query in this case doesn't use `zerver_usermessage`,
# and isn't yet limited to messages the user is entitled to see!
#
# This is OK only because we've made sure this is a narrow that
# will cause us to limit the query appropriately elsewhere.
# See `ok_to_include_history` for details.
#
# Note that is_web_public_query=True goes here, since
# include_history is semantically correct for is_web_public_query.
need_message = True
need_user_message = False
elif narrow is None:
# We need to limit to messages the user has received, but we don't actually
# need any fields from Message
need_message = False
need_user_message = True
else:
need_message = True
need_user_message = True
query: SelectBase
query, inner_msg_id_col = get_base_query_for_search(
user_profile=user_profile,
need_message=need_message,
need_user_message=need_user_message,
)
query, is_search = add_narrow_conditions(
user_profile=user_profile,
inner_msg_id_col=inner_msg_id_col,
query=query,
narrow=narrow,
realm=realm,
is_web_public_query=is_web_public_query,
)
if narrow is not None:
# Add some metadata to our logging data for narrows
verbose_operators = []
@ -199,56 +154,19 @@ def get_messages_backend(
assert log_data is not None
log_data["extra"] = "[{}]".format(",".join(verbose_operators))
with get_sqlalchemy_connection() as sa_conn:
if anchor is None:
# `anchor=None` corresponds to the anchor="first_unread" parameter.
anchor = find_first_unread_anchor(
sa_conn,
user_profile,
narrow,
)
anchored_to_left = anchor == 0
# Set value that will be used to short circuit the after_query
# altogether and avoid needless conditions in the before_query.
anchored_to_right = anchor >= LARGER_THAN_MAX_MESSAGE_ID
if anchored_to_right:
num_after = 0
first_visible_message_id = get_first_visible_message_id(realm)
query = limit_query_to_range(
query=query,
query_info = fetch_messages(
narrow=narrow,
user_profile=user_profile,
realm=realm,
is_web_public_query=is_web_public_query,
anchor=anchor,
num_before=num_before,
num_after=num_after,
anchor=anchor,
anchored_to_left=anchored_to_left,
anchored_to_right=anchored_to_right,
id_col=inner_msg_id_col,
first_visible_message_id=first_visible_message_id,
)
main_query = query.subquery()
query = (
select(*main_query.c)
.select_from(main_query)
.order_by(column("message_id", Integer).asc())
)
# This is a hack to tag the query we use for testing
query = query.prefix_with("/* get_messages */")
rows = list(sa_conn.execute(query).fetchall())
query_info = post_process_limited_query(
rows=rows,
num_before=num_before,
num_after=num_after,
anchor=anchor,
anchored_to_left=anchored_to_left,
anchored_to_right=anchored_to_right,
first_visible_message_id=first_visible_message_id,
)
anchor = query_info.anchor
include_history = query_info.include_history
is_search = query_info.is_search
rows = query_info.rows
# The following is a little messy, but ensures that the code paths