From 0a4cc60e66924544d9835f8f876b10bec846f81c Mon Sep 17 00:00:00 2001
From: Robert Rollins <rrollins@caltech.edu>
Date: Thu, 27 Feb 2020 14:50:55 -0800
Subject: [PATCH] Implemented qinspect as a python context manager.

---
 qinspect/manager.py    | 31 +++++++++++++++++++++++++++++++
 qinspect/middleware.py |  1 +
 2 files changed, 32 insertions(+)
 create mode 100644 qinspect/manager.py

diff --git a/qinspect/manager.py b/qinspect/manager.py
new file mode 100644
index 0000000..99aa476
--- /dev/null
+++ b/qinspect/manager.py
@@ -0,0 +1,31 @@
+import time
+from django.db import connection
+from qinspect.middleware import QueryInspectMiddleware, cfg
+
+
+class QueryInspectContextManager():
+    """
+    Code called within this context manager is profiled by django-queryinspect.
+    The results will be presented when the context manager exits.
+    """
+
+    def __enter__(self):
+        if cfg['enabled']:
+            self.request_start = time.time()
+            self.conn_queries_len = len(connection.queries)
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        if cfg['enabled']:
+            request_time = time.time() - self.request_start
+
+            infos = QueryInspectMiddleware.get_query_infos(
+                connection.queries[self.conn_queries_len:])
+
+            num_duplicates = QueryInspectMiddleware.check_duplicates(infos)
+            QueryInspectMiddleware.check_stddev_limit(infos)
+            QueryInspectMiddleware.check_absolute_limit(infos)
+            # Can call ootput_states() with response={} because the only thing
+            # it does to response is add keys.
+            QueryInspectMiddleware.output_stats(
+                infos, num_duplicates, request_time, {})
diff --git a/qinspect/middleware.py b/qinspect/middleware.py
index cb74e25..8525668 100644
--- a/qinspect/middleware.py
+++ b/qinspect/middleware.py
@@ -140,6 +140,7 @@ def check_duplicates(cls, infos):
 
         return n
 
+    @classmethod
     def check_stddev_limit(cls, infos):
         total = sum(qi.time for qi in infos)
         n = len(infos)