zulip/zerver/lib/fix_unreads.py

import logging
import time
from typing import Callable, List

from django.db import connection
from django.db.backends.utils import CursorWrapper
from psycopg2.sql import SQL

from zerver.models import UserProfile

"""
NOTE!  Be careful modifying this library, as it is used
in a migration, and it needs to be valid for the state
of the database that is in place when the 0104_fix_unreads
migration runs.
"""

logger = logging.getLogger("zulip.fix_unreads")
logger.setLevel(logging.WARNING)


def update_unread_flags(cursor: CursorWrapper, user_message_ids: List[int]) -> None:
    query = SQL(
        """
        UPDATE zerver_usermessage
        SET flags = flags | 1
        WHERE id IN %(user_message_ids)s
    """
    )

    cursor.execute(query, {"user_message_ids": tuple(user_message_ids)})


def get_timing(message: str, f: Callable[[], None]) -> None:
    start = time.time()
    logger.info(message)
    f()
    elapsed = time.time() - start
    logger.info("elapsed time: %.03f\n", elapsed)


def fix_unsubscribed(cursor: CursorWrapper, user_profile: UserProfile) -> None:
    recipient_ids = []

    def find_recipients() -> None:
        query = SQL(
            """
            SELECT
                zerver_subscription.recipient_id
            FROM
                zerver_subscription
            INNER JOIN zerver_recipient ON (
                zerver_recipient.id = zerver_subscription.recipient_id
            )
            WHERE (
                zerver_subscription.user_profile_id = %(user_profile_id)s AND
                zerver_recipient.type = 2 AND
                (NOT zerver_subscription.active)
            )
        """
        )
        cursor.execute(query, {"user_profile_id": user_profile.id})
        rows = cursor.fetchall()
        for row in rows:
            recipient_ids.append(row[0])
        logger.info("%s", recipient_ids)

    get_timing(
        "get recipients",
        find_recipients,
    )

    if not recipient_ids:
        return

    user_message_ids = []

    def find() -> None:
        query = SQL(
            """
            SELECT
                zerver_usermessage.id
            FROM
                zerver_usermessage
            INNER JOIN zerver_message ON (
                zerver_message.id = zerver_usermessage.message_id
            )
            WHERE (
                zerver_usermessage.user_profile_id = %(user_profile_id)s AND
                (zerver_usermessage.flags & 1) = 0 AND
                zerver_message.recipient_id in %(recipient_ids)s
            )
        """
        )

        cursor.execute(
            query,
            {
                "user_profile_id": user_profile.id,
                "recipient_ids": tuple(recipient_ids),
            },
        )
        rows = cursor.fetchall()
        for row in rows:
            user_message_ids.append(row[0])
        logger.info("rows found: %d", len(user_message_ids))

    get_timing(
        "finding unread messages for non-active streams",
        find,
    )

    if not user_message_ids:
        return

    def fix() -> None:
        update_unread_flags(cursor, user_message_ids)

    get_timing(
        "fixing unread messages for non-active streams",
        fix,
    )


def fix(user_profile: UserProfile) -> None:
    logger.info("\n---\nFixing %s:", user_profile.id)
    with connection.cursor() as cursor:
        fix_unsubscribed(cursor, user_profile)