1from flask import current_app
2from flask_debugtoolbar.panels import DebugPanel
3from jinja2 import ChoiceLoader, PackageLoader
4
5from flask_mongoengine import operation_tracker
6
7package_loader = PackageLoader("flask_mongoengine", "templates")
8
9
10def _maybe_patch_jinja_loader(jinja_env):
11    """Patch the jinja_env loader to include flaskext.mongoengine
12    templates folder if necessary.
13    """
14    if not isinstance(jinja_env.loader, ChoiceLoader):
15        jinja_env.loader = ChoiceLoader([jinja_env.loader, package_loader])
16    elif package_loader not in jinja_env.loader.loaders:
17        jinja_env.loader.loaders.append(package_loader)
18
19
20class MongoDebugPanel(DebugPanel):
21    """Panel that shows information about MongoDB operations (including stack)
22
23    Adapted from https://github.com/hmarr/django-debug-toolbar-mongo
24    """
25
26    name = "MongoDB"
27    has_content = True
28
29    def __init__(self, *args, **kwargs):
30        super(MongoDebugPanel, self).__init__(*args, **kwargs)
31        _maybe_patch_jinja_loader(self.jinja_env)
32        operation_tracker.install_tracker()
33
34    def process_request(self, request):
35        operation_tracker.reset()
36
37    def nav_title(self):
38        return "MongoDB"
39
40    def nav_subtitle(self):
41        attrs = ["queries", "inserts", "updates", "removes"]
42        ops = sum(
43            sum((1 for o in getattr(operation_tracker, a) if not o["internal"]))
44            for a in attrs
45        )
46        total_time = sum(
47            sum(o["time"] for o in getattr(operation_tracker, a)) for a in attrs
48        )
49        return "{0} operations in {1:.2f}ms".format(ops, total_time)
50
51    def title(self):
52        return "MongoDB Operations"
53
54    def url(self):
55        return ""
56
57    def content(self):
58        context = self.context.copy()
59        context["queries"] = operation_tracker.queries
60        context["inserts"] = operation_tracker.inserts
61        context["updates"] = operation_tracker.updates
62        context["removes"] = operation_tracker.removes
63        context["slow_query_limit"] = current_app.config.get(
64            "MONGO_DEBUG_PANEL_SLOW_QUERY_LIMIT", 100
65        )
66        return self.render("panels/mongo-panel.html", context)
67