Add support for Druid Basic Auth to SQLAlchemy

pydruid/db/api.py

@@ -25,6 +25,8 @@ def connect(
         port=8082,
         path='/druid/v2/sql/',
         scheme='http',
+        user=None,
+        password=None,
         context=None,
         header=False,
     ):  # noqa: E125
@@ -36,7 +38,7 @@ def connect(
 
     """
     context = context or {}
-    return Connection(host, port, path, scheme, context, header)
+    return Connection(host, port, path, scheme, user, password, context, header)
 
 
 def check_closed(f):
@@ -105,6 +107,8 @@ def __init__(
         port=8082,
         path='/druid/v2/sql/',
         scheme='http',
+        user=None,
+        password=None,
         context=None,
         header=False,
     ):
@@ -115,6 +119,8 @@ def __init__(
         self.closed = False
         self.cursors = []
         self.header = header
+        self.user = user
+        self.password = password
 
     @check_closed
     def close(self):
@@ -138,7 +144,8 @@ def commit(self):
     @check_closed
     def cursor(self):
         """Return a new Cursor Object using the connection."""
-        cursor = Cursor(self.url, self.context, self.header)
+        cursor = Cursor(self.url, self.user, self.password, self.context,
+                        self.header)
         self.cursors.append(cursor)
 
         return cursor
@@ -159,10 +166,14 @@ class Cursor(object):
 
     """Connection cursor."""
 
-    def __init__(self, url, context=None, header=False):
+    def __init__(self, url, user=None, password=None, context=None,
+                 header=False):
         self.url = url
         self.context = context or {}
         self.header = header
+        self.url = url
+        self.user = user
+        self.password = password
 
         # This read/write attribute specifies the number of rows to fetch at a
         # time with .fetchmany(). It defaults to 1 meaning to fetch a single
@@ -288,7 +299,10 @@ def _stream_query(self, query):
             'header': self.header,
         }
 
-        r = requests.post(self.url, stream=True, headers=headers, json=payload)
+        auth = requests.auth.HTTPBasicAuth(self.user,
+                                           self.password) if self.user else None
+        r = requests.post(self.url, stream=True, headers=headers, json=payload,
+                          auth=auth)
         if r.encoding is None:
             r.encoding = 'utf-8'
         # raise any error messages

pydruid/db/sqlalchemy.py

@@ -95,6 +95,8 @@ class DruidDialect(default.DefaultDialect):
     name = 'druid'
     scheme = 'http'
     driver = 'rest'
+    user = None
+    password = None
     preparer = DruidIdentifierPreparer
     statement_compiler = DruidCompiler
     type_compiler = DruidTypeCompiler
@@ -120,6 +122,8 @@ def create_connect_args(self, url):
         kwargs = {
             'host': url.host,
             'port': url.port or 8082,
+            'user': url.username or None,
+            'password': url.password or None,
             'path': url.database,
             'scheme': self.scheme,
             'context': self.context,

tests/db/test_cursor.py

@@ -54,11 +54,12 @@ def test_context(self, requests_post_mock):
         query = 'SELECT * FROM table'
         context = {'source': 'unittest'}
 
-        cursor = Cursor(url, context)
+        cursor = Cursor(url, user=None, password=None, context=context)
         cursor.execute(query)
 
         requests_post_mock.assert_called_with(
             'http://example.com/',
+            auth=None,
             stream=True,
             headers={'Content-Type': 'application/json'},
             json={'query': query, 'context': context, 'header': False},