diff --git a/redash/query_runner/__init__.py b/redash/query_runner/__init__.py index c839d92493..2851587f78 100644 --- a/redash/query_runner/__init__.py +++ b/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 = {} diff --git a/redash/query_runner/jql.py b/redash/query_runner/jql.py index 61d8fc6598..04c9c8fefc 100644 --- a/redash/query_runner/jql.py +++ b/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) - diff --git a/redash/query_runner/url.py b/redash/query_runner/url.py index c99289cca4..cfc1b03864 100644 --- a/redash/query_runner/url.py +++ b/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) diff --git a/tests/query_runner/test_http.py b/tests/query_runner/test_http.py new file mode 100644 index 0000000000..e4d88c24a8 --- /dev/null +++ b/tests/query_runner/test_http.py @@ -0,0 +1,136 @@ +import mock +from unittest import TestCase + +import requests +from redash.query_runner import BaseHTTPQueryRunner + + +class RequiresAuthQueryRunner(BaseHTTPQueryRunner): + requires_authentication = True + + +class TestBaseHTTPQueryRunner(TestCase): + + def test_requires_authentication_default(self): + self.assertFalse(BaseHTTPQueryRunner.requires_authentication) + schema = BaseHTTPQueryRunner.configuration_schema() + self.assertNotIn('username', schema['required']) + self.assertNotIn('password', schema['required']) + + def test_requires_authentication_true(self): + schema = RequiresAuthQueryRunner.configuration_schema() + self.assertIn('username', schema['required']) + self.assertIn('password', schema['required']) + + def test_get_auth_with_values(self): + query_runner = BaseHTTPQueryRunner({ + 'username': 'username', + 'password': 'password' + }) + self.assertEqual(query_runner.get_auth(), ('username', 'password')) + + def test_get_auth_empty(self): + query_runner = BaseHTTPQueryRunner({}) + self.assertIsNone(query_runner.get_auth()) + + def test_get_auth_empty_requires_authentication(self): + query_runner = RequiresAuthQueryRunner({}) + self.assertRaisesRegexp( + ValueError, + "Username and Password required", + query_runner.get_auth + ) + + @mock.patch('requests.get') + def test_get_response_success(self, mock_get): + mock_response = mock.Mock() + mock_response.status_code = 200 + mock_response.text = "Success" + mock_get.return_value = mock_response + + url = 'https://example.com/' + query_runner = BaseHTTPQueryRunner({}) + response, error = query_runner.get_response(url) + mock_get.assert_called_once_with(url, auth=None) + self.assertEqual(response.status_code, 200) + self.assertIsNone(error) + + @mock.patch('requests.get') + def test_get_response_success_custom_auth(self, mock_get): + mock_response = mock.Mock() + mock_response.status_code = 200 + mock_response.text = "Success" + mock_get.return_value = mock_response + + url = 'https://example.com/' + query_runner = BaseHTTPQueryRunner({}) + auth = ('username', 'password') + response, error = query_runner.get_response(url, auth=auth) + mock_get.assert_called_once_with(url, auth=auth) + self.assertEqual(response.status_code, 200) + self.assertIsNone(error) + + @mock.patch('requests.get') + def test_get_response_failure(self, mock_get): + mock_response = mock.Mock() + mock_response.status_code = 301 + mock_response.text = "Redirect" + mock_get.return_value = mock_response + + url = 'https://example.com/' + query_runner = BaseHTTPQueryRunner({}) + response, error = query_runner.get_response(url) + mock_get.assert_called_once_with(url, auth=None) + self.assertIn(query_runner.response_error, error) + + @mock.patch('requests.get') + def test_get_response_httperror_exception(self, mock_get): + mock_response = mock.Mock() + mock_response.status_code = 500 + mock_response.text = "Server Error" + http_error = requests.HTTPError() + mock_response.raise_for_status.side_effect = http_error + mock_get.return_value = mock_response + + url = 'https://example.com/' + query_runner = BaseHTTPQueryRunner({}) + response, error = query_runner.get_response(url) + mock_get.assert_called_once_with(url, auth=None) + self.assertIsNotNone(error) + self.assertIn("Failed to execute query", error) + + @mock.patch('requests.get') + def test_get_response_requests_exception(self, mock_get): + mock_response = mock.Mock() + mock_response.status_code = 500 + mock_response.text = "Server Error" + exception_message = "Some requests exception" + requests_exception = requests.RequestException(exception_message) + mock_response.raise_for_status.side_effect = requests_exception + mock_get.return_value = mock_response + + url = 'https://example.com/' + query_runner = BaseHTTPQueryRunner({}) + response, error = query_runner.get_response(url) + mock_get.assert_called_once_with(url, auth=None) + self.assertIsNotNone(error) + self.assertEqual(exception_message, error) + + @mock.patch('requests.get') + def test_get_response_generic_exception(self, mock_get): + mock_response = mock.Mock() + mock_response.status_code = 500 + mock_response.text = "Server Error" + exception_message = "Some generic exception" + exception = ValueError(exception_message) + mock_response.raise_for_status.side_effect = exception + mock_get.return_value = mock_response + + url = 'https://example.com/' + query_runner = BaseHTTPQueryRunner({}) + self.assertRaisesRegexp( + ValueError, + exception_message, + query_runner.get_response, + url + )