Make URLLibFetcher aware of basic auth info in scheduler URL.

luigi/rpc.py

@@ -24,9 +24,11 @@
 import logging
 import socket
 import time
+import base64
 
+from luigi import six
 from luigi.six.moves.urllib.parse import urljoin, urlencode, urlparse
-from luigi.six.moves.urllib.request import urlopen
+from luigi.six.moves.urllib.request import urlopen, Request
 from luigi.six.moves.urllib.error import URLError
 
 from luigi import configuration
@@ -71,9 +73,31 @@ def __init__(self, message, sub_exception=None):
 class URLLibFetcher(object):
     raises = (URLError, socket.timeout)
 
+    def _create_request(self, full_url, body=None):
+        # when full_url contains basic auth info, extract it and set the Authorization header
+        url = urlparse(full_url)
+        if url.username:
+            # base64 encoding of username:password
+            auth = base64.b64encode(six.b('{}:{}'.format(url.username, url.password or '')))
+            if six.PY3:
+                auth = auth.decode('utf-8')
+
+            # update full_url and create a request object with the auth header set
+            full_url = url._replace(netloc=url.netloc.split('@', 1)[-1]).geturl()
+            req = Request(full_url)
+            req.add_header('Authorization', 'Basic {}'.format(auth))
+        else:
+            req = Request(full_url)
+
+        # add the request body
+        if body:
+            req.data = urlencode(body).encode('utf-8')
+
+        return req
+
     def fetch(self, full_url, body, timeout):
-        body = urlencode(body).encode('utf-8')
-        return urlopen(full_url, body, timeout).read().decode('utf-8')
+        req = self._create_request(full_url, body=body)
+        return urlopen(req, timeout=timeout).read().decode('utf-8')
 
 
 class RequestsFetcher(object):

test/rpc_test.py

@@ -29,6 +29,7 @@
 import socket
 from multiprocessing import Process, Queue
 import requests
+from luigi import six
 
 
 class RemoteSchedulerTest(unittest.TestCase):
@@ -168,3 +169,45 @@ def check_session(q):
         p.join()
 
         self.assertTrue(q.get(), 'the requests.Session should have changed in the new process')
+
+
+class URLLibFetcherTest(ServerTestBase):
+
+    def test_url_with_basic_auth(self):
+        fetcher = luigi.rpc.URLLibFetcher()
+
+        # without password
+        req = fetcher._create_request('http://user@localhost')
+        self.assertTrue(req.has_header('Authorization'))
+        self.assertEqual(req.get_header('Authorization'), 'Basic dXNlcjo=')
+        self.assertEqual(req.get_full_url(), 'http://localhost')
+
+        # empty password (same as above)
+        req = fetcher._create_request('http://user:@localhost')
+        self.assertTrue(req.has_header('Authorization'))
+        self.assertEqual(req.get_header('Authorization'), 'Basic dXNlcjo=')
+        self.assertEqual(req.get_full_url(), 'http://localhost')
+
+        # with password
+        req = fetcher._create_request('http://user:pass@localhost')
+        self.assertTrue(req.has_header('Authorization'))
+        self.assertEqual(req.get_header('Authorization'), 'Basic dXNlcjpwYXNz')
+        self.assertEqual(req.get_full_url(), 'http://localhost')
+
+    def test_url_without_basic_auth(self):
+        fetcher = luigi.rpc.URLLibFetcher()
+        req = fetcher._create_request('http://localhost')
+
+        self.assertFalse(req.has_header('Authorization'))
+        self.assertEqual(req.get_full_url(), 'http://localhost')
+
+    def test_body_encoding(self):
+        fetcher = luigi.rpc.URLLibFetcher()
+
+        # with body
+        req = fetcher._create_request('http://localhost', body={'foo': 'bar baz/test'})
+        self.assertEqual(req.data, six.b('foo=bar+baz%2Ftest'))
+
+        # without body
+        req = fetcher._create_request('http://localhost')
+        self.assertIsNone(req.data)