updates based on PR comments

washort · Jul 14, 2017 · 3a0ec4a · 3a0ec4a
1 parent 4af5b2e
commit 3a0ec4a
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 9 deletions.
diff --git a/redash/cli/data_sources.py b/redash/cli/data_sources.py
@@ -57,7 +57,7 @@ def test(name, organization='default'):
             name, data_source.id)
         try:
             data_source.query_runner.test_connection()
-        except Exception, e:
+        except Exception as e:
             print "Failure: {}".format(e)
             exit(1)
         else:
@@ -82,7 +82,7 @@ def get_data_source_version(name, organization='default'):
             name, data_source.id)
         try:
             info = data_source.query_runner.get_data_source_version()
-        except Exception, e:
+        except Exception as e:
             print "Failure: {}".format(e)
             exit(1)
         else:

diff --git a/redash/query_runner/__init__.py b/redash/query_runner/__init__.py
@@ -77,16 +77,21 @@ def get_data_source_version(self):
         if self.data_source_version_query is None:
             raise NotImplementedError
         data, error = self.run_query(self.data_source_version_query, None)
-        version = json.loads(data)['rows'][0]['version']
-        if(self.data_source_version_post_process == "split by space take second"):
+
+        if error is not None:
+            raise Exception(error)
+
+        try:
+            version = json.loads(data)['rows'][0]['version']
+        except KeyError as e:
+            raise Exception(e)
+
+        if self.data_source_version_post_process == "split by space take second":
             version = version.split(" ")[1]
-        elif(self.data_source_version_post_process == "split by space take last"):
+        elif self.data_source_version_post_process == "split by space take last":
             version = version.split(" ")[-1]
-        elif(self.data_source_version_post_process == "none"):
+        elif self.data_source_version_post_process == "none":
             version = version
-
-        if error is not None:
-            raise Exception(error)
 
         return version