Files
zulip/zilencer/management/commands/profile_request.py
PIG208 460119986b management: Fix typing for management scripts.
There are some remaining errors related to the django `Manager[T]` and
the `List[T]` type that we use to annotate the `Manage[T]` objects.
2021-08-20 05:54:18 -07:00

70 lines
2.2 KiB
Python

import cProfile
import logging
import tempfile
from os import path
from typing import Any, Dict
from django.contrib.sessions.backends.base import SessionBase
from django.core.management.base import CommandParser
from django.http import HttpRequest, HttpResponse
from zerver.lib.management import ZulipBaseCommand
from zerver.lib.request import get_request_notes
from zerver.lib.test_helpers import HostRequestMock
from zerver.middleware import LogRequests
from zerver.models import UserMessage, UserProfile
from zerver.views.message_fetch import get_messages_backend
request_logger = LogRequests()
class MockSession:
def __init__(self) -> None:
self.modified = False
class MockRequest(HttpRequest):
def __init__(self, user: UserProfile) -> None:
self.user = user
self.path = "/"
self.method = "POST"
self.META = {"REMOTE_ADDR": "127.0.0.1"}
anchor = (
UserMessage.objects.filter(user_profile=self.user).order_by("-message")[200].message_id
)
self.REQUEST = {
"anchor": anchor,
"num_before": 1200,
"num_after": 200,
}
self.POST = self.REQUEST
self.GET: Dict[Any, Any] = {}
self.session = MockSession()
def get_full_path(self) -> str:
return self.path
def profile_request(request: HttpRequest) -> HttpResponse:
request_logger.process_request(request)
prof = cProfile.Profile()
prof.enable()
ret = get_messages_backend(request, request.user, apply_markdown=True)
prof.disable()
with tempfile.NamedTemporaryFile(prefix="profile.data.", delete=False) as stats_file:
prof.dump_stats(stats_file.name)
request_logger.process_response(request, ret)
logging.info("Profiling data written to %s", stats_file.name)
return ret
class Command(ZulipBaseCommand):
def add_arguments(self, parser: CommandParser) -> None:
parser.add_argument("email", metavar="<email>", help="Email address of the user")
self.add_realm_args(parser)
def handle(self, *args: Any, **options: Any) -> None:
realm = self.get_realm(options)
user = self.get_user(options["email"], realm)
profile_request(MockRequest(user))