Merge pull request #1349 from CTFd/in-house-export-serialization

* Remove `datafreeze` dependency
* Closes #1348

CTFd/utils/exports/__init__.py

@@ -5,12 +5,9 @@ import re
 import tempfile
 import zipfile
 
-import datafreeze
 import dataset
 import six
 from alembic.util import CommandError
-from datafreeze.format import SERIALIZERS
-from datafreeze.format.fjson import JSONEncoder, JSONSerializer
 from flask import current_app as app
 from flask_migrate import upgrade
 from sqlalchemy.exc import OperationalError, ProgrammingError
@@ -20,6 +17,7 @@ from CTFd import __version__ as CTFD_VERSION
 from CTFd.cache import cache
 from CTFd.models import db, get_class_by_tablename
 from CTFd.utils import get_app_config, set_config
+from CTFd.utils.exports.freeze import freeze_export
 from CTFd.utils.migrations import (
     create_database,
     drop_database,
@@ -29,52 +27,6 @@ from CTFd.utils.migrations import (
 from CTFd.utils.uploads import get_uploader
 
 
-class CTFdSerializer(JSONSerializer):
-    """
-    Slightly modified datafreeze serializer so that we can properly
-    export the CTFd database into a zip file.
-    """
-
-    def close(self):
-        for path, result in self.buckets.items():
-            result = self.wrap(result)
-
-            if self.fileobj is None:
-                fh = open(path, "wb")
-            else:
-                fh = self.fileobj
-
-            # Certain databases (MariaDB) store JSON as LONGTEXT.
-            # Before emitting a file we should standardize to valid JSON (i.e. a dict)
-            # See Issue #973
-            for i, r in enumerate(result["results"]):
-                data = r.get("requirements")
-                if data:
-                    try:
-                        if isinstance(data, six.string_types):
-                            result["results"][i]["requirements"] = json.loads(data)
-                    except ValueError:
-                        pass
-
-            data = json.dumps(
-                result, cls=JSONEncoder, indent=self.export.get_int("indent")
-            )
-
-            callback = self.export.get("callback")
-            if callback:
-                data = "%s && %s(%s);" % (callback, callback, data)
-
-            if six.PY3:
-                fh.write(bytes(data, encoding="utf-8"))
-            else:
-                fh.write(data)
-            if self.fileobj is None:
-                fh.close()
-
-
-SERIALIZERS["ctfd"] = CTFdSerializer  # Load the custom serializer
-
-
 def export_ctf():
     # TODO: For some unknown reason dataset is only able to see alembic_version during tests.
     # Even using a real sqlite database. This makes this test impossible to pass in sqlite.
@@ -89,7 +41,7 @@ def export_ctf():
     for table in tables:
         result = db[table].all()
         result_file = six.BytesIO()
-        datafreeze.freeze(result, format="ctfd", fileobj=result_file)
+        freeze_export(result, fileobj=result_file)
         result_file.seek(0)
         backup_zip.writestr("db/{}.json".format(table), result_file.read())
 

CTFd/utils/exports/freeze.py

@@ -0,0 +1,11 @@
+from CTFd.utils.exports.serializers import JSONSerializer
+from sqlalchemy.exc import ProgrammingError, OperationalError
+
+
+def freeze_export(result, fileobj):
+    try:
+        query = result
+        serializer = JSONSerializer(query, fileobj)
+        serializer.serialize()
+    except (OperationalError, ProgrammingError) as e:
+        raise OperationalError("Invalid query: %s" % e)

CTFd/utils/exports/serializers.py

@@ -0,0 +1,53 @@
+import json
+import six
+from collections import defaultdict, OrderedDict
+from datetime import datetime, date
+from decimal import Decimal
+
+
+class JSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, (datetime, date)):
+            return obj.isoformat()
+        if isinstance(obj, Decimal):
+            return str(obj)
+
+
+class JSONSerializer(object):
+    def __init__(self, query, fileobj):
+        self.query = query
+        self.fileobj = fileobj
+        self.buckets = defaultdict(list)
+
+    def serialize(self):
+        for row in self.query:
+            self.write(None, row)
+        self.close()
+
+    def write(self, path, result):
+        self.buckets[path].append(result)
+
+    def wrap(self, result):
+        result = OrderedDict(
+            [("count", len(result)), ("results", result), ("meta", {})]
+        )
+        return result
+
+    def close(self):
+        for path, result in self.buckets.items():
+            result = self.wrap(result)
+
+            # Certain databases (MariaDB) store JSON as LONGTEXT.
+            # Before emitting a file we should standardize to valid JSON (i.e. a dict)
+            # See Issue #973
+            for i, r in enumerate(result["results"]):
+                data = r.get("requirements")
+                if data:
+                    try:
+                        if isinstance(data, six.string_types):
+                            result["results"][i]["requirements"] = json.loads(data)
+                    except ValueError:
+                        pass
+
+            data = json.dumps(result, cls=JSONEncoder, separators=(",", ":"))
+            self.fileobj.write(data.encode("utf-8"))

requirements.txt

@@ -13,13 +13,10 @@ itsdangerous==1.1.0
 requests>=2.20.0
 PyMySQL==0.9.3
 gunicorn==19.9.0
-banal==0.4.2
-normality==2.0.0
 dataset==1.1.2
 mistune==0.8.4
 netaddr==0.7.19
 redis==3.3.11
-datafreeze==0.1.0
 gevent==1.4.0
 python-dotenv==0.10.3
 flask-restx==0.1.1