CodeCow/zulip
1
0
Fork 0
mirror of https://github.com/zulip/zulip.git synced 2026-07-21 21:05:48 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

narrow: Rewrite message-fetch query using Django ORM.

Browse Source 
Our unsupported hack for sharing a PostgreSQL connection between
Django and SQLAlchemy has left us unable to upgrade to current
versions of SQLAlchemy and Psycopg.  This is the only place we use
SQLAlchemy, so let’s remove it.  While Django’s ORM may or may not
have been comprehensive enough to support this logic when it was
originally written, it has no trouble with it today.

I rewrote almost all of this by hand after Claude made a mess of
everything.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
Signed-off-by: Anders Kaseorg <[email protected]>
This commit is contained in:
Anders Kaseorg
2026-06-08 13:53:24 +05:30
committed by Aman Agrawal
co-authored by Claude Opus 4.7
parent 8212f313e9
commit c286f39d46
9 changed files with 820 additions and 974 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/linter_lib/custom_check.py
-1
View File
@@ -11,7 +11,6 @@ from zulint.custom_rules import Rule, RuleList
FILES_WITH_LEGACY_SUBJECT = {
# This basically requires a big DB migration:
"zerver/lib/topic.py",
"zerver/lib/topic_sqlalchemy.py",
# This is tied to legacy events.
"zerver/lib/event_types.py",
# This is for backward compatibility.
zerver/lib/narrow.py
+341 -390
View File
File diff suppressed because it is too large Load Diff
zerver/lib/sqlalchemy_utils.py
-42
View File
@@ -1,42 +0,0 @@
from collections.abc import Iterator
from contextlib import contextmanager
import sqlalchemy
from django.db import connection
from sqlalchemy.engine import Connection, Engine
from typing_extensions import override
from zerver.lib.db import TimeTrackingConnection
# This is a Pool that doesn't close connections. Therefore it can be used with
# existing Django database connections.
class NonClosingPool(sqlalchemy.pool.NullPool):
@override
def status(self) -> str:
return "NonClosingPool"
def _do_return_conn(self, conn: sqlalchemy.engine.base.Connection) -> None:
pass
sqlalchemy_engine: Engine | None = None
@contextmanager
def get_sqlalchemy_connection() -> Iterator[Connection]:
global sqlalchemy_engine
if sqlalchemy_engine is None:
def get_dj_conn() -> TimeTrackingConnection:
connection.ensure_connection()
return connection.connection
sqlalchemy_engine = sqlalchemy.create_engine(
"postgresql://",
creator=get_dj_conn,
poolclass=NonClosingPool,
pool_reset_on_return=None,
)
with sqlalchemy_engine.connect().execution_options(autocommit=False) as sa_connection:
yield sa_connection
zerver/lib/test_runner.py
+2 -7
View File
@@ -24,7 +24,6 @@ from scripts.lib.zulip_tools import (
)
from zerver.lib import test_helpers, upload
from zerver.lib.partial import partial
from zerver.lib.sqlalchemy_utils import get_sqlalchemy_connection
from zerver.lib.test_fixtures import BACKEND_DATABASE_TEMPLATE
from zerver.lib.test_helpers import append_instrumentation_data, write_instrumentation_reports
@@ -391,12 +390,8 @@ class Runner(DiscoverRunner):
destroy_test_databases(_worker_id)
create_test_databases(_worker_id)
# We have to do the next line to avoid flaky scenarios where we
# run a single test and getting an SA connection causes data from
# a Django connection to be rolled back mid-test.
with get_sqlalchemy_connection():
result = self.run_suite(suite)
assert isinstance(result, TextTestResult)
result = self.run_suite(suite)
assert isinstance(result, TextTestResult)
self.teardown_test_environment()
failed = self.suite_result(suite, result)
if not failed:
zerver/lib/topic.py
+34 -2
View File
@@ -4,14 +4,25 @@ from typing import Any
import orjson
from django.db import connection
from django.db.models import F, Func, JSONField, Q, QuerySet, Subquery, TextField, Value
from django.db.models import (
Exists,
F,
Func,
JSONField,
OuterRef,
Q,
QuerySet,
Subquery,
TextField,
Value,
)
from django.db.models.functions import Cast
from django.utils.translation import gettext as _
from django.utils.translation import override as override_language
from zerver.lib.types import EditHistoryEvent, StreamMessageEditRequest
from zerver.lib.utils import assert_is_not_none
from zerver.models import Message, Reaction, UserMessage, UserProfile
from zerver.models import Message, Reaction, UserMessage, UserProfile, UserTopic
# Only use these constants for events.
ORIG_TOPIC = "orig_subject"
@@ -393,3 +404,24 @@ def get_topic_display_name(topic_name: str, language: str) -> str:
with override_language(language):
return _(Message.EMPTY_TOPIC_FALLBACK_NAME)
return topic_name
def topic_match_q(topic_name: str) -> Q:
return Q(subject__iexact=topic_name, is_channel_message=True)
def get_resolved_topic_condition_q() -> Q:
return Q(subject__startswith=RESOLVED_TOPIC_PREFIX, is_channel_message=True)
def get_followed_topic_condition_q(user_id: int) -> Q:
return Q(
Exists(
UserTopic.objects.filter(
user_profile_id=user_id,
visibility_policy=UserTopic.VisibilityPolicy.FOLLOWED,
topic_name__iexact=OuterRef("subject"),
recipient=OuterRef("recipient"),
)
)
)
zerver/lib/topic_sqlalchemy.py
-47
View File
@@ -1,47 +0,0 @@
from sqlalchemy.sql import ColumnElement, and_, column, func, literal, literal_column, select, table
from sqlalchemy.types import Boolean, Text
from zerver.lib.topic import RESOLVED_TOPIC_PREFIX
from zerver.models import UserTopic
def topic_match_sa(topic_name: str) -> ColumnElement[Boolean]:
# _sa is short for SQLAlchemy, which we use mostly for
# queries that search messages
topic_cond = and_(
func.upper(column("subject", Text)) == func.upper(literal(topic_name)),
column("is_channel_message", Boolean),
)
return topic_cond
def get_resolved_topic_condition_sa() -> ColumnElement[Boolean]:
resolved_topic_cond = and_(
column("subject", Text).startswith(RESOLVED_TOPIC_PREFIX),
column("is_channel_message", Boolean),
)
return resolved_topic_cond
def topic_column_sa() -> ColumnElement[Text]:
return column("subject", Text)
def get_followed_topic_condition_sa(user_id: int) -> ColumnElement[Boolean]:
follow_topic_cond = (
select(1)
.select_from(table("zerver_usertopic"))
.where(
and_(
literal_column("zerver_usertopic.user_profile_id") == literal(user_id),
literal_column("zerver_usertopic.visibility_policy")
== literal(UserTopic.VisibilityPolicy.FOLLOWED),
func.upper(literal_column("zerver_usertopic.topic_name"))
== func.upper(literal_column("zerver_message.subject")),
literal_column("zerver_message.is_channel_message", Boolean),
literal_column("zerver_usertopic.recipient_id")
== literal_column("zerver_message.recipient_id"),
)
)
).exists()
return follow_topic_cond
zerver/lib/user_topics.py
+25 -27
View File
@@ -5,14 +5,12 @@ from datetime import datetime
from typing import TypedDict
from django.db import connection, transaction
from django.db.models import QuerySet
from django.db.models import Q, QuerySet
from django.utils.timezone import now as timezone_now
from psycopg2.sql import SQL, Literal
from sqlalchemy.sql import ClauseElement, and_, column, not_, or_
from sqlalchemy.types import Integer
from zerver.lib.timestamp import datetime_to_timestamp
from zerver.lib.topic_sqlalchemy import topic_match_sa
from zerver.lib.topic import topic_match_q
from zerver.lib.types import UserTopicDict
from zerver.models import Recipient, Subscription, UserProfile, UserTopic
from zerver.models.streams import get_stream
@@ -228,9 +226,7 @@ def topic_has_visibility_policy(
return has_visibility_policy
def exclude_stream_and_topic_mutes(
conditions: list[ClauseElement], user_profile: UserProfile, stream_id: int | None
) -> list[ClauseElement]:
def exclude_stream_and_topic_mutes(user_profile: UserProfile, stream_id: int | None) -> Q:
# Note: Unlike get_topic_mutes, here we always want to
# consider topics in deactivated streams, so they are
# never filtered from the query in this method.
@@ -244,26 +240,31 @@ def exclude_stream_and_topic_mutes(
# by not considering topic mutes outside the stream.
query = query.filter(stream_id=stream_id)
excluded_topic_rows = query.values(
"recipient_id",
"topic_name",
excluded_topic_rows = list(
query.values(
"recipient_id",
"topic_name",
)
)
conditions = ~Q(pk__in=[]) # Always true.
class RecipientTopicDict(TypedDict):
recipient_id: int
topic_name: str
def topic_cond(row: RecipientTopicDict) -> ClauseElement:
def topic_cond(row: RecipientTopicDict) -> Q:
recipient_id = row["recipient_id"]
topic_name = row["topic_name"]
stream_cond = column("recipient_id", Integer) == recipient_id
topic_cond = topic_match_sa(topic_name)
return and_(stream_cond, topic_cond)
assert isinstance(topic_name, str)
return Q(recipient_id=recipient_id) & topic_match_q(topic_name)
# Add this query later to reduce the number of messages it has to run on.
if excluded_topic_rows:
exclude_muted_topics_condition = not_(or_(*map(topic_cond, excluded_topic_rows)))
conditions = [*conditions, exclude_muted_topics_condition]
muted_topics_q = topic_cond(excluded_topic_rows[0])
for row in excluded_topic_rows[1:]:
muted_topics_q |= topic_cond(row)
conditions &= ~muted_topics_q
# Channel-level muting only applies when looking at views that
# include multiple channels, since we do want users to be able to
@@ -296,24 +297,21 @@ def exclude_stream_and_topic_mutes(
)
# Exclude muted_recipient_ids unless they match include_followed_or_unmuted_topics_condition
muted_stream_condition = column("recipient_id", Integer).in_(muted_recipient_ids)
muted_stream_condition = Q(recipient_id__in=muted_recipient_ids)
if included_topic_rows:
include_followed_or_unmuted_topics_condition = or_(
*map(topic_cond, included_topic_rows)
)
include_followed_or_unmuted_topics_condition = topic_cond(included_topic_rows[0])
for row in included_topic_rows[1:]: # nocoverage
include_followed_or_unmuted_topics_condition |= topic_cond(row)
exclude_muted_streams_condition = not_(
and_(
muted_stream_condition,
not_(include_followed_or_unmuted_topics_condition),
)
exclude_muted_streams_condition = ~(
muted_stream_condition & ~include_followed_or_unmuted_topics_condition
)
else:
# If no included topics, exclude all muted streams
exclude_muted_streams_condition = not_(muted_stream_condition)
exclude_muted_streams_condition = ~muted_stream_condition
conditions = [*conditions, exclude_muted_streams_condition]
conditions &= exclude_muted_streams_condition
return conditions
zerver/tests/test_message_fetch.py
+393 -432
View File
File diff suppressed because it is too large Load Diff
zerver/views/message_fetch.py
+25 -26
View File
@@ -4,11 +4,10 @@ from typing import Annotated
from django.conf import settings
from django.contrib.auth.models import AnonymousUser
from django.db import connection, transaction
from django.db.models import F, Func, TextField
from django.http import HttpRequest, HttpResponse
from django.utils.translation import gettext as _
from pydantic import Json, NonNegativeInt
from sqlalchemy.sql import column, func
from sqlalchemy.types import Integer, Text
from zerver.context_processors import get_valid_realm_from_request
from zerver.lib.exceptions import (
@@ -29,9 +28,7 @@ from zerver.lib.narrow import (
)
from zerver.lib.request import RequestNotes
from zerver.lib.response import json_success
from zerver.lib.sqlalchemy_utils import get_sqlalchemy_connection
from zerver.lib.topic import MATCH_TOPIC
from zerver.lib.topic_sqlalchemy import topic_column_sa
from zerver.lib.topic import DB_TOPIC_NAME, MATCH_TOPIC
from zerver.lib.typed_endpoint import ApiParamConfig, typed_endpoint
from zerver.models import UserMessage, UserProfile
@@ -347,15 +344,12 @@ def messages_in_narrow_backend(
msg_ids = [message_id for message_id in msg_ids if message_id >= first_visible_message_id]
# This query is limited to messages the user has access to because they
# actually received them, as reflected in `zerver_usermessage`.
query, inner_msg_id_col = get_base_query_for_search(
user_profile.realm_id, user_profile, need_user_message=True
)
query = query.where(column("message_id", Integer).in_(msg_ids))
query = get_base_query_for_search(user_profile.realm_id, user_profile, need_user_message=True)
query = query.filter(id__in=msg_ids)
cleaned_narrow = clean_narrow_for_message_fetch(narrow, user_profile.realm, user_profile)
query, is_search, _is_dm_narrow = add_narrow_conditions(
user_profile=user_profile,
inner_msg_id_col=inner_msg_id_col,
query=query,
narrow=cleaned_narrow,
is_web_public_query=False,
@@ -364,24 +358,29 @@ def messages_in_narrow_backend(
if not is_search:
# `add_narrow_conditions` adds the following columns only if narrow has search operands.
query = query.add_columns(
func.escape_html(topic_column_sa(), type_=Text).label("escaped_topic_name"),
column("rendered_content", Text),
query = query.annotate(
escaped_topic_name=Func(
F(DB_TOPIC_NAME), function="escape_html", output_field=TextField()
),
)
search_fields = {}
with get_sqlalchemy_connection() as sa_conn:
for row in sa_conn.execute(query).mappings():
message_id = row["message_id"]
escaped_topic_name: str = row["escaped_topic_name"]
rendered_content: str = row["rendered_content"]
content_matches = row.get("content_matches", [])
topic_matches = row.get("topic_matches", [])
search_fields[str(message_id)] = get_search_fields(
rendered_content,
escaped_topic_name,
content_matches,
topic_matches,
)
for row in query.values(
"id",
"escaped_topic_name",
"rendered_content",
*["content_matches", "topic_matches"] if is_search else [],
):
message_id = row["id"]
escaped_topic_name: str = row["escaped_topic_name"]
rendered_content: str = row["rendered_content"]
content_matches = row.get("content_matches", [])
topic_matches = row.get("topic_matches", [])
search_fields[str(message_id)] = get_search_fields(
rendered_content,
escaped_topic_name,
content_matches,
topic_matches,
)
return json_success(request, data={"messages": search_fields})