Merge branch 'master' into add-tls-auth

.gitignore

@@ -14,3 +14,4 @@ __pycache__
 *~
 .tox
 env
+venv

pydruid/db/api.py

@@ -20,16 +20,24 @@ class Type(object):
     BOOLEAN = 3
 
 
-def connect(host='localhost', port=8082, path='/druid/v2/sql/', scheme='http',
-            user=None, password=None):
+def connect(
+    host='localhost',
+    port=8082,
+    path='/druid/v2/sql/',
+    scheme='http',
+    user=None,
+    password=None,
+    context=None,
+        ):
     """
     Constructor for creating a connection to the database.
 
         >>> conn = connect('localhost', 8082)
         >>> curs = conn.cursor()
 
     """
-    return Connection(host, port, path, scheme, user, password)
+    context = context or {}
+    return Connection(host, port, path, scheme, user, password, context)
 
 
 def check_closed(f):
@@ -100,10 +108,12 @@ def __init__(
         scheme='http',
         user=None,
         password=None,
+        context=None,
     ):
         netloc = '{host}:{port}'.format(host=host, port=port)
         self.url = parse.urlunparse(
             (scheme, netloc, path, None, None, None))
+        self.context = context or {}
         self.closed = False
         self.cursors = []
         self.user = user
@@ -131,7 +141,7 @@ def commit(self):
     @check_closed
     def cursor(self):
         """Return a new Cursor Object using the connection."""
-        cursor = Cursor(self.url, self.user, self.password)
+        cursor = Cursor(self.url, self.user, self.password, self.context)
         self.cursors.append(cursor)
 
         return cursor
@@ -152,10 +162,11 @@ class Cursor(object):
 
     """Connection cursor."""
 
-    def __init__(self, url, user=None, password=None):
+    def __init__(self, url, user=None, password=None, context=None):
         self.url = url
         self.user = user
         self.password = password
+        self.context = context or {}
 
         # 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
@@ -269,7 +280,7 @@ def _stream_query(self, query):
         self.description = None
 
         headers = {'Content-Type': 'application/json'}
-        payload = {'query': query}
+        payload = {'query': query, 'context': self.context}
         auth = requests.auth.HTTPBasicAuth(self.user,
                                            self.password) if self.user else None
         r = requests.post(self.url, stream=True, headers=headers, json=payload,
@@ -279,7 +290,14 @@ def _stream_query(self, query):
 
         # raise any error messages
         if r.status_code != 200:
-            payload = r.json()
+            try:
+                payload = r.json()
+            except Exception:
+                payload = {
+                    'error': 'Unknown error',
+                    'errorClass': 'Unknown',
+                    'errorMessage': r.text,
+                }
             msg = (
                 '{error} ({errorClass}): {errorMessage}'.format(**payload)
             )

pydruid/db/sqlalchemy.py

@@ -110,6 +110,10 @@ class DruidDialect(default.DefaultDialect):
     description_encoding = None
     supports_native_boolean = True
 
+    def __init__(self, context=None, *args, **kwargs):
+        super(DruidDialect, self).__init__(*args, **kwargs)
+        self.context = context or {}
+
     @classmethod
     def dbapi(cls):
         return pydruid.db
@@ -122,6 +126,7 @@ def create_connect_args(self, url):
             'password': url.password or None,
             'path': url.database,
             'scheme': self.scheme,
+            'context': self.context,
         }
         return ([], kwargs)
 

tests/db/test_cursor.py

@@ -43,6 +43,27 @@ def test_execute_empty_result(self, requests_post_mock):
         expected = []
         self.assertEquals(result, expected)
 
+    @patch('requests.post')
+    def test_context(self, requests_post_mock):
+        response = Response()
+        response.status_code = 200
+        response.raw = BytesIO(b'[]')
+        requests_post_mock.return_value = response
+
+        url = 'http://example.com/'
+        query = 'SELECT * FROM table'
+        context = {'source': 'unittest'}
+
+        cursor = Cursor(url, context)
+        cursor.execute(query)
+
+        requests_post_mock.assert_called_with(
+            'http://example.com/',
+            stream=True,
+            headers={'Content-Type': 'application/json'},
+            json={'query': query, 'context': context},
+        )
+
 
 if __name__ == '__main__':
     unittest.main()