Merge "Add redirection handling to openstack.session.Session"

openstack/session.py

@@ -38,10 +38,14 @@ class Session(requests.Session):
 
     _user_agent = DEFAULT_USER_AGENT
 
+    REDIRECT_STATUSES = (301, 302, 303, 305, 307)
+    DEFAULT_REDIRECT_LIMIT = 30
+
     def __init__(
             self,
             user_agent=None,
             verify=True,
+            redirect=DEFAULT_REDIRECT_LIMIT,
     ):
         """Wraps requests.Session to add some OpenStack-specific features
 
@@ -51,6 +55,11 @@ class Session(requests.Session):
         :param boolean/string verify: If ``True``, the SSL cert will be
                                       verified. A CA_BUNDLE path can also be
                                       provided.
+        :param boolean/integer redirect: (integer) The maximum number of
+                                         redirections followed in a request.
+                                         (boolean) No redirections if False,
+                                         requests.Session handles redirection
+                                         if True. (optional)
 
         User agent handling is as follows:
 
@@ -65,12 +74,18 @@ class Session(requests.Session):
         if user_agent:
             self._user_agent = user_agent
         self.verify = verify
+        self._redirect = redirect
 
-    def request(self, method, url, **kwargs):
+    def request(self, method, url, redirect=None, **kwargs):
         """Send a request
 
         :param string method: Request HTTP method
         :param string url: Request URL
+        :param boolean/integer redirect: (integer) The maximum number of
+                                         redirections followed in a request.
+                                         (boolean) No redirections if False,
+                                         requests.Session handles redirection
+                                         if True. (optional)
 
         The following additional kw args are supported:
         :param object json: Request body to be encoded as JSON
@@ -101,12 +116,63 @@ class Session(requests.Session):
         else:
             headers.setdefault('User-Agent', DEFAULT_USER_AGENT)
 
+        if redirect is None:
+            redirect = self._redirect
+
+        if isinstance(redirect, bool) and redirect:
+            # Fall back to requests redirect handling
+            kwargs['allow_redirects'] = True
+        else:
+            # Force disable requests redirect handling, we will manage
+            # redirections below
+            kwargs['allow_redirects'] = False
+
         self._log_request(method, url, **kwargs)
 
+        resp = self._send_request(method, url, redirect, **kwargs)
+
+        self._log_response(resp)
+
+        return resp
+
+    def _send_request(self, method, url, redirect, **kwargs):
+        # NOTE(jamielennox): We handle redirection manually because the
+        # requests lib follows some browser patterns where it will redirect
+        # POSTs as GETs for certain statuses which is not want we want for an
+        # API. See: https://en.wikipedia.org/wiki/Post/Redirect/Get
+
         resp = super(Session, self).request(method, url, **kwargs)
 
         self._log_response(resp)
 
+        if resp.status_code in self.REDIRECT_STATUSES:
+            # Be careful here in python True == 1 and False == 0
+            if isinstance(redirect, bool):
+                redirect_allowed = redirect
+            else:
+                redirect -= 1
+                redirect_allowed = redirect >= 0
+
+            if redirect_allowed:
+                try:
+                    location = resp.headers['location']
+                except KeyError:
+                    _logger.warn(
+                        "Redirection from %s failed, no location provided",
+                        resp.url,
+                    )
+                else:
+                    new_resp = self._send_request(
+                        method,
+                        location,
+                        redirect,
+                        **kwargs
+                    )
+
+                    new_resp.history = list(new_resp.history)
+                    new_resp.history.insert(0, resp)
+                    resp = new_resp
+
         return resp
 
     def _log_request(self, method, url, **kwargs):

openstack/tests/test_session.py

@@ -17,6 +17,8 @@ import six
 import fixtures
 import httpretty
 
+import requests
+
 from openstack import session
 from openstack.tests import base
 
@@ -24,6 +26,7 @@ from openstack.tests import base
 fake_url = 'http://www.root.url'
 fake_request = 'Now is the time...'
 fake_response = 'for the quick brown fox...'
+fake_redirect = 'redirect text'
 
 fake_record1 = {
     'key1': {
@@ -429,3 +432,94 @@ class TestSessionDebug(TestSessionBase):
         for k, v in six.iteritems(headers):
             self.assertIn(k, self.log_fixture.output)
             self.assertIn(v, self.log_fixture.output)
+
+
+class TestSessionRedirects(TestSessionBase):
+
+    REDIRECT_CHAIN = [
+        'http://myhost:3445/',
+        'http://anotherhost:6555/',
+        'http://thirdhost/',
+        'http://finaldestination:55/',
+    ]
+
+    def setup_redirects(
+            self,
+            method=httpretty.GET,
+            status=305,
+            redirect_kwargs={},
+            final_kwargs={},
+    ):
+        redirect_kwargs.setdefault('body', fake_redirect)
+
+        for s, d in zip(self.REDIRECT_CHAIN, self.REDIRECT_CHAIN[1:]):
+            httpretty.register_uri(
+                method,
+                s,
+                status=status,
+                location=d,
+                **redirect_kwargs
+            )
+
+        final_kwargs.setdefault('status', 200)
+        final_kwargs.setdefault('body', fake_response)
+        httpretty.register_uri(method, self.REDIRECT_CHAIN[-1], **final_kwargs)
+
+    @httpretty.activate
+    def test_get_redirect(self):
+        self.setup_redirects()
+        sess = session.Session()
+        resp = sess.get(self.REDIRECT_CHAIN[-2])
+        self.assertResponseOK(resp)
+
+    @httpretty.activate
+    def test_post_keeps_correct_method(self):
+        self.setup_redirects(method=httpretty.POST, status=301)
+        sess = session.Session()
+        resp = sess.post(self.REDIRECT_CHAIN[-2])
+        self.assertResponseOK(resp)
+
+    @httpretty.activate
+    def test_redirect_forever(self):
+        self.setup_redirects()
+        sess = session.Session()
+        resp = sess.get(self.REDIRECT_CHAIN[0])
+        self.assertResponseOK(resp)
+        # Request history length is 1 less than the source chain due to the
+        # last response not being a redirect and not added to the history.
+        self.assertEqual(len(self.REDIRECT_CHAIN) - 1, len(resp.history))
+
+    @httpretty.activate
+    def test_no_redirect(self):
+        self.setup_redirects()
+        sess = session.Session(redirect=False)
+        resp = sess.get(self.REDIRECT_CHAIN[0])
+        self.assertEqual(305, resp.status_code)
+        self.assertEqual(self.REDIRECT_CHAIN[0], resp.url)
+
+    @httpretty.activate
+    def test_redirect_limit(self):
+        self.setup_redirects()
+        for i in (1, 2):
+            sess = session.Session(redirect=i)
+            resp = sess.get(self.REDIRECT_CHAIN[0])
+            self.assertResponseOK(resp, status=305, body=fake_redirect)
+            self.assertEqual(self.REDIRECT_CHAIN[i], resp.url)
+
+    @httpretty.activate
+    def test_history_matches_requests(self):
+        self.setup_redirects(status=301)
+        sess = session.Session(redirect=True)
+        req_resp = requests.get(
+            self.REDIRECT_CHAIN[0],
+            allow_redirects=True,
+        )
+
+        ses_resp = sess.get(self.REDIRECT_CHAIN[0])
+
+        self.assertEqual(type(ses_resp.history), type(req_resp.history))
+        self.assertEqual(len(ses_resp.history), len(req_resp.history))
+
+        for r, s in zip(req_resp.history, ses_resp.history):
+            self.assertEqual(s.url, r.url)
+            self.assertEqual(s.status_code, r.status_code)