From 9497d8b6191998072b3fafb72c4110db387a8a5e Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Tue, 5 Mar 2019 19:09:07 +0200
Subject: [PATCH 1/2] Pass context to Druid (#152)

* Pass context to Druid

* Trigger checks.

* Fix lint
---
 .gitignore               |  1 +
 pydruid/db/api.py        | 20 +++++++++++++++-----
 pydruid/db/sqlalchemy.py |  5 +++++
 tests/db/test_cursor.py  | 21 +++++++++++++++++++++
 4 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/.gitignore b/.gitignore
index a73bd31d..11c93964 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,3 +14,4 @@ __pycache__
 *~
 .tox
 env
+venv
diff --git a/pydruid/db/api.py b/pydruid/db/api.py
index 9bfdbded..43c19ea2 100644
--- a/pydruid/db/api.py
+++ b/pydruid/db/api.py
@@ -20,7 +20,13 @@ 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',
+    context=None,
+        ):
     """
     Constructor for creating a connection to the database.
 
@@ -28,7 +34,8 @@ def connect(host='localhost', port=8082, path='/druid/v2/sql/', scheme='http'):
         >>> curs = conn.cursor()
 
     """
-    return Connection(host, port, path, scheme)
+    context = context or {}
+    return Connection(host, port, path, scheme, context)
 
 
 def check_closed(f):
@@ -97,10 +104,12 @@ def __init__(
         port=8082,
         path='/druid/v2/sql/',
         scheme='http',
+        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 = []
 
@@ -126,7 +135,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.context)
         self.cursors.append(cursor)
 
         return cursor
@@ -147,8 +156,9 @@ class Cursor(object):
 
     """Connection cursor."""
 
-    def __init__(self, url):
+    def __init__(self, url, context=None):
         self.url = url
+        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
@@ -262,7 +272,7 @@ def _stream_query(self, query):
         self.description = None
 
         headers = {'Content-Type': 'application/json'}
-        payload = {'query': query}
+        payload = {'query': query, 'context': self.context}
         r = requests.post(self.url, stream=True, headers=headers, json=payload)
         if r.encoding is None:
             r.encoding = 'utf-8'
diff --git a/pydruid/db/sqlalchemy.py b/pydruid/db/sqlalchemy.py
index ef233d31..9f5edd18 100644
--- a/pydruid/db/sqlalchemy.py
+++ b/pydruid/db/sqlalchemy.py
@@ -108,6 +108,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
@@ -118,6 +122,7 @@ def create_connect_args(self, url):
             'port': url.port or 8082,
             'path': url.database,
             'scheme': self.scheme,
+            'context': self.context,
         }
         return ([], kwargs)
 
diff --git a/tests/db/test_cursor.py b/tests/db/test_cursor.py
index f1d6599f..b9560716 100644
--- a/tests/db/test_cursor.py
+++ b/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()

From fa7f4f2d415e5901517900595b06548f6e6b8227 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Wed, 6 Mar 2019 07:15:25 +0200
Subject: [PATCH 2/2] Improve error message (#150)

* Improve error message

* Fix error payload

* Fix lint
---
 pydruid/db/api.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/pydruid/db/api.py b/pydruid/db/api.py
index 43c19ea2..5e59b1df 100644
--- a/pydruid/db/api.py
+++ b/pydruid/db/api.py
@@ -279,7 +279,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)
             )