Coverage for backend/core/managers.py: 91%

1from django.db import models 

2 

3from authentication.user.models import User 

4from core.auxiliary.enums.FlowsheetTemplateType import FlowsheetTemplateType 

5from core.auxiliary.models.Flowsheet import Flowsheet 

6from core.validation import get_current_flowsheet, cache_result 

7from authentication.user.AccessTable import AccessTable 

8from typing import TypeVar 

9t = TypeVar('t', bound=models.Model, covariant=True) 

10 

11 

12def verify_flowsheet_access(user: User, flowsheet_id: int) -> bool: 

13 """ 

14 Verify if the user has access to the flowsheet either directly or through public template edit access. 

15 

16 :param user: User object 

17 :param flowsheet_id: Flowsheet ID to check access for 

18 

19 :return: True if the user has access, False otherwise 

20 """ 

21 user_has_direct_access = AccessTable.objects.filter( 

22 user_id=user.id, 

23 flowsheet_id=flowsheet_id 

24 ).exists() 

25 

26 user_can_edit_template = False 

27 

28 # Check if this is a public flowsheet template that the user can edit (only admins can edit public templates) 

29 if user.is_staff and not user_has_direct_access: 29 ↛ 30line 29 didn't jump to line 30 because the condition on line 29 was never true

30 flowsheet: Flowsheet = Flowsheet.objects.get(id=flowsheet_id) 

31 user_can_edit_template = flowsheet.flowsheet_template_type == FlowsheetTemplateType.PublicTemplate 

32 

33 access_granted = user_has_direct_access or user_can_edit_template 

34 

35 cache_result(has_access=access_granted) 

36 

37 return access_granted 

38 

39 

40class AccessControlManager(models.Manager): 

41 """ 

42 Custom object manager to enforce access control based on the current flowsheet and user context. 

43 This manager overrides the default queryset to filter objects based on the flowsheet. 

44 

45 By default, object modification methods (update, delete, etc.) will be scoped only to the 

46 active flowsheet that the user has access to. Creation methods check explicitly if the user has 

47 access to the flowsheet before allowing creation. 

48 """ 

49 

50 def __init__(self): 

51 super().__init__() 

52 self.flowsheet_related_name = "flowsheet" 

53 

54 def create(self, **kwargs): 

55 ctx = get_current_flowsheet() or {} 

56 

57 flowsheet = ctx.get("flowsheet") 

58 user = ctx.get("user") 

59 has_access = ctx.get("has_access") 

60 

61 # This is only the case when tests uses the objects directly without going through a view 

62 if flowsheet == None or user == None: 

63 return super().create(**kwargs) 

64 

65 if has_access is None: 65 ↛ 67line 65 didn't jump to line 67 because the condition on line 65 was never true

66 # Check access and cache the result 

67 has_access = verify_flowsheet_access(user=user, flowsheet_id=flowsheet) 

68 

69 kwargs[f"{self.flowsheet_related_name}_id"] = flowsheet 

70 

71 if has_access: 71 ↛ 74line 71 didn't jump to line 74 because the condition on line 71 was always true

72 return super().create(**kwargs) 

73 else: 

74 raise PermissionError("User does not have access to this flowsheet.") 

75 

76 def bulk_create(self, objs, **kwargs): 

77 ctx = get_current_flowsheet() or {} 

78 

79 flowsheet = ctx.get("flowsheet") 

80 user = ctx.get("user") 

81 has_access = ctx.get("has_access") 

82 

83 # This is only the case when tests uses the objects directly without going through a view 

84 if flowsheet == None or user == None: 

85 return super().bulk_create(objs, **kwargs) 

86 

87 if has_access is None: 

88 # Check access and cache the result 

89 has_access = verify_flowsheet_access(user=user, flowsheet_id=flowsheet) 

90 

91 # Set flowsheet ID on each object before bulk create 

92 for obj in objs: 

93 setattr(obj, f"{self.flowsheet_related_name}_id", flowsheet) 

94 

95 if has_access: 95 ↛ 98line 95 didn't jump to line 98 because the condition on line 95 was always true

96 return super().bulk_create(objs, **kwargs) 

97 else: 

98 raise PermissionError("User does not have access to this flowsheet.") 

99 

100 

101 def get_queryset(self): 

102 """ 

103 Override the get_queryset method to filter objects based on the current flowsheet and user access. 

104 """ 

105 qs = super().get_queryset() 

106 ctx = get_current_flowsheet() or {} 

107 

108 flowsheet = ctx.get("flowsheet") 

109 user = ctx.get("user") 

110 has_access = ctx.get("has_access") 

111 

112 # This is only the case when tests uses the objects directly without going through a view 

113 if flowsheet == None or user == None: 

114 return qs 

115 

116 # Use cached access result if present 

117 if has_access is None: 

118 # Check access and cache the result 

119 has_access = verify_flowsheet_access(user=user, flowsheet_id=flowsheet) 

120 

121 filter_kwargs = {f"{self.flowsheet_related_name}_id": flowsheet} 

122 

123 if has_access: 

124 return qs.filter(**filter_kwargs) 

125 else: 

126 return qs.none() 

127 

128class SoftDeleteManager(AccessControlManager): 

129 def __init__(self): 

130 super().__init__() 

131 def get_queryset(self): 

132 # filter out deleted objects 

133 return super().get_queryset().filter(is_deleted=False) 

134 

135 def include_deleted(self): 

136 # include deleted objects 

137 return super().get_queryset() 

138 

139 

140 

141def include_soft_deleted[t](objects: models.Manager[t]) -> models.QuerySet[t]: 

142 """ 

143 In some operations, including copying a flowsheet, 

144 we need to include the soft deleted objects 

145 Example usage: 

146 include_soft_deleted(SimulationObject.objects).filter(flowsheet_id=1) 

147 """ 

148 if objects is not None and isinstance(objects, SoftDeleteManager): 

149 return objects.include_deleted().all() 

150 else: 

151 return objects.all() 

152 

153 

154