diff --git a/redash/handlers/query_results.py b/redash/handlers/query_results.py index 815a7a154d..28f8485b8b 100644 --- a/redash/handlers/query_results.py +++ b/redash/handlers/query_results.py @@ -12,7 +12,7 @@ from redash.tasks.queries import enqueue_query from redash.utils import (collect_parameters_from_request, gen_query_hash, json_dumps, utcnow, to_filename) from redash.models.parameterized_query import ParameterizedQuery, InvalidParameterError, dropdown_values -from redash.serializers import serialize_query_result_to_csv, serialize_query_result_to_xlsx +from redash.serializers import serialize_query_result, serialize_query_result_to_csv, serialize_query_result_to_xlsx def error_response(message): @@ -42,7 +42,7 @@ def run_query(query, parameters, data_source, query_id, max_age=0): query_result = models.QueryResult.get_latest(data_source, query.text, max_age) if query_result: - return {'query_result': query_result.to_dict()} + return {'query_result': serialize_query_result(query_result, current_user.is_api_user())} else: job = enqueue_query(query.text, data_source, current_user.id, current_user.is_api_user(), metadata={ "Username": repr(current_user) if current_user.is_api_user() else current_user.email, diff --git a/redash/serializers/__init__.py b/redash/serializers/__init__.py index ba51ad0eb1..9223aaaac4 100644 --- a/redash/serializers/__init__.py +++ b/redash/serializers/__init__.py @@ -12,7 +12,7 @@ from redash.utils import json_loads from redash.models.parameterized_query import ParameterizedQuery -from .query_result import serialize_query_result_to_csv, serialize_query_result_to_xlsx +from .query_result import serialize_query_result, serialize_query_result_to_csv, serialize_query_result_to_xlsx def public_widget(widget): diff --git a/redash/serializers/query_result.py b/redash/serializers/query_result.py index 3a6dd7e123..32eef9789f 100644 --- a/redash/serializers/query_result.py +++ b/redash/serializers/query_result.py @@ -1,7 +1,7 @@ import cStringIO import csv import xlsxwriter -from funcy import rpartial +from funcy import rpartial, project from dateutil.parser import isoparse as parse_date from redash.utils import json_loads, UnicodeWriter from redash.query_runner import (TYPE_BOOLEAN, TYPE_DATE, TYPE_DATETIME) @@ -57,6 +57,14 @@ def _get_column_lists(columns): return fieldnames, special_columns +def serialize_query_result(query_result, is_api_user): + if is_api_user: + publicly_needed_keys = ['data', 'retrieved_at'] + return project(query_result.to_dict(), publicly_needed_keys) + else: + return query_result.to_dict() + + def serialize_query_result_to_csv(query_result): s = cStringIO.StringIO() diff --git a/tests/serializers/test_query_results.py b/tests/serializers/test_query_results.py index ff77d20646..3f9e7d0381 100644 --- a/tests/serializers/test_query_results.py +++ b/tests/serializers/test_query_results.py @@ -6,7 +6,7 @@ from redash import models from redash.utils import utcnow, json_dumps -from redash.serializers import serialize_query_result_to_csv +from redash.serializers import serialize_query_result, serialize_query_result_to_csv data = { @@ -24,6 +24,19 @@ ] } +class QueryResultSerializationTest(BaseTestCase): + def test_serializes_all_keys_for_authenticated_users(self): + query_result = self.factory.create_query_result(data=json_dumps({})) + serialized = serialize_query_result(query_result, False) + self.assertSetEqual(set(query_result.to_dict().keys()), + set(serialized.keys())) + + def test_doesnt_serialize_sensitive_keys_for_unauthenticated_users(self): + query_result = self.factory.create_query_result(data=json_dumps({})) + serialized = serialize_query_result(query_result, True) + self.assertSetEqual(set(['data', 'retrieved_at']), + set(serialized.keys())) + class CsvSerializationTest(BaseTestCase): def get_csv_content(self): query_result = self.factory.create_query_result(data=json_dumps(data))