Serialization - Import/Export: manage put an release of locks

1 files changed, 44 insertions, 18 deletions

diff --git a/ishtar_common/serializers.py b/ishtar_common/serializers.py
index 4b06a9be3..0914a297a 100644
--- a/ishtar_common/serializers.py
+++ b/ishtar_common/serializers.py
@@ -122,7 +122,8 @@ def directory_serialization(archive=False, return_empty_types=False,
 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):
+                           find_queryset=None, warehouse_queryset=None,
+                           put_locks=False, lock_user=None):
     result_queryset = {}
     get_queryset_attr = None
     if operation_queryset:
@@ -164,6 +165,12 @@ def document_serialization(archive=False, return_empty_types=False,
 
     result = generic_get_results([models.Document], "documents",
                                  result_queryset=result_queryset)
+    if put_locks:
+        q = models.Document.objects
+        if result_queryset:
+            q = result_queryset["Document"]
+        q.update(locked=True, lock_user=lock_user)
+
     media_archive = None
     if archive:
         media_archive = generic_archive_files([models.Document],
@@ -215,7 +222,8 @@ def full_serialization(operation_queryset=None, site_queryset=None,
                        warehouse_queryset=None, archive=True, no_geo=True,
                        info=None, export_types=True, export_conf=True,
                        export_importers=True, export_geo=True, export_dir=True,
-                       export_docs=True, export_items=True):
+                       export_docs=True, export_items=True, put_locks=False,
+                       lock_user=None):
     archive_name = None
     if export_types:
         # print("type")
@@ -242,7 +250,9 @@ def full_serialization(operation_queryset=None, site_queryset=None,
             archive=archive, archive_name=archive_name,
             operation_queryset=operation_queryset, site_queryset=site_queryset,
             cr_queryset=cr_queryset, find_queryset=find_queryset,
-            warehouse_queryset=warehouse_queryset)
+            warehouse_queryset=warehouse_queryset,
+            put_locks=put_locks, lock_user=lock_user
+        )
     if export_items:
         # print("operation")
         archive_name = operation_serialization(
@@ -250,32 +260,33 @@ def full_serialization(operation_queryset=None, site_queryset=None,
             archive_name=archive_name, operation_queryset=operation_queryset,
             site_queryset=site_queryset, cr_queryset=cr_queryset,
             find_queryset=find_queryset, warehouse_queryset=warehouse_queryset,
-            no_geo=no_geo)
+            no_geo=no_geo, put_locks=put_locks, lock_user=lock_user)
         # print("cr")
         cr_serialization(
             archive=archive,
             archive_name=archive_name, operation_queryset=operation_queryset,
             site_queryset=site_queryset, cr_queryset=cr_queryset,
             find_queryset=find_queryset, warehouse_queryset=warehouse_queryset,
-            no_geo=no_geo)
+            no_geo=no_geo, put_locks=put_locks, lock_user=lock_user)
         # print("find")
         find_serialization(
             archive=archive,
             archive_name=archive_name, operation_queryset=operation_queryset,
             site_queryset=site_queryset, cr_queryset=cr_queryset,
             find_queryset=find_queryset, warehouse_queryset=warehouse_queryset,
-            no_geo=no_geo)
+            no_geo=no_geo, put_locks=put_locks, lock_user=lock_user)
         # print("warehouse")
         warehouse_serialization(
             archive=archive,
             archive_name=archive_name, operation_queryset=operation_queryset,
             site_queryset=site_queryset, cr_queryset=cr_queryset,
             find_queryset=find_queryset, warehouse_queryset=warehouse_queryset,
-            no_geo=no_geo)
+            no_geo=no_geo, put_locks=put_locks, lock_user=lock_user)
     return archive_name
 
 
-def restore_serialized(archive_name, user=None, delete_existing=False):
+def restore_serialized(archive_name, user=None, delete_existing=False,
+                       release_locks=False):
     with zipfile.ZipFile(archive_name, "r") as zip_file:
         # check version
         info = json.loads(zip_file.read("info.json").decode("utf-8"))
@@ -323,13 +334,15 @@ def restore_serialized(archive_name, user=None, delete_existing=False):
                     # regenerate labels, add a new version, etc.
                     historized = hasattr(model, "history_modifier") and (
                         hasattr(model, "history_creator"))
+                    releasing_locks = hasattr(model, "locked") and (
+                        release_locks)
                     need_resave = hasattr(model, "CACHED_LABELS") or \
                                   hasattr(model, "cached_label") or \
-                                  (user and historized)
+                                  releasing_locks or (user and historized)
                     idx = -1
                     for idx, obj in enumerate(deserialize("json", data)):
                         extra_attrs = {}
-                        if historized:
+                        if historized or hasattr(model, "locked"):
                             keys = obj.object.natural_key()
                             old_obj = None
                             try:
@@ -337,14 +350,20 @@ def restore_serialized(archive_name, user=None, delete_existing=False):
                                     *keys)
                             except model.DoesNotExist:
                                 pass
-                            if old_obj and (old_obj.history_creator or
-                                            old_obj.history_modifier):
-                                extra_attrs = {
-                                    "history_modifier_id":
-                                        old_obj.history_modifier_id,
-                                    "history_creator_id":
-                                        old_obj.history_creator_id
-                                }
+                            if old_obj:
+                                if historized and (old_obj.history_creator or
+                                                   old_obj.history_modifier):
+                                    extra_attrs = {
+                                        "history_modifier_id":
+                                            old_obj.history_modifier_id,
+                                        "history_creator_id":
+                                            old_obj.history_creator_id
+                                    }
+                                if hasattr(model, "locked") and old_obj.locked:
+                                    extra_attrs.update({
+                                        "locked": old_obj.locked,
+                                        "lock_user": old_obj.lock_user,
+                                    })
                         obj.save()
                         if need_resave or extra_attrs:
                             obj = model.objects.get(id=obj.object.id)
@@ -356,10 +375,17 @@ def restore_serialized(archive_name, user=None, delete_existing=False):
                                         "history_creator_id"]
                                 else:
                                     obj.history_creator = user
+                                if extra_attrs and \
+                                        "locked" in extra_attrs:
+                                    obj.locked = extra_attrs["locked"]
+                                    obj.lock_user = extra_attrs["lock_user"]
                             elif extra_attrs:
                                 for k in extra_attrs:
                                     setattr(obj, k, extra_attrs[k])
                                 obj.skip_history_when_saving = True
+                            if releasing_locks:
+                                obj.locked = False
+                                obj.lock_user = None
                             obj._no_move = True
                             obj.save()
                     if idx >= 0: