Serialization: manage documenst (with filters) - refactoring

1 files changed, 69 insertions, 191 deletions

diff --git a/ishtar_common/serializers.py b/ishtar_common/serializers.py
index 034828835..c6cd118c7 100644
--- a/ishtar_common/serializers.py
+++ b/ishtar_common/serializers.py
@@ -1,8 +1,4 @@
-from collections import OrderedDict
-from copy import deepcopy
-import datetime
 import json
-import importlib
 import os
 import tempfile
 import zipfile
@@ -11,13 +7,19 @@ from zipfile import ZipFile
 
 from django.apps import apps
 from django.conf import settings
-from django.contrib.sites.models import Site
-from django.core.serializers import deserialize, serialize
+from django.core.serializers import deserialize
 from django.db.models import Q
 
-from version import get_version
 from . import models
-from archaeological_warehouse import models as warehouse_models
+
+from ishtar_common.serializers_utils import generic_get_results, \
+    archive_serialization, generic_archive_files, SERIALIZATION_VERSION, \
+    get_model_from_filename
+
+from archaeological_operations.serializers import operation_serialization
+from archaeological_context_records.serializers import cr_serialization
+from archaeological_finds.serializers import find_serialization
+from archaeological_warehouse.serializers import warehouse_serialization
 
 
 class PublicSerializer(serializers.BaseSerializer):
@@ -25,189 +27,6 @@ class PublicSerializer(serializers.BaseSerializer):
         return obj.public_representation()
 
 
-SERIALIZATION_VERSION = "1.0"
-
-
-def get_model_from_filename(filename):
-    filename = filename.split(".")[0]  # remove extension
-    module_name, model_name = filename.split("__")
-    module = importlib.import_module(module_name + ".models")
-    return getattr(module, model_name)
-
-
-def serialization_info():
-    site = Site.objects.get_current()
-    return {
-        "serialize-version": SERIALIZATION_VERSION,
-        "ishtar-version": get_version(),
-        "domain": site.domain,
-        "name": site.name,
-        "date": datetime.datetime.now().isoformat()
-    }
-
-
-def archive_serialization(result, archive_dir=None, archive=False,
-                          return_empty_types=False, archive_name=None):
-    """
-    Serialize all types models to JSON
-    Used for import and export scripts
-
-    :param result: serialization results
-    :param archive_dir: directory inside the archive (default None)
-    :param return_empty_types: if True instead of serialization return empty
-    types (default False)
-    :param archive: if True return a zip file containing all the file serialized
-    (default False)
-    :param archive_name: path to the archive if not provided a new archive is
-    created
-    :return: string containing the json serialization of types unless
-    return_empty_types or archive is set to True
-    """
-    if archive and return_empty_types:
-        raise ValueError("archive and return_empty_types are incompatible")
-    if return_empty_types:
-        return [k for k in result if not result[k]]
-    if not archive:
-        return result
-    archive_created = False
-    if not archive_name:
-        archive_created = True
-        tmpdir = tempfile.mkdtemp(prefix="ishtarexport-") + os.sep
-        archive_name = tmpdir + "ishtar-{}.zip".format(
-            datetime.date.today().strftime("%Y-%m-%d")
-        )
-    if not archive_name.endswith(".zip"):
-        archive_name += ".zip"
-    mode = "w" if archive_created else "a"
-    with tempfile.TemporaryDirectory() as tmpdirname:
-        if archive_dir:
-            os.mkdir(tmpdirname + os.sep + archive_dir)
-
-        with ZipFile(archive_name, mode) as current_zip:
-            if archive_created:
-                base_filename = "info.json"
-                filename = tmpdirname + os.sep + base_filename
-                with open(filename, "w") as json_file:
-                    json_file.write(
-                        json.dumps(serialization_info(), indent=2)
-                    )
-                current_zip.write(filename, arcname=base_filename)
-
-            for dir_name, model_name in result:
-                base_filename = model_name + ".json"
-                filename = tmpdirname + os.sep + base_filename
-                with open(filename, "w") as json_file:
-                    json_file.write(result[(dir_name, model_name)])
-                arcname = base_filename
-                if dir_name:
-                    arcname = dir_name + os.sep + base_filename
-                current_zip.write(filename, arcname=arcname)
-    return archive_name
-
-
-def generic_get_results(model_list, dirname, no_geo=True,
-                        result_queryset=None):
-    result = OrderedDict()
-    for model in model_list:
-        base_model_name = model.__name__
-        model_name = str(model.__module__).split(".")[0] + "__" + \
-                     base_model_name
-
-        if result_queryset and base_model_name in result_queryset:
-            base_q = result_queryset[base_model_name]
-        else:
-            base_q = model.objects
-        q = base_q
-        recursion = None
-        if hasattr(model, "parent"):
-            recursion = "parent"
-        elif hasattr(model, "inverse_relation"):
-            recursion = "inverse_relation"
-        elif hasattr(model, "children"):
-            recursion = "children__id"
-        if recursion:
-            q = q.filter(**{recursion + "__isnull": True})
-
-        key = (dirname, model_name)
-        result[key] = serialize(
-            "json", q.distinct().all(),
-            indent=2,
-            use_natural_foreign_keys=True, use_natural_primary_keys=True,
-        )
-
-        if recursion:
-            serialized = [item["id"] for item in q.values("id").all()]
-            recursion_in = recursion
-            if not recursion.endswith("_id"):
-                recursion_in += "_id"
-            recursion_in += "__in"
-            q = base_q.filter(**{recursion_in: serialized}
-                              ).exclude(id__in=serialized)
-            while q.count():
-                v = serialize(
-                    "json", q.all(), indent=2, use_natural_foreign_keys=True,
-                    use_natural_primary_keys=True)
-                new_result = json.loads(result[key])
-                new_result += json.loads(v)
-                result[key] = json.dumps(new_result, indent=2)
-                serialized += [item["id"] for item in q.values("id").all()]
-                q = base_q.filter(**{recursion_in: serialized}
-                                  ).exclude(id__in=serialized)
-            # managed circular
-            q = base_q.exclude(id__in=serialized)
-            if q.count():
-                v = serialize(
-                    "json", q.all(), indent=2, use_natural_foreign_keys=True,
-                    use_natural_primary_keys=True)
-                result_to_add = json.loads(v)
-                result_cleaned = deepcopy(result_to_add)
-                for res in result_cleaned:  # first add with no recursion
-                    res["fields"][recursion] = None
-                new_result = json.loads(result[key])
-                new_result += result_cleaned
-                new_result += result_to_add
-                result[key] = json.dumps(new_result, indent=2)
-
-        excluded_fields = ["history_modifier", "history_creator", "imports"]
-        if hasattr(model, "SERIALIZATION_EXCLUDE"):
-            excluded_fields = list(model.SERIALIZATION_EXCLUDE)
-        if no_geo:
-            excluded_fields += ["center", "limit"] + [
-                field.name for field in models.GeoItem._meta.get_fields()
-            ]
-        if excluded_fields:
-            new_result = json.loads(result[key])
-            for idx in range(len(new_result)):
-                for excluded_field in excluded_fields:
-                    if excluded_field in new_result[idx]["fields"]:
-                        new_result[idx]["fields"].pop(excluded_field)
-            result[key] = json.dumps(new_result, indent=2)
-    return result
-
-
-def generic_archive_files(model_list, archive_name=None):
-    result = []
-    for model in model_list:
-        if hasattr(model, "SERIALIZATION_FILES"):
-            for item in model.objects.all():
-                for attr in model.SERIALIZATION_FILES:
-                    media = getattr(item, attr)
-                    result.append((media.path, media.name))
-
-    archive_created = False
-    if not archive_name:
-        archive_created = True
-        tmpdir = tempfile.mkdtemp(prefix="ishtarexport-") + os.sep
-        archive_name = tmpdir + "media.zip"
-    if not archive_name.endswith(".zip"):
-        archive_name += ".zip"
-    mode = "w" if archive_created else "a"
-    with ZipFile(archive_name, mode) as current_zip:
-        for media_path, name in result:
-            current_zip.write(media_path, arcname=name)
-    return archive_name
-
-
 TYPE_MODEL_EXCLUDE = ["Area", "OperationTypeOld"]
 
 
