Coverage for utilities/management/commands/calculate_cached_counts.py: 0%

1from collections import defaultdict 

2 

3from django.core.management.base import BaseCommand 

4 

5from netbox.registry import registry 

6from utilities.counters import update_counts 

7 

8 

9class Command(BaseCommand): 

10 help = "Force a recalculation of all cached counter fields" 

11 

12 @staticmethod 

13 def collect_models(): 

14 """ 

15 Query the registry to find all models which have one or more counter fields. Return a mapping of counter fields 

16 to related query names for each model. 

17 """ 

18 models = defaultdict(dict) 

19 

20 for model, field_mappings in registry['counter_fields'].items(): 

21 for field_name, counter_name in field_mappings.items(): 

22 fk_field = model._meta.get_field(field_name) # Interface.device 

23 parent_model = fk_field.related_model # Device 

24 related_query_name = fk_field.related_query_name() # 'interfaces' 

25 models[parent_model][counter_name] = related_query_name 

26 

27 return models 

28 

29 def handle(self, *model_names, **options): 

30 for model, mappings in self.collect_models().items(): 

31 for field_name, related_query in mappings.items(): 

32 update_counts(model, field_name, related_query) 

33 

34 self.stdout.write(self.style.SUCCESS('Finished.'))