Merge pull request #6 from andrewsg/user_middleware

Wrap user apps with user middleware at load time

Author: andrewsg

multicore_runtime/wsgi_config.py

@@ -45,11 +45,43 @@ def app_for_script(script):
       logging.exception('Failed to import %s: %s', script, err)
       return None
     else:
-      return app
+      return app_wrapped_in_user_middleware(app)
+
+
+def app_wrapped_in_user_middleware(app):
+  """Returns the input WSGI app, wrapped in appengine_config middleware."""
+  add_middleware = get_add_middleware_from_appengine_config()
+  if add_middleware:
+    return add_middleware(app)
+  else:
+    return app
+
+
+def get_add_middleware_from_appengine_config():
+  """Tries to import appengine_config and return middleware; fails silently.
+
+  `appengine_config` is optionally part of GAE-compatible user code.
+  If the user chooses to define the `webapp_add_wsgi_middleware` function
+  there, then we will use it to wrap app scripts. This is used for e.g.
+  appstats. Many user apps do not have or need appengine_config, so failure
+  to import it or failure to find `webapp_add_wsgi_middleware` is not an
+  error.
+
+  Returns:
+    The appengine_config.webapp_add_wsgi_middleware function or None.
+  """
+  try:
+    import appengine_config  # pylint: disable=g-import-not-at-top
+    try:
+      return appengine_config.webapp_add_wsgi_middleware
+    except AttributeError:
+      return None
+  except ImportError:
+    return None
 
 
 def load_user_scripts_into_handlers(handlers):
-  """Preloads user scripts.
+  """Preloads user scripts, wrapped in appengine_config middleware if present.
 
   Args:
     handlers: appinfo_external.handlers data as provided by get_module_config()

multicore_runtime/wsgi_config_test.py

@@ -0,0 +1,56 @@
+# Copyright 2015 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import unittest
+
+from mock import patch
+from werkzeug.test import Client
+from werkzeug.wrappers import Request
+from werkzeug.wrappers import Response
+
+from wsgi_config import app_for_script
+
+
+@Request.application
+def salutation_world(request):
+  salutation = request.args.get('salutation', 'Hello')
+  return Response('%s World!' % salutation)
+
+
+def goodbye_world_middleware(app):
+  def goodbye_wrapper(wsgi_env, start_response):
+    wsgi_env['QUERY_STRING'] = 'salutation=Goodbye'
+    return app(wsgi_env, start_response)
+  return goodbye_wrapper
+
+
+class AppConfigTestCase(unittest.TestCase):
+
+  def test_app_for_script(self):
+    with patch('wsgi_config.get_add_middleware_from_appengine_config',
+               return_value=None):
+      app = app_for_script('wsgi_config_test.salutation_world')
+    client = Client(app, Response)
+    response = client.get('/?salutation=Hello')
+    self.assertEqual(response.status_code, 200)
+    self.assertEqual(response.data, 'Hello World!')
+
+  def test_app_for_script_with_middleware(self):
+    with patch('wsgi_config.get_add_middleware_from_appengine_config',
+               return_value=goodbye_world_middleware):
+      app = app_for_script('wsgi_config_test.salutation_world')
+    client = Client(app, Response)
+    response = client.get('/?salutation=Hello')
+    self.assertEqual(response.status_code, 200)
+    self.assertEqual(response.data, 'Goodbye World!')