mirror of https://github.com/zulip/zulip.git
69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
import cProfile
|
|
import logging
|
|
import tempfile
|
|
from typing import Any
|
|
|
|
from django.contrib.sessions.backends.base import SessionBase
|
|
from django.core.management.base import CommandParser
|
|
from django.http import HttpRequest, HttpResponseBase
|
|
from typing_extensions import override
|
|
|
|
from zerver.lib.management import ZulipBaseCommand
|
|
from zerver.lib.request import RequestNotes
|
|
from zerver.lib.test_helpers import HostRequestMock
|
|
from zerver.middleware import LogRequests
|
|
from zerver.models import UserMessage
|
|
from zerver.views.message_fetch import get_messages_backend
|
|
|
|
|
|
class MockSession(SessionBase):
|
|
def __init__(self) -> None:
|
|
self.modified = False
|
|
|
|
|
|
def profile_request(request: HttpRequest, num_before: int, num_after: int) -> HttpResponseBase:
|
|
def get_response(request: HttpRequest) -> HttpResponseBase:
|
|
return prof.runcall(
|
|
get_messages_backend,
|
|
request,
|
|
request.user,
|
|
num_before=num_before,
|
|
num_after=num_after,
|
|
apply_markdown=True,
|
|
)
|
|
|
|
prof = cProfile.Profile()
|
|
with tempfile.NamedTemporaryFile(prefix="profile.data.", delete=False) as stats_file:
|
|
response = LogRequests(get_response)(request)
|
|
assert isinstance(response, HttpResponseBase) # async responses not supported here for now
|
|
prof.dump_stats(stats_file.name)
|
|
logging.info("Profiling data written to %s", stats_file.name)
|
|
return response
|
|
|
|
|
|
class Command(ZulipBaseCommand):
|
|
@override
|
|
def add_arguments(self, parser: CommandParser) -> None:
|
|
parser.add_argument("email", metavar="<email>", help="Email address of the user")
|
|
self.add_realm_args(parser)
|
|
|
|
@override
|
|
def handle(self, *args: Any, **options: Any) -> None:
|
|
realm = self.get_realm(options)
|
|
user = self.get_user(options["email"], realm)
|
|
anchor = UserMessage.objects.filter(user_profile=user).order_by("-message")[200].message_id
|
|
mock_request = HostRequestMock(
|
|
post_data={
|
|
"anchor": anchor,
|
|
"num_before": 1200,
|
|
"num_after": 200,
|
|
},
|
|
user_profile=user,
|
|
meta_data={"REMOTE_ADDR": "127.0.0.1"},
|
|
path="/",
|
|
)
|
|
mock_request.session = MockSession()
|
|
RequestNotes.get_notes(mock_request).log_data = None
|
|
|
|
profile_request(mock_request, num_before=1200, num_after=200)
|