Updated middleware and views to match upstream

flaskext/gae_mini_profiler/__init__.py

@@ -146,18 +146,30 @@ def _render(self, template_name, context):
     def _request_view(self):
         """Renders the request stats."""
 
-        request_id = request.args['request_id']
-        stats = profiler.RequestStats.get(request_id)
-
-        if not stats:
-            return u''
-
-        dict_request_stats = {}
-        for property in profiler.RequestStats.serialized_properties:
-            dict_request_stats[property] = \
-                    stats.__getattribute__(property)
-
-        return jsonify(dict_request_stats)
+        request_ids = request.args['request_ids']
+
+        stats_list = []
+        for request_id in request_ids.split(','):
+            request_stats = profiler.RequestStats.get(request_id)
+
+            if request_stats and not request_stats.disabled:
+                dict_request_stats = {}
+                for property in profiler.RequestStats.serialized_properties:
+                    dict_request_stats[property] = \
+                            request_stats.__getattribute__(property)
+
+                stats_list.append(dict_request_stats)
+
+                # Don't show temporary redirect profiles more than once
+                # automatically, as they are tied to URL params and may be
+                # copied around easily.
+                if request_stats.temporary_redirect:
+                    request_stats.disabled = True
+                    request_stats.store()
+
+        # For security reasons, we return an object instead of a list as it is
+        # dont in the upstream module.
+        return jsonify(stats=stats_list)
 
     def _share_view(self):
         """Renders the shared stats view."""

flaskext/gae_mini_profiler/profiler.py

@@ -7,14 +7,12 @@
     :license: MIT, see LICENSE for more details
 """
 
-
 import datetime
-import logging
 import os
 import pickle
+import re
 import simplejson
 import StringIO
-import sys
 from types import GeneratorType
 import zlib
 
@@ -24,26 +22,43 @@
 # request_id is a per-request identifier accessed by a couple other pieces of gae_mini_profiler
 request_id = None
 
-
 class RequestStatsHandler(RequestHandler):
 
     def get(self):
 
         self.response.headers["Content-Type"] = "application/json"
 
-        request_stats = RequestStats.get(self.request.get("request_id"))
-        if not request_stats:
-            return
+        list_request_ids = []
+
+        request_ids = self.request.get("request_ids")
+        if request_ids:
+            list_request_ids = request_ids.split(",")
+
+        list_request_stats = []
+
+        for request_id in list_request_ids:
+
+            request_stats = RequestStats.get(request_id)
 
-        dict_request_stats = {}
-        for property in RequestStats.serialized_properties:
-            dict_request_stats[property] = request_stats.__getattribute__(property)
+            if request_stats and not request_stats.disabled:
 
-        self.response.out.write(simplejson.dumps(dict_request_stats))
+                dict_request_stats = {}
+                for property in RequestStats.serialized_properties:
+                    dict_request_stats[property] = request_stats.__getattribute__(property)
+
+                list_request_stats.append(dict_request_stats)
+
+                # Don't show temporary redirect profiles more than once automatically, as they are
+                # tied to URL params and may be copied around easily.
+                if request_stats.temporary_redirect:
+                    request_stats.disabled = True
+                    request_stats.store()
+
+        self.response.out.write(simplejson.dumps(list_request_stats))
 
 class RequestStats(object):
 
-    serialized_properties = ["request_id", "url", "url_short", "s_dt", "profiler_results", "appstats_results"]
+    serialized_properties = ["request_id", "url", "url_short", "s_dt", "profiler_results", "appstats_results", "temporary_redirect"]
 
     def __init__(self, request_id, environ, middleware):
         self.request_id = request_id
@@ -61,6 +76,9 @@ def __init__(self, request_id, environ, middleware):
         self.profiler_results = RequestStats.calc_profiler_results(middleware)
         self.appstats_results = RequestStats.calc_appstats_results(middleware)
 
+        self.temporary_redirect = middleware.temporary_redirect
+        self.disabled = False
+
     def store(self):
         # Store compressed results so we stay under the memcache 1MB limit
         pickled = pickle.dumps(self)
@@ -102,7 +120,9 @@ def short_method_fmt(s):
     def short_rpc_file_fmt(s):
         if not s:
             return ""
-        return s[s.find("/"):]
+        if "/" in s:
+            return s[s.find("/"):]
+        return s
 
     @staticmethod
     def calc_profiler_results(middleware):
@@ -155,15 +175,22 @@ def calc_appstats_results(middleware):
             calls = []
             service_totals_dict = {}
             likely_dupes = False
+            end_offset_last = 0
 
             dict_requests = {}
 
             appstats_key = long(middleware.recorder.start_timestamp * 1000)
 
             for trace in middleware.recorder.traces:
                 total_call_count += 1
+
                 total_time += trace.duration_milliseconds()
 
+                # Don't accumulate total RPC time for traces that overlap asynchronously
+                if trace.start_offset_milliseconds() < end_offset_last:
+                    total_time -= (end_offset_last - trace.start_offset_milliseconds())
+                end_offset_last = trace.start_offset_milliseconds() + trace.duration_milliseconds()
+
                 service_prefix = trace.service_call_name()
 
                 if "." in service_prefix:
@@ -232,6 +259,7 @@ def __init__(self, app):
         self.app_clean = app
         self.prof = None
         self.recorder = None
+        self.temporary_redirect = False
 
     def __call__(self, environ, start_response):
 
@@ -242,18 +270,28 @@ def __call__(self, environ, start_response):
         self.app = self.app_clean
         self.prof = None
         self.recorder = None
+        self.temporary_redirect = False
 
         if self.should_profile(environ):
 
             # Set a random ID for this request so we can look up stats later
             import base64
-            import os
-            request_id = base64.urlsafe_b64encode(os.urandom(15))
+            request_id = base64.urlsafe_b64encode(os.urandom(5))
 
             # Send request id in headers so jQuery ajax calls can pick
             # up profiles.
             def profiled_start_response(status, headers, exc_info = None):
+
+                if status.startswith("302 "):
+                    # Temporary redirect. Add request identifier to redirect location
+                    # so next rendered page can show this request's profile.
+                    headers = ProfilerWSGIMiddleware.headers_with_modified_redirect(environ, headers)
+                    self.temporary_redirect = True
+
+                # Append headers used when displaying profiler results from ajax requests
                 headers.append(("X-MiniProfiler-Id", request_id))
+                headers.append(("X-MiniProfiler-QS", environ.get("QUERY_STRING")))
+
                 return start_response(status, headers, exc_info)
 
             # Configure AppStats output, keeping a high level of request
@@ -303,3 +341,32 @@ def wrapped_appstats_app(environ, start_response):
             result = self.app(environ, start_response)
             for value in result:
                 yield value
+
+    @staticmethod
+    def headers_with_modified_redirect(environ, headers):
+        headers_modified = []
+
+        for header in headers:
+            if header[0] == "Location":
+                reg = re.compile("mp-r-id=([^&]+)")
+
+                # Keep any chain of redirects around
+                request_id_chain = request_id
+                match = reg.search(environ.get("QUERY_STRING"))
+                if match:
+                    request_id_chain = ",".join([match.groups()[0], request_id])
+
+                # Remove any pre-existing miniprofiler redirect id
+                location = header[1]
+                location = reg.sub("", location)
+
+                # Add current request id as miniprofiler redirect id
+                location += ("&" if "?" in location else "?")
+                location = location.replace("&&", "&")
+                location += "mp-r-id=%s" % request_id_chain
+
+                headers_modified.append((header[0], location))
+            else:
+                headers_modified.append(header)
+
+        return headers_modified

tests/extension.py

@@ -165,19 +165,27 @@ def test_request_view(self):
             from flaskext.gae_mini_profiler import GAEMiniProfiler
             properties = ['a', 'b', 'c']
 
-            request_id = mock.Sentinel()
-            request.args = {'request_id': request_id}
+            request_ids = ['1', '2', '3']
+            request.args = {'request_ids': ','.join(request_ids)}
 
             stats = profiler.RequestStats.get.return_value
             stats.__getattribute__ = lambda x: x
+            stats.disabled = False
+            stats.temporary_redirect = False
+
             profiler.RequestStats.serialized_properties = properties
             app = mock.Mock()
             ext = GAEMiniProfiler(app)
             ext._request_view()
 
-            profiler.RequestStats.get.assertCalledOnceWith(request_id)
-            jsonify.assert_called_once_with(
-                dict(zip(properties, properties)))
+            self.assertEquals(len(request_ids),
+                              profiler.RequestStats.get.call_count)
+
+            expected_result = []
+            for id in request_ids:
+                expected_result.append(dict(zip(properties, properties)))
+
+            jsonify.assert_called_once_with(stats=expected_result)
         finally:
             request_patcher.stop()
             profiler_patcher.stop()