Coverage for extras/models/scripts.py: 48%

1import inspect 

2import logging 

3from functools import cached_property 

4 

5from django.contrib.contenttypes.fields import GenericRelation 

6from django.contrib.contenttypes.models import ContentType 

7from django.db import models, router, transaction 

8from django.db.models import Q 

9from django.urls import reverse 

10from django.utils.translation import gettext_lazy as _ 

11 

12from core.choices import ManagedFileRootPathChoices 

13from core.models import ManagedFile 

14from extras.utils import is_script 

15from netbox.models.features import EventRulesMixin, JobsMixin, batch_delete_jobs 

16from utilities.querysets import RestrictedQuerySet 

17 

18from .mixins import PythonModuleMixin 

19 

20__all__ = ( 

21 'Script', 

22 'ScriptModule', 

23) 

24 

25logger = logging.getLogger('netbox.data_backends') 

26 

27 

28class Script(EventRulesMixin, JobsMixin): 

29 name = models.CharField( 

30 verbose_name=_('name'), 

31 max_length=79, # Maximum length for a Python class name 

32 editable=False, 

33 ) 

34 module = models.ForeignKey( 

35 to='extras.ScriptModule', 

36 on_delete=models.CASCADE, 

37 related_name='scripts', 

38 editable=False 

39 ) 

40 is_executable = models.BooleanField( 

41 default=True, 

42 verbose_name=_('is executable'), 

43 editable=False 

44 ) 

45 events = GenericRelation( 

46 'extras.EventRule', 

47 content_type_field='action_object_type', 

48 object_id_field='action_object_id' 

49 ) 

50 

51 def __str__(self): 

52 return self.name 

53 

54 objects = RestrictedQuerySet.as_manager() 

55 

56 class Meta: 

57 ordering = ('module', 'name') 

58 constraints = ( 

59 models.UniqueConstraint( 

60 fields=('name', 'module'), 

61 name='extras_script_unique_name_module' 

62 ), 

63 ) 

64 indexes = ( 

65 models.Index(fields=('module', 'name')), # Default ordering 

66 ) 

67 verbose_name = _('script') 

68 verbose_name_plural = _('scripts') 

69 

70 def get_absolute_url(self): 

71 return reverse('extras:script', args=[self.pk]) 

72 

73 @property 

74 def result(self): 

75 return self.jobs.all().order_by('-created').first() 

76 

77 @cached_property 

78 def python_class(self): 

79 return self.module.module_scripts.get(self.name) 

80 

81 def delete(self, soft_delete=False, **kwargs): 

82 if soft_delete and self.jobs.exists(): 

83 self.is_executable = False 

84 self.save() 

85 else: 

86 super().delete(**kwargs) 

87 self.id = None 

88 

89 

90class ScriptModuleManager(models.Manager.from_queryset(RestrictedQuerySet)): 

91 

92 def get_queryset(self): 

93 return super().get_queryset().filter( 

94 Q(file_root=ManagedFileRootPathChoices.SCRIPTS) | Q(file_root=ManagedFileRootPathChoices.REPORTS)) 

95 

96 

97class ScriptModule(PythonModuleMixin, JobsMixin, ManagedFile): 

98 """ 

99 Proxy model for script module files. 

100 """ 

101 objects = ScriptModuleManager() 

102 error = None 

103 

104 event_rules = GenericRelation( 

105 to='extras.EventRule', 

106 content_type_field='action_object_type', 

107 object_id_field='action_object_id', 

108 for_concrete_model=False 

109 ) 

110 

111 class Meta: 

112 proxy = True 

113 ordering = ('file_root', 'file_path') 

114 verbose_name = _('script module') 

115 verbose_name_plural = _('script modules') 

116 

117 def get_absolute_url(self): 

118 return reverse('extras:script_list') 

119 

120 def __str__(self): 

121 return self.python_name 

122 

123 def delete(self, using=None, *args, **kwargs): 

124 # Job is imported here rather than at module level to avoid a circular import 

125 # (core.models.jobs -> core.signals -> extras.events -> extras.models -> this module). 

126 from core.models import Job 

127 

128 # Deleting a ScriptModule cascades (via the Script.module FK) to its child Scripts, and 

129 # Django's collector would materialize every one of those Scripts' Jobs to delete them. 

130 # A module's scripts can accumulate thousands of jobs, exhausting memory. Batch-delete 

131 # the child Scripts' jobs up front, in a single queryset (no per-script loop), before 

132 # delegating to the cascade. The transaction rolls the job deletions back if the parent 

133 # delete fails; note it does not cover ManagedFile.delete removing the file from disk, 

134 # which happens before the DB delete and is not transactional. See #22812. 

135 using = using or router.db_for_write(self.__class__, instance=self) 

136 with transaction.atomic(using=using): 

137 script_type = ContentType.objects.get_for_model(Script, for_concrete_model=False) 

138 child_jobs = Job.objects.using(using).filter( 

139 object_type=script_type, 

140 object_id__in=self.scripts.values_list('pk', flat=True), 

141 ) 

142 batch_delete_jobs(child_jobs) 

143 return super().delete(using, *args, **kwargs) 

144 delete.alters_data = True 

145 

146 @property 

147 def ordered_scripts(self): 

148 script_objects = {s.name: s for s in self.scripts.all()} 

149 ordered = [ 

150 script_objects.pop(sc) for sc in self.module_scripts.keys() if sc in script_objects 

151 ] 

152 ordered.extend(script_objects.values()) 

153 return ordered 

154 

155 @cached_property 

156 def module_scripts(self): 

157 

158 def _get_name(cls): 

159 # For child objects in submodules use the full import path w/o the root module as the name 

160 return cls.full_name.split(".", maxsplit=1)[1] 

161 

162 try: 

163 module = self.get_module() 

164 except Exception as e: 

165 self.error = e 

166 logger.error(f"Failed to load script: {self.python_name} error: {e}") 

167 module = None 

168 

169 scripts = {} 

170 ordered = getattr(module, 'script_order', []) 

171 

172 for cls in ordered: 

173 scripts[_get_name(cls)] = cls 

174 for name, cls in inspect.getmembers(module, is_script): 

175 if cls not in ordered: 

176 scripts[_get_name(cls)] = cls 

177 

178 return scripts 

179 

180 def sync_classes(self): 

181 """ 

182 Syncs the file-based module to the database, adding and removing individual Script objects 

183 in the database as needed. 

184 """ 

185 if self.id: 

186 db_classes = { 

187 script.name: script for script in self.scripts.all() 

188 } 

189 else: 

190 db_classes = {} 

191 

192 db_classes_set = set(db_classes.keys()) 

193 module_classes_set = set(self.module_scripts.keys()) 

194 

195 # remove any existing db classes if they are no longer in the file 

196 removed = db_classes_set - module_classes_set 

197 for name in removed: 

198 db_classes[name].delete(soft_delete=True) 

199 

200 added = module_classes_set - db_classes_set 

201 for name in added: 

202 Script.objects.create( 

203 module=self, 

204 name=name, 

205 is_executable=True, 

206 ) 

207 sync_classes.alters_data = True 

208 

209 def sync_data(self): 

210 super().sync_data() 

211 sync_data.alters_data = True 

212 

213 def save(self, *args, **kwargs): 

214 self.file_root = ManagedFileRootPathChoices.SCRIPTS 

215 super().save(*args, **kwargs) 

216 

217 # Sync script classes after the module has been saved. This is the 

218 # single intended synchronization path for ScriptModule saves. 

219 self.sync_classes()