Add support for Druid Basic Auth to SQLAlchemy

This adds support for basic authentication (in the same format as the other sqlalchemy drivers) `druid://USER:PASSWORD@localhost:8082/druid/v2/sql/` Signed-off-by: Don Bowman <[email protected]>
druid-io · Feb 1, 2019 · 32ea4e1 · 32ea4e1
1 parent 0a59a70
commit 32ea4e1
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 5 deletions.
diff --git a/pydruid/db/api.py b/pydruid/db/api.py
@@ -20,15 +20,16 @@ class Type(object):
     BOOLEAN = 3
 
 
-def connect(host='localhost', port=8082, path='/druid/v2/sql/', scheme='http'):
+def connect(host='localhost', port=8082, path='/druid/v2/sql/', scheme='http',
+            user='', password=''):
     """
     Constructor for creating a connection to the database.
 
         >>> conn = connect('localhost', 8082)
         >>> curs = conn.cursor()
 
     """
-    return Connection(host, port, path, scheme)
+    return Connection(host, port, path, scheme, user, password)
 
 
 def check_closed(f):
@@ -97,12 +98,16 @@ def __init__(
         port=8082,
         path='/druid/v2/sql/',
         scheme='http',
+        user='',
+        password='',
     ):
         netloc = '{host}:{port}'.format(host=host, port=port)
         self.url = parse.urlunparse(
             (scheme, netloc, path, None, None, None))
         self.closed = False
         self.cursors = []
+        self.user = user
+        self.password = password
 
     @check_closed
     def close(self):
@@ -126,7 +131,7 @@ def commit(self):
     @check_closed
     def cursor(self):
         """Return a new Cursor Object using the connection."""
-        cursor = Cursor(self.url)
+        cursor = Cursor(self.url, self.user, self.password)
         self.cursors.append(cursor)
 
         return cursor
@@ -147,8 +152,10 @@ class Cursor(object):
 
     """Connection cursor."""
 
-    def __init__(self, url):
+    def __init__(self, url, user='', password=''):
         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
@@ -263,7 +270,9 @@ def _stream_query(self, query):
 
         headers = {'Content-Type': 'application/json'}
         payload = {'query': query}
-        r = requests.post(self.url, stream=True, headers=headers, json=payload)
+        auth = requests.auth.HTTPBasicAuth(self.user, self.password)
+        r = requests.post(self.url, stream=True, headers=headers, json=payload,
+                          auth=auth)
         if r.encoding is None:
             r.encoding = 'utf-8'
 

diff --git a/pydruid/db/sqlalchemy.py b/pydruid/db/sqlalchemy.py
@@ -95,6 +95,8 @@ class DruidDialect(default.DefaultDialect):
     name = 'druid'
     scheme = 'http'
     driver = 'rest'
+    user = ''
+    password = ''
     preparer = DruidIdentifierPreparer
     statement_compiler = DruidCompiler
     type_compiler = DruidTypeCompiler
@@ -116,6 +118,8 @@ def create_connect_args(self, url):
         kwargs = {
             'host': url.host,
             'port': url.port or 8082,
+            'user': url.username or '',
+            'password': url.password or '',
             'path': url.database,
             'scheme': self.scheme,
         }