Code cleanup on batch operations.

2 files changed, 21 insertions, 56 deletions

diff --git a/app/utils/batch/batch_op.py b/app/utils/batch/batch_op.py
index ef99138..d0f0a31 100644
--- a/app/utils/batch/batch_op.py
+++ b/app/utils/batch/batch_op.py
@@ -13,13 +13,6 @@
 
 """Batch operation classes."""
 
-try:
-    import simplejson as json
-except ImportError:
-    import json
-
-import bson
-
 import handlers.common as hcommon
 import handlers.count as hcount
 import models
@@ -119,11 +112,10 @@ class BatchOperation(object):
             self.operation = utils.db.find_one
             self.args = [
                 self._database[self._collection],
-                self.document_id,
+                self.document_id
             ]
             self.kwargs = {
-                'fields': hcommon.get_query_fields(self.query_args_func)
-            }
+                "fields": hcommon.get_query_fields(self.query_args_func)}
         else:
             # Get the spec and perform the query, can perform an aggregation
             # as well.
@@ -140,10 +132,10 @@ class BatchOperation(object):
                     unique
                 ]
                 self.kwargs = {
-                    'sort': sort,
-                    'fields': fields,
-                    'match': spec,
-                    'limit': self._limit
+                    "sort": sort,
+                    "fields": fields,
+                    "match": spec,
+                    "limit": self._limit
                 }
             else:
                 self.operation = utils.db.find_and_count
@@ -153,9 +145,9 @@ class BatchOperation(object):
                     self._skip,
                 ]
                 self.kwargs = {
-                    'spec': spec,
-                    'fields': fields,
-                    'sort': sort
+                    "spec": spec,
+                    "fields": fields,
+                    "sort": sort
                 }
 
     def _prepare_post_operation(self):
@@ -186,31 +178,12 @@ class BatchOperation(object):
 
             json_obj = {
                 models.RESULT_KEY: res,
-                models.COUNT_KEY: count
-            }
+                models.COUNT_KEY: count}
 
             if self._limit is not None:
                 json_obj[models.LIMIT_KEY] = self._limit
 
-            try:
-                # Doing like this otherwise the result returned will be a
-                # string and not a real JSON object (since we have to serialize
-                # it here).
-                response[models.RESULT_KEY] = [
-                    json.loads(
-                        json.dumps(
-                            json_obj,
-                            default=bson.json_util.default,
-                            separators=(",", ":")
-                        )
-                    )
-                ]
-            except TypeError:
-                response[models.RESULT_KEY] = []
-                utils.LOG.error(
-                    "Error serializing JSON object. Objects to serialize: "
-                    "%s, %s", type(res), type(count)
-                )
+            response[models.RESULT_KEY] = [json_obj]
         else:
             response[models.RESULT_KEY] = result
 
@@ -219,8 +192,8 @@ class BatchOperation(object):
     def run(self):
         """Prepare and run this operation.
 
-        This method will invoke the `operation` attribute as a function, passing
-        the defined `args` and `kwargs` parameters.
+        This method will invoke the `operation` attribute as a function,
+        passing the defined `args` and `kwargs` parameters.
         """
         self.prepare_operation()
 
@@ -235,8 +208,7 @@ class BatchBootOperation(BatchOperation):
 
     def __init__(self, collection, database, operation_id=None):
         super(BatchBootOperation, self).__init__(
-            collection, database, operation_id
-        )
+            collection, database, operation_id)
         self.valid_keys = hcommon.BOOT_VALID_KEYS
 
 
@@ -245,8 +217,7 @@ class BatchJobOperation(BatchOperation):
 
     def __init__(self, collection, database, operation_id=None):
         super(BatchJobOperation, self).__init__(
-            collection, database, operation_id
-        )
+            collection, database, operation_id)
         self.valid_keys = hcommon.JOB_VALID_KEYS
 
 
@@ -255,8 +226,7 @@ class BatchDefconfigOperation(BatchOperation):
 
     def __init__(self, collection, database, operation_id=None):
         super(BatchDefconfigOperation, self).__init__(
-            collection, database, operation_id
-        )
+            collection, database, operation_id)
         self.valid_keys = hcommon.DEFCONFIG_VALID_KEYS
 
 
@@ -265,8 +235,7 @@ class BatchCountOperation(BatchOperation):
 
     def __init__(self, collection, database, operation_id=None):
         super(BatchCountOperation, self).__init__(
-            collection, database, operation_id
-        )
+            collection, database, operation_id)
         self.valid_keys = hcommon.COUNT_VALID_KEYS
 
     def _prepare_get_operation(self):
diff --git a/app/utils/batch/common.py b/app/utils/batch/common.py
index 64c6ebd..cf59f1b 100644
--- a/app/utils/batch/common.py
+++ b/app/utils/batch/common.py
@@ -43,16 +43,13 @@ def create_batch_operation(json_obj, db_options):
 
             if collection == models.COUNT_COLLECTION:
                 batch_op = batchop.BatchCountOperation(
-                    collection, database, operation_id=operation_id
-                )
+                    collection, database, operation_id=operation_id)
             elif collection == models.BOOT_COLLECTION:
                 batch_op = batchop.BatchBootOperation(
-                    collection, database, operation_id=operation_id
-                )
+                    collection, database, operation_id=operation_id)
             elif collection == models.JOB_COLLECTION:
                 batch_op = batchop.BatchJobOperation(
-                    collection, database, operation_id=operation_id
-                )
+                    collection, database, operation_id=operation_id)
             elif collection == models.DEFCONFIG_COLLECTION:
                 batch_op = batchop.BatchDefconfigOperation(
                     collection, database, operation_id=operation_id)
@@ -61,8 +58,7 @@ def create_batch_operation(json_obj, db_options):
                     collection, database, operation_id=operation_id)
 
             batch_op.query_args = get_batch_query_args(
-                get_func(models.QUERY_KEY, None)
-            )
+                get_func(models.QUERY_KEY, None))
             batch_op.document_id = get_func(models.DOCUMENT_ID_KEY, None)
             batch_op.query_args_func = batch_op.query_args.get
             batch_op.method = get_func(models.METHOD_KEY, None)