diff --git a/redash/query_runner/athena.py b/redash/query_runner/athena.py index 5d3415d5a..a07952934 100644 --- a/redash/query_runner/athena.py +++ b/redash/query_runner/athena.py @@ -1,10 +1,11 @@ import json import logging import os +import simplejson from redash.query_runner import * from redash.settings import parse_boolean -from redash.utils import JSONEncoder +from redash.utils import SimpleJSONEncoder logger = logging.getLogger(__name__) ANNOTATE_QUERY = parse_boolean(os.environ.get('ATHENA_ANNOTATE_QUERY', 'true')) @@ -194,7 +195,7 @@ class Athena(BaseQueryRunner): 'athena_query_id': athena_query_id } } - json_data = json.dumps(data, cls=JSONEncoder) + json_data = simplejson.dumps(data, ignore_nan=True, cls=SimpleJSONEncoder) error = None except KeyboardInterrupt: if cursor.query_id: diff --git a/redash/utils/__init__.py b/redash/utils/__init__.py index 1bee42850..727fe07a2 100644 --- a/redash/utils/__init__.py +++ b/redash/utils/__init__.py @@ -10,6 +10,7 @@ import hashlib import pytz import pystache import os +import simplejson from funcy import distinct, select_values from sqlalchemy.orm.query import Query @@ -66,23 +67,43 @@ def generate_token(length): return ''.join(rand.choice(chars) for x in range(length)) -class JSONEncoder(json.JSONEncoder): +class JSONEncoderMixin: """Custom JSON encoding class, to handle Decimal and datetime.date instances.""" - def default(self, o): + def process_default(self, o): # Some SQLAlchemy collections are lazy. if isinstance(o, Query): - return list(o) + return True, list(o) if isinstance(o, decimal.Decimal): - return float(o) + return True, float(o) if isinstance(o, (datetime.date, datetime.time)): - return o.isoformat() + return True, o.isoformat() if isinstance(o, datetime.timedelta): - return str(o) + return True, str(o) - super(JSONEncoder, self).default(o) + return False, None # default processing + + +class JSONEncoder(JSONEncoderMixin, json.JSONEncoder): + """Adapter for `json.dumps`.""" + + def default(self, o): + processed, result = self.process_default(o) + if not processed: + result = super(JSONEncoder, self).default(o) + return result + + +class SimpleJSONEncoder(JSONEncoderMixin, simplejson.JSONEncoder): + """Adapter for `simplejson.dumps`.""" + + def default(self, o): + processed, result = self.process_default(o) + if not processed: + super(SimpleJSONEncoder, self).default(o) + return result def json_dumps(data):