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

1 files changed, 11 insertions, 1 deletions

diff --git a/archaeological_warehouse/serializers.py b/archaeological_warehouse/serializers.py
index ff7c23778..b0cf4d96f 100644
--- a/archaeological_warehouse/serializers.py
+++ b/archaeological_warehouse/serializers.py
@@ -60,7 +60,8 @@ def warehouse_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,
-                            get_queryset=False, no_geo=True):
+                            get_queryset=False, no_geo=True,
+                            put_locks=False, lock_user=None):
     result_queryset = {}
     if operation_queryset:
         operation_ids = operation_queryset.values_list("id", flat=True)
@@ -100,6 +101,15 @@ def warehouse_serialization(archive=False, return_empty_types=False,
         return result_queryset
     result = generic_get_results(WAREHOUSE_MODEL_LIST, "warehouse",
                                  result_queryset=result_queryset, no_geo=no_geo)
+    if put_locks:
+        for model in WAREHOUSE_MODEL_LIST:
+            if not hasattr(model, "locked"):
+                continue
+            q = model.objects
+            if result_queryset and model.__name__ in result_queryset:
+                q = result_queryset[model.__name__]
+            q.update(locked=True, lock_user=lock_user)
+
     full_archive = archive_serialization(
         result, archive_dir="warehouse", archive=archive,
         return_empty_types=return_empty_types, archive_name=archive_name,