From 4a0612328ea90f9ad571274e3abf685221103c04 Mon Sep 17 00:00:00 2001
From: Arik Fraimovich <arik@arikfr.com>
Date: Mon, 15 May 2017 16:25:18 +0300
Subject: [PATCH] TreasureData: improve error handling and upgrade client.

---
 redash/query_runner/treasuredata.py | 30 +++++++++++++----------------
 requirements_all_ds.txt             |  2 +-
 2 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/redash/query_runner/treasuredata.py b/redash/query_runner/treasuredata.py
index 43a95cfe9b2..56894482af8 100644
--- a/redash/query_runner/treasuredata.py
+++ b/redash/query_runner/treasuredata.py
@@ -1,9 +1,9 @@
 import json
+import logging
 
-from redash.utils import JSONEncoder
 from redash.query_runner import *
+from redash.utils import JSONEncoder
 
-import logging
 logger = logging.getLogger(__name__)
 
 try:
@@ -103,21 +103,17 @@ def run_query(self, query, user):
 
         cursor = connection.cursor()
 
-        try:
-            cursor.execute(query)
-            columns_data = [(row[0], cursor.show_job()['hive_result_schema'][i][1]) for i,row in enumerate(cursor.description)]
-
-            columns = [{'name': col[0],
-                'friendly_name': col[0],
-                'type': TD_TYPES_MAPPING.get(col[1], None)} for col in columns_data]
-
-            rows = [dict(zip(([c[0] for c in columns_data]), r)) for i, r in enumerate(cursor.fetchall())]
-            data = {'columns': columns, 'rows': rows}
-            json_data = json.dumps(data, cls=JSONEncoder)
-            error = None
-        except Exception, ex:
-            json_data = None
-            error = ex.message
+        cursor.execute(query)
+        columns_data = [(row[0], cursor.show_job()['hive_result_schema'][i][1]) for i,row in enumerate(cursor.description)]
+
+        columns = [{'name': col[0],
+            'friendly_name': col[0],
+            'type': TD_TYPES_MAPPING.get(col[1], None)} for col in columns_data]
+
+        rows = [dict(zip(([c[0] for c in columns_data]), r)) for i, r in enumerate(cursor.fetchall())]
+        data = {'columns': columns, 'rows': rows}
+        json_data = json.dumps(data, cls=JSONEncoder)
+        error = None
 
         return json_data, error
 
diff --git a/requirements_all_ds.txt b/requirements_all_ds.txt
index 03f2206866e..c4ce4147789 100644
--- a/requirements_all_ds.txt
+++ b/requirements_all_ds.txt
@@ -8,7 +8,7 @@ pyhive==0.3.0
 pymongo==3.2.1
 pyOpenSSL==0.14
 vertica-python==0.5.1
-td-client==0.4.1
+td-client==0.8.0
 pymssql==2.1.3
 dql==0.5.16
 dynamo3==0.4.7