@@ -294,6 +113,65 @@ def directory_serialization(archive=False, return_empty_types=False,
     return full_archive
 
 
+def document_serialization(archive=False, return_empty_types=False,
+                           archive_name=None, operation_queryset=None,
+                           site_queryset=None, cr_queryset=None,
+                           find_queryset=None, warehouse_queryset=None):
+    result_queryset = {}
+    get_queryset_attr = None
+    if operation_queryset:
+        get_queryset_attr = {"operation_queryset": operation_queryset,
+                             "get_queryset": True}
+    elif site_queryset:
+        get_queryset_attr = {"site_queryset": site_queryset,
+                             "get_queryset": True}
+    elif cr_queryset:
+        get_queryset_attr = {"cr_queryset": cr_queryset,
+                             "get_queryset": True}
+    elif find_queryset:
+        get_queryset_attr = {"find_queryset": find_queryset,
+                             "get_queryset": True}
+    elif warehouse_queryset:
+        get_queryset_attr = {"warehouse_queryset": warehouse_queryset,
+                             "get_queryset": True}
+
+    if get_queryset_attr:
+        queries = operation_serialization(**get_queryset_attr)
+        queries.update(cr_serialization(**get_queryset_attr))
+        queries.update(find_serialization(**get_queryset_attr))
+        queries.update(warehouse_serialization(**get_queryset_attr))
+        q = None
+        for model, attr in (
+                ("Operation", "operations"),
+                ("ArchaeologicalSite", "sites"),
+                ("ContextRecord", "context_records"),
+                ("Find", "finds"),
+                ("Warehouse", "warehouses"),
+                ("Container", "containers")):
+            values = queries[model].values_list("id", flat=True)
+            base_q = Q(**{attr + "__id__in": values})
+            if not q:
+                q = base_q
+            else:
+                q |= base_q
+        result_queryset["Document"] = models.Document.objects.filter(q)
+
+    result = generic_get_results([models.Document], "documents",
+                                 result_queryset=result_queryset)
+    media_archive = None
+    if archive:
+        media_archive = generic_archive_files([models.Document],
+                                              result_queryset=result_queryset)
+    full_archive = archive_serialization(
+        result, archive_dir="documents", archive=archive,
+        return_empty_types=return_empty_types, archive_name=archive_name)
+    if not media_archive:
+        return full_archive
+    with ZipFile(full_archive, 'a') as current_zip:
+        current_zip.write(media_archive, arcname="media.zip")
+    return full_archive
+
+
 def restore_serialized(archive_name, delete_existing=False):
     with zipfile.ZipFile(archive_name, "r") as zip_file:
         # check version