Bulk assign sequence numbers when saving Tasks and WorkOrders

Adds new optional keyword arguments to safely assign sequential reference
numbers to Task and WorkOrder models.

Connects #3049

opentreemap/treemap/instance.py

@@ -771,6 +771,54 @@ def raise_errors(errors):
         if errors:
             raise_errors(errors)
 
+    @transaction.atomic
+    def _bulk_assign_task_reference_numbers(self, tasks):
+        instance = Instance.objects.select_for_update().get(id=self.id)
+        for task in tasks:
+            task.reference_number = instance.task_sequence_number
+            instance.task_sequence_number += 1
+        instance.save(update_fields=['task_sequence_number'])
+
+    @transaction.atomic
+    def _bulk_assign_work_order_reference_numbers(self, work_orders):
+        instance = Instance.objects.select_for_update().get(id=self.id)
+        for work_order in work_orders:
+            work_order.reference_number = instance.work_order_sequence_number
+            instance.work_order_sequence_number += 1
+        instance.save(update_fields=['work_order_sequence_number'])
+
+    @transaction.atomic
+    def get_next_task_sequence(self, how_many=1):
+        """
+        Return next sequence value and increment counter by `how_many`.
+
+        Arguments:
+        how_many - Amount of sequence values to allocate
+        """
+        if how_many < 1:
+            raise ValueError('how_many must be >= 1')
+        qs = Instance.objects.filter(id=self.id).select_for_update()
+        result = qs[0].task_sequence_number
+        next_value = F('task_sequence_number') + how_many
+        qs.update(task_sequence_number=next_value)
+        return result
+
+    @transaction.atomic
+    def get_next_work_order_sequence(self, how_many=1):
+        """
+        Return next sequence value and increment counter by `how_many`.
+
+        Arguments:
+        how_many - Amount of sequence values to allocate
+        """
+        if how_many < 1:
+            raise ValueError('how_many must be >= 1')
+        qs = Instance.objects.filter(id=self.id).select_for_update()
+        result = qs[0].work_order_sequence_number
+        next_value = F('work_order_sequence_number') + how_many
+        qs.update(work_order_sequence_number=next_value)
+        return result
+
     def save(self, *args, **kwargs):
         self.full_clean()
 

opentreemap/works_management/tests.py

@@ -7,7 +7,7 @@
 from django.utils import timezone
 
 from treemap.audit import UserTrackingException
-from treemap.models import Plot, Instance
+from treemap.models import Plot
 from treemap.tests.base import OTMTestCase
 from treemap.tests import make_instance, make_commander_user
 
@@ -61,12 +61,12 @@ def setUp(self):
         self.instance = make_instance()
         self.user = make_commander_user(self.instance)
         self.team = make_team(self.instance)
+        self.plot = make_plot(self.instance, self.user)
 
     def test_work_order(self):
         w = make_work_order(self.instance, self.user)
 
-        plot = make_plot(self.instance, self.user)
-        t = make_task(self.instance, self.user, plot)
+        t = make_task(self.instance, self.user, self.plot)
         t.work_order = w
         t.save_with_user(self.user)
         self.assertEqual(1, w.tasks.count())
@@ -79,16 +79,15 @@ def test_work_order(self):
         w = WorkOrder.objects.get(id=w.id)
         self.assertTrue(w.updated_at > old_updated_at)
 
-    def test_cannot_call_save_on_workorder(self):
+    def test_cannot_call_save_on_work_order(self):
         # WorkOrder is Auditible and should only allow calling save_with_user.
         w = make_work_order(self.instance, self.user)
         with self.assertRaises(UserTrackingException):
             w.save()
 
     def test_cannot_call_save_on_task(self):
         # Task is Auditible and should only allow calling save_with_user.
-        t = make_task(self.instance, self.user,
-                      make_plot(self.instance, self.user))
+        t = make_task(self.instance, self.user, self.plot)
         with self.assertRaises(UserTrackingException):
             t.save()