Coverage for backend/django/core/auxiliary/models/ObjectTypeCounter.py: 96%

1# core/auxiliary/models/ObjectTypeCounter.py 

2from django.db import models 

3from django.db import transaction 

4from django.db.models import F 

5 

6from core.auxiliary.models.FlowsheetState import FlowsheetState 

7from core.managers import AccessControlManager, AllFlowsheetStatesManager 

8 

9class ObjectTypeCounter(models.Model): 

10 """ 

11 Keeps a per-type counter for objects inside a flowsheet. 

12 Example: Heater1, Heater2, ST1, ST2, etc. 

13 """ 

14 objects = AccessControlManager() 

15 all_states = AllFlowsheetStatesManager() 

16 flowsheet_state = models.ForeignKey( 

17 FlowsheetState, on_delete=models.CASCADE, related_name="type_counters" 

18 ) 

19 object_type = models.CharField(max_length=64) # e.g., "heater", "compressor", "stream" 

20 next_index = models.IntegerField(default=1) 

21 

22 class Meta: 

23 unique_together = ("flowsheet_state", "object_type") 

24 

25 def __str__(self): 

26 return f"{self.flowsheet_state.flowsheet.name} - {self.object_type}: {self.next_index}" 

27 

28 @classmethod 

29 def next_for(cls, flowsheet_state: FlowsheetState, object_type: str) -> int: 

30 """ 

31 Atomically reserve and return the next index for (flowsheet, object_type). 

32 Safe even when multiple users create objects at the same time. 

33 """ 

34 with transaction.atomic(): 

35 counter, _ = cls.objects.select_for_update().get_or_create( 

36 flowsheet_state=flowsheet_state, 

37 object_type=object_type, 

38 defaults={"next_index": 1}, 

39 ) 

40 current = counter.next_index 

41 counter.next_index = F("next_index") + 1 

42 counter.save(update_fields=["next_index"]) 

43 counter.refresh_from_db(fields=["next_index"]) 

44 return current