Added solr route cookie for solr-service requests (#143)

* Added solr route cookie for solr-service requests

When a request is directed to solr-service, the user token is extracted
from the request and its solr route (which link to a specific solr
instance) is retreived from redis (if it exists) before passing the
request to solr-service.

Solr-service response is also inspected for "Set-Cookie" instructions
with the solr route, in case it is present the solr route is stored in
redis using the user's token as key.

* Solr route is also checked if it exists as a remote service

* Decoupling affinity enhancement/solr_route/search from adsws

* Less generic affinity decorator

.gitignore

@@ -79,3 +79,4 @@ test/mocha/sandbox.spec.js
 .idea
 
 .vagrant
+python

adsws/api/config.py

@@ -39,3 +39,5 @@
 
 API_PROXYVIEW_HEADERS = {'Cache-Control': 'public, max-age=600'}
 REMOTE_PROXY_ALLOWED_HEADERS = ['Content-Type', 'Content-Disposition']
+
+AFFINITY_ENHANCED_ENDPOINTS = {"/search": "sroute",} # keys: deploy paths, value: cookie

adsws/api/discoverer/affinity.py

@@ -0,0 +1,111 @@
+import Cookie
+from flask import request
+from functools import wraps
+from werkzeug.datastructures import Headers
+from werkzeug.datastructures import ImmutableTypeConversionDict
+
+def _get_route(storage, route_redis_prefix, user_token):
+    """
+    Obtains the solr route from redis for a given user token. It piggybacks the
+    existing rate limiter extension connection, and if it fails the request
+    will not be stopped (answering a request has a higher priority than
+    assigning a solr instance).
+    """
+    try:
+        route = storage.get(route_redis_prefix+user_token)
+    except:
+        route = None
+    return route
+
+def _set_route(storage, route_redis_prefix, user_token, route, route_redis_expiration_time):
+    """
+    Sets the solr route from redis for a given user token. It piggybacks the
+    existing rate limiter extension connection, and if it fails the request
+    will not be stopped (answering a request has a higher priority than
+    assigning a solr instance).
+
+    Keys in redis will expire in N seconds to reduce chances of saturation
+    of a particular solr and to automatically clear entries in redis.
+    """
+    try:
+        storage.setex(route_redis_prefix+user_token, route, route_redis_expiration_time)
+    except:
+        pass
+
+def _build_updated_cookies(request, user_token, route, name):
+    """
+    Based on the current request, create updated headers and cookies content
+    attributes.
+    """
+    # Interpret cookie header content
+    cookies_header = Cookie.SimpleCookie()
+    currrent_cookie_content = request.headers.get('cookie', None)
+    if currrent_cookie_content:
+        cookies_header.load(currrent_cookie_content.encode("utf8"))
+    # Interpret cookies attribute (immutable dict) as a normal dict
+    if request.cookies:
+        cookies = dict(request.cookies)
+    else:
+        cookies = {}
+    # Update cookie structures
+    if route:
+        # Create/update solr route
+        cookies_header[name] = route.encode("utf8")
+        cookies[name] = route
+    else:
+        # Discard non-registered solr route if it is present
+        cookies_header.pop(name, None)
+        cookies.pop(name, None)
+    # Transform cookies structures into the format that request requires
+    cookies_header_content = cookies_header.output(header="", sep=";")
+    cookies_content = ImmutableTypeConversionDict(cookies)
+    return cookies_header_content, cookies_content
+
+def affinity_decorator(storage, name="sroute"):
+    """
+    Assign a cookie that will be used by solr ingress to send request to
+    a specific solr instance for the same user, maximizing the use of solr
+    cache capabilities.
+
+    The storage should be a redis connection.
+    """
+    route_redis_prefix="token:{}:".format(name)
+    route_redis_expiration_time=86400 # 1 day
+
+    def real_affinity_decorator(f):
+        @wraps(f)
+        def decorated_function(*args, **kwargs):
+            # Obtain user token, giving priority to forwarded authorization field (used when a microservice uses its own token)
+            user_token = request.headers.get('X-Forwarded-Authorization', None)
+            if user_token is None:
+                user_token = request.headers.get('Authorization', None)
+            if user_token and len(user_token) > 7: # This should be always true
+                user_token = user_token[7:] # Get rid of "Bearer:" or "Bearer "
+                route = _get_route(storage, route_redis_prefix, user_token)
+                cookies_header_content, cookies_content = _build_updated_cookies(request, user_token, route, name)
+                # Update request cookies (header and cookies attributes)
+                request.headers = Headers(request.headers)
+                request.headers.set('cookie', cookies_header_content)
+                request.cookies = cookies_content
+
+            r = f(*args, **kwargs)
+            if type(r) is tuple and len(r) > 2:
+                response_headers = r[2]
+            elif hasattr(r, 'headers'):
+                response_headers = r.headers
+            else:
+                response_headers = None
+
+            if user_token and response_headers:
+                set_cookie = response_headers.get('Set-Cookie', None)
+                if set_cookie:
+                    # If solr issued a set cookie, store the value in redis linked to the user token
+                    cookie = Cookie.SimpleCookie()
+                    cookie.load(set_cookie.encode("utf8"))
+                    route = cookie.get(name, None)
+                    if route:
+                        _set_route(storage, route_redis_prefix, user_token, route.value, route_redis_expiration_time)
+            return r
+        return decorated_function
+    return real_affinity_decorator
+

adsws/api/discoverer/utils.py

@@ -1,6 +1,7 @@
 import os
 import requests
 import json
+import Cookie
 from flask.ext.headers import headers
 from flask import request
 from views import ProxyView
@@ -11,6 +12,7 @@
 from adsws.ext.ratelimiter import ratelimit, limit_func, scope_func, key_func
 from flask.ext.consulate import ConsulService
 from functools import wraps
+from .affinity import affinity_decorator
 
 def local_app_context(local_app):
     """
@@ -61,6 +63,9 @@ def bootstrap_local_module(service_uri, deploy_path, app):
         # ensure the current_app matches local_app and not API app
         view = local_app_context(local_app)(view)
 
+        if deploy_path in local_app.config.get('AFFINITY_ENHANCED_ENDPOINTS', []):
+            view = affinity_decorator(ratelimit._storage.storage, name=local_app.config['AFFINITY_ENHANCED_ENDPOINTS'].get(deploy_path))(view)
+
         # Decorate the view with ratelimit
         if hasattr(attr_base, 'rate_limit'):
             d = attr_base.rate_limit[0]
@@ -169,6 +174,9 @@ def bootstrap_remote_service(service_uri, deploy_path, app):
             properties.setdefault('rate_limit', [1000, 86400])
             properties.setdefault('scopes', [])
 
+            if deploy_path in app.config.get('AFFINITY_ENHANCED_ENDPOINTS', []):
+                view = affinity_decorator(ratelimit._storage.storage, name=app.config['AFFINITY_ENHANCED_ENDPOINTS'].get(deploy_path))(view)
+
             # Decorate the view with ratelimit.
             d = properties['rate_limit'][0]
             view = ratelimit.shared_limit_and_check(

adsws/tests/test_affinity.py

@@ -0,0 +1,68 @@
+from unittest import TestCase
+from adsws.api.discoverer import affinity
+from flask.ext.restful import Resource
+import flask
+from flask_restful import Resource, Api
+import mock
+
+class SetCookieView(Resource):
+    """
+    Returns a good HTTP answer with a coockie set in the headers
+    """
+    storage = None
+    @affinity.affinity_decorator(storage, name="sroute")
+    def get(self):
+        return {}, 200, {'Set-Cookie': 'sroute=aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa; Path=/; HttpOnly'}
+
+class DontSetCookieView(Resource):
+    """
+    Returns a good HTTP answer without any coockie set in the headers
+    """
+    storage = None
+    @affinity.affinity_decorator(storage, name="sroute")
+    def get(self):
+        return {}, 200, {}
+
+
+class AffinityRouteTestCase(TestCase):
+    """
+    Tests solr route decorator
+    """
+
+    def setUp(self):
+        super(self.__class__, self).setUp()
+        app = flask.Flask(__name__)
+        api = Api(app)
+        api.add_resource(SetCookieView, '/set_cookie')
+        api.add_resource(DontSetCookieView, '/dont_set_cookie')
+        self.app = app.test_client()
+
+
+    def tearDown(self):
+        super(self.__class__, self).tearDown()
+
+
+    def test_set_cookie(self):
+        """
+        Test that the cookie is set
+        """
+        affinity._get_route = mock.Mock(return_value="zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz")
+        affinity._set_route = mock.Mock()
+        rv = self.app.get('/set_cookie', headers=[['Authorization', "Bearer:TTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTT"]])
+        self.assertIn('Set-Cookie', rv.headers)
+        self.assertEquals(rv.headers['Set-Cookie'], 'sroute=aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa; Path=/; HttpOnly')
+        affinity._get_route.assert_called_once()
+        affinity._set_route.assert_called_once()
+
+
+    def test_set_cookie(self):
+        """
+        Test that no cookie is set
+        """
+        affinity._get_route = mock.Mock(return_value=None)
+        affinity._set_route = mock.Mock()
+        rv = self.app.get('/dont_set_cookie', headers=[['Authorization', "Bearer:TTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTTT"]])
+        self.assertNotIn('Set-Cookie', rv.headers)
+        affinity._get_route.assert_called_once()
+        affinity._set_route.assert_not_called()
+