Support authentication for URL data source (re #330)

redash/query_runner/__init__.py

@@ -1,14 +1,16 @@
-import sys
 import logging
 import json
+import sys
+
+import requests
 
-from collections import OrderedDict
 from redash import settings
 
 logger = logging.getLogger(__name__)
 
 __all__ = [
     'BaseQueryRunner',
+    'BaseHTTPQueryRunner',
     'InterruptException',
     'BaseSQLQueryRunner',
     'TYPE_DATETIME',
@@ -90,7 +92,7 @@ def get_data_source_version(self):
             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":
@@ -169,6 +171,107 @@ def _get_tables_stats(self, tables_dict):
                 tables_dict[t]['size'] = res[0]['cnt']
 
 
+class BaseHTTPQueryRunner(BaseQueryRunner):
+    response_error = "Endpoint returned unexpected status code"
+    requires_authentication = False
+    url_title = 'URL base path'
+    username_title = 'HTTP Basic Auth Username'
+    password_title = 'HTTP Basic Auth Password'
+
+    @classmethod
+    def configuration_schema(cls):
+        schema = {
+            'type': 'object',
+            'properties': {
+                'url': {
+                    'type': 'string',
+                    'title': cls.url_title,
+                },
+                'username': {
+                    'type': 'string',
+                    'title': cls.username_title,
+                },
+                'password': {
+                    'type': 'string',
+                    'title': cls.password_title,
+                },
+                "doc_url": {
+                    "type": "string",
+                    "title": "Documentation URL",
+                    "default": cls.default_doc_url,
+                },
+                "toggle_table_string": {
+                    "type": "string",
+                    "title": "Toggle Table String",
+                    "default": "_v",
+                    "info": (
+                        "This string will be used to toggle visibility of "
+                        "tables in the schema browser when editing a query "
+                        "in order to remove non-useful tables from sight."
+                    ),
+                }
+            },
+            'required': ['url'],
+            'secret': ['password']
+        }
+        if cls.requires_authentication:
+            schema['required'] += ['username', 'password']
+        return schema
+
+    def get_auth(self):
+        username = self.configuration.get('username')
+        password = self.configuration.get('password')
+        if username and password:
+            return (username, password)
+        if self.requires_authentication:
+            raise ValueError("Username and Password required")
+        else:
+            return None
+
+    def get_response(self, url, auth=None, **kwargs):
+        # Get authentication values if not given
+        if auth is None:
+            auth = self.get_auth()
+
+        # Then call requests to get the response from the given endpoint
+        # URL optionally, with the additional requests parameters.
+        error = None
+        response = None
+        try:
+            response = requests.get(url, auth=auth, **kwargs)
+            # Raise a requests HTTP exception with the appropriate reason
+            # for 4xx and 5xx response status codes which is later caught
+            # and passed back.
+            response.raise_for_status()
+
+            # Any other responses (e.g. 2xx and 3xx):
+            if response.status_code != 200:
+                error = '{} ({}).'.format(
+                    self.response_error,
+                    response.status_code,
+                )
+
+        except requests.HTTPError as exc:
+            logger.exception(exc)
+            error = (
+                "Failed to execute query. "
+                "Return Code: {} Reason: {}".format(
+                    response.status_code,
+                    response.text
+                )
+            )
+        except requests.RequestException as exc:
+            # Catch all other requests exceptions and return the error.
+            logger.exception(exc)
+            error = str(exc)
+        except Exception as exc:
+            # Catch any other exceptions, log it and reraise it.
+            logger.exception(exc)
+            raise sys.exc_info()[1], None, sys.exc_info()[2]
+
+        return response, error
+
+
 query_runners = {}
 
 

redash/query_runner/jql.py

@@ -1,5 +1,4 @@
 import json
-import requests
 import re
 
 from collections import OrderedDict
@@ -137,41 +136,15 @@ def get_dict_output_field_name(cls,field_name, member_name):
         return None
 
 
-class JiraJQL(BaseQueryRunner):
+class JiraJQL(BaseHTTPQueryRunner):
     noop_query = '{"queryType": "count"}'
     default_doc_url = ("https://confluence.atlassian.com/jirasoftwarecloud/"
                        "advanced-searching-764478330.html")
-
-    @classmethod
-    def configuration_schema(cls):
-        return {
-            'type': 'object',
-            'properties': {
-                'url': {
-                    'type': 'string',
-                    'title': 'JIRA URL'
-                },
-                'username': {
-                    'type': 'string',
-                },
-                'password': {
-                    'type': 'string'
-                },
-                "doc_url": {
-                    "type": "string",
-                    "title": "Documentation URL",
-                    "default": cls.default_doc_url
-                },
-                "toggle_table_string": {
-                    "type": "string",
-                    "title": "Toggle Table String",
-                    "default": "_v",
-                    "info": "This string will be used to toggle visibility of tables in the schema browser when editing a query in order to remove non-useful tables from sight."
-                }
-            },
-            'required': ['url', 'username', 'password'],
-            'secret': ['password']
-        }
+    response_error = "JIRA returned unexpected status code"
+    requires_authentication = True
+    url_title = 'JIRA URL'
+    username_title = 'Username'
+    password_title = 'Password'
 
     @classmethod
     def name(cls):
@@ -199,13 +172,9 @@ def run_query(self, query, user):
             else:
                 query['maxResults'] = query.get('maxResults', 1000)
 
-            response = requests.get(jql_url, params=query, auth=(self.configuration.get('username'), self.configuration.get('password')))
-
-            if response.status_code == 401 or response.status_code == 403:
-                return None, "Authentication error. Please check username/password."
-
-            if response.status_code != 200:
-                return None, "JIRA returned unexpected status code ({})".format(response.status_code)
+            response, error = self.get_response(jql_url, params=query)
+            if error is not None:
+                return None, error
 
             data = response.json()
 
@@ -219,4 +188,3 @@ def run_query(self, query, user):
             return None, "Query cancelled by user."
 
 register(JiraJQL)
-

redash/query_runner/url.py

@@ -1,34 +1,10 @@
-import requests
-from redash.query_runner import BaseQueryRunner, register
+from redash.query_runner import BaseHTTPQueryRunner, register
 
 
-class Url(BaseQueryRunner):
+class Url(BaseHTTPQueryRunner):
     default_doc_url = ("http://redash.readthedocs.io/en/latest/"
                        "datasources.html#url")
 
-    @classmethod
-    def configuration_schema(cls):
-        return {
-            'type': 'object',
-            'properties': {
-                'url': {
-                    'type': 'string',
-                    'title': 'URL base path'
-                },
-                "doc_url": {
-                    "type": "string",
-                    "title": "Documentation URL",
-                    "default": cls.default_doc_url
-                },
-                "toggle_table_string": {
-                    "type": "string",
-                    "title": "Toggle Table String",
-                    "default": "_v",
-                    "info": "This string will be used to toggle visibility of tables in the schema browser when editing a query in order to remove non-useful tables from sight."
-                }
-            }
-        }
-
     @classmethod
     def annotate_query(cls):
         return False
@@ -40,7 +16,6 @@ def run_query(self, query, user):
         base_url = self.configuration.get("url", None)
 
         try:
-            error = None
             query = query.strip()
 
             if base_url is not None and base_url != "":
@@ -52,20 +27,17 @@ def run_query(self, query, user):
 
             url = base_url + query
 
-            response = requests.get(url)
-            response.raise_for_status()
-            json_data = response.content.strip()
+            response, error = self.get_response(url)
+            if error is not None:
+                return None, error
 
-            if not json_data:
-                error = "Got empty response from '{}'.".format(url)
+            json_data = response.content.strip()
 
-            return json_data, error
-        except requests.RequestException as e:
-            return None, str(e)
+            if json_data:
+                return json_data, None
+            else:
+                return None, "Got empty response from '{}'.".format(url)
         except KeyboardInterrupt:
-            error = "Query cancelled by user."
-            json_data = None
-
-        return json_data, error
+            return None, "Query cancelled by user."
 
 register(Url)