Some bytes/str issues in py3 w/ zlib and json

superset/sql_lab.py

@@ -6,7 +6,6 @@
 import pandas as pd
 import sqlalchemy
 import uuid
-import zlib
 
 from sqlalchemy.pool import NullPool
 from sqlalchemy.orm import sessionmaker
@@ -90,8 +89,8 @@ def handle_error(msg):
         executed_sql = superset_query.as_create_table(query.tmp_table_name)
         query.select_as_cta_used = True
     elif (
-            query.limit and superset_query.is_select() and
-            db_engine_spec.limit_method == LimitMethod.WRAP_SQL):
+                    query.limit and superset_query.is_select() and
+                    db_engine_spec.limit_method == LimitMethod.WRAP_SQL):
         executed_sql = database.wrap_sql_limit(executed_sql, query.limit)
         query.limit_used = True
     try:
@@ -170,7 +169,7 @@ def handle_error(msg):
     if store_results:
         key = '{}'.format(uuid.uuid4())
         logging.info("Storing results in results backend, key: {}".format(key))
-        results_backend.set(key, zlib.compress(payload))
+        results_backend.set(key, utils.zlib_compress(payload))
         query.results_key = key
 
     session.flush()

superset/utils.py

@@ -15,7 +15,9 @@
 import smtplib
 import sqlalchemy as sa
 import signal
+import sys
 import uuid
+import zlib
 
 from builtins import object
 from datetime import date, datetime, time
@@ -41,7 +43,7 @@
 
 logging.getLogger('MARKDOWN').setLevel(logging.INFO)
 
-
+PY3K = sys.version_info >= (3, 0)
 EPOCH = datetime(1970, 1, 1)
 DTTM_ALIAS = '__timestamp'
 
@@ -580,3 +582,44 @@ def setup_cache(app, cache_config):
     """Setup the flask-cache on a flask app"""
     if cache_config and cache_config.get('CACHE_TYPE') != 'null':
         return Cache(app, config=cache_config)
+
+
+def zlib_compress(data):
+    """
+    compress things in a py2/3 safe fashion
+
+    >>> json_str = '{"test": 1}'
+    >>> blob = zlib_compress(json_str)
+    """
+
+    if PY3K:
+        if isinstance(data, str):
+            return zlib.compress(bytes(data, "utf-8"))
+        else:
+            return zlib.compress(data)
+    else:
+        return zlib.compress(data)
+
+
+def zlib_uncompress_to_string(blob):
+    """
+    uncompress things to a string in a py2/3 safe fashion
+    >>> json_str = '{"test": 1}'
+    >>> blob = zlib_compress(json_str)
+    >>> got_str = zlib_uncompress_to_string(blob)
+    >>> got_str == json_str
+    True
+    """
+
+    if PY3K:
+        decompressed = ""
+        if isinstance(blob, bytes):
+            decompressed = zlib.decompress(blob)
+        else:
+            decompressed = zlib.decompress(bytes(blob, "utf-8"))
+
+        if isinstance(decompressed, str):
+            return decompressed
+        return decompressed.decode("utf-8")
+    else:
+        return zlib.decompress(blob)