Fix serializer

in-house-export-serialization
Kevin Chung 2020-04-27 01:49:56 -04:00
parent 593fed300b
commit 89ade49695
1 changed files with 4 additions and 4 deletions

View File

@ -1,6 +1,6 @@
import json import json
import six import six
from collections import OrderedDict from collections import defaultdict, OrderedDict
from CTFd.utils.exports.encoders import JSONEncoder from CTFd.utils.exports.encoders import JSONEncoder
@ -9,7 +9,7 @@ class JSONSerializer(object):
def __init__(self, query, fileobj): def __init__(self, query, fileobj):
self.query = query self.query = query
self.fileobj = fileobj self.fileobj = fileobj
self.buckets = [] self.buckets = defaultdict(list)
def serialize(self): def serialize(self):
for row in self.query: for row in self.query:
@ -17,7 +17,7 @@ class JSONSerializer(object):
self.close() self.close()
def write(self, path, result): def write(self, path, result):
self.buckets.append([result]) self.buckets[path].append(result)
def wrap(self, result): def wrap(self, result):
result = OrderedDict([("count", len(result)), ("results", result)]) result = OrderedDict([("count", len(result)), ("results", result)])
@ -25,7 +25,7 @@ class JSONSerializer(object):
return result return result
def close(self): def close(self):
for result in self.buckets: for path, result in self.buckets.items():
result = self.wrap(result) result = self.wrap(result)
# Certain databases (MariaDB) store JSON as LONGTEXT. # Certain databases (MariaDB) store JSON as LONGTEXT.