Athena: Switch to simple_json to serialize NaN/Infinity values as nulls.

Fixes #2544.

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:

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):