upstream-zulip/zerver/management/commands/enqueue_file.py

import sys
from argparse import ArgumentParser
from typing import IO, Any

import orjson
from typing_extensions import override

from zerver.lib.management import ZulipBaseCommand
from zerver.lib.queue import queue_json_publish_rollback_unsafe


def error(*args: Any) -> None:
    raise Exception("We cannot enqueue because settings.USING_RABBITMQ is False.")


def enqueue_file(queue_name: str, f: IO[str]) -> None:
    for line in f:
        line = line.strip()
        try:
            payload = line.split("\t")[1]
        except IndexError:
            payload = line

        print(f"Queueing to queue {queue_name}: {payload}")

        # Verify that payload is valid json.
        data = orjson.loads(payload)

        # This is designed to use the `error` method rather than
        # the call_consume_in_tests flow.
        queue_json_publish_rollback_unsafe(queue_name, data, error)


class Command(ZulipBaseCommand):
    help = """Read JSON lines from a file and enqueue them to a worker queue.

Each line in the file should either be a JSON payload or two tab-separated
fields, the second of which is a JSON payload.  (The latter is to accommodate
the format of error files written by queue workers that catch exceptions--their
first field is a timestamp that we ignore.)

You can use "-" to represent stdin.
"""

    @override
    def add_arguments(self, parser: ArgumentParser) -> None:
        parser.add_argument(
            "queue_name", metavar="<queue>", help="name of worker queue to enqueue to"
        )
        parser.add_argument(
            "file_name", metavar="<file>", help="name of file containing JSON lines"
        )

    @override
    def handle(self, *args: Any, **options: str) -> None:
        queue_name = options["queue_name"]
        file_name = options["file_name"]

        if file_name == "-":
            enqueue_file(queue_name, sys.stdin)
        else:
            with open(file_name) as f:
                enqueue_file(queue_name, f)