mirror of
https://github.com/netbox-community/netbox.git
synced 2024-05-10 07:54:54 +00:00
116 lines
4.5 KiB
Python
116 lines
4.5 KiB
Python
from cacheops import invalidate_model
|
|
from django.apps import apps
|
|
from django.core.management.base import BaseCommand, CommandError
|
|
|
|
from utilities.fields import NaturalOrderingField
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = "Recalculate natural ordering values for the specified models"
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument(
|
|
'args', metavar='app_label.ModelName', nargs='*',
|
|
help='One or more specific models (each prefixed with its app_label) to renaturalize',
|
|
)
|
|
|
|
def _get_models(self, names):
|
|
"""
|
|
Compile a list of models to be renaturalized. If no names are specified, all models which have one or more
|
|
NaturalOrderingFields will be included.
|
|
"""
|
|
models = []
|
|
|
|
if names:
|
|
# Collect all NaturalOrderingFields present on the specified models
|
|
for name in names:
|
|
try:
|
|
app_label, model_name = name.split('.')
|
|
except ValueError:
|
|
raise CommandError(
|
|
f"Invalid format: {name}. Models must be specified in the form app_label.ModelName."
|
|
)
|
|
try:
|
|
app_config = apps.get_app_config(app_label)
|
|
except LookupError as e:
|
|
raise CommandError(str(e))
|
|
try:
|
|
model = app_config.get_model(model_name)
|
|
except LookupError:
|
|
raise CommandError(f"Unknown model: {app_label}.{model_name}")
|
|
fields = [
|
|
field for field in model._meta.concrete_fields if type(field) is NaturalOrderingField
|
|
]
|
|
if not fields:
|
|
raise CommandError(
|
|
f"Invalid model: {app_label}.{model_name} does not employ natural ordering"
|
|
)
|
|
models.append(
|
|
(model, fields)
|
|
)
|
|
|
|
else:
|
|
# Find *all* models with NaturalOrderingFields
|
|
for app_config in apps.get_app_configs():
|
|
for model in app_config.models.values():
|
|
fields = [
|
|
field for field in model._meta.concrete_fields if type(field) is NaturalOrderingField
|
|
]
|
|
if fields:
|
|
models.append(
|
|
(model, fields)
|
|
)
|
|
|
|
return models
|
|
|
|
def handle(self, *args, **options):
|
|
|
|
models = self._get_models(args)
|
|
|
|
if options['verbosity']:
|
|
self.stdout.write(f"Renaturalizing {len(models)} models.")
|
|
|
|
for model, fields in models:
|
|
for field in fields:
|
|
target_field = field.target_field
|
|
naturalize = field.naturalize_function
|
|
count = 0
|
|
|
|
# Print the model and field name
|
|
if options['verbosity']:
|
|
self.stdout.write(
|
|
f"{model._meta.label}.{field.target_field} ({field.name})... ",
|
|
ending='\n' if options['verbosity'] >= 2 else ''
|
|
)
|
|
self.stdout.flush()
|
|
|
|
# Find all unique values for the field
|
|
queryset = model.objects.values_list(target_field, flat=True).order_by(target_field).distinct()
|
|
for value in queryset:
|
|
naturalized_value = naturalize(value, max_length=field.max_length)
|
|
|
|
if options['verbosity'] >= 2:
|
|
self.stdout.write(f" {value} -> {naturalized_value}", ending='')
|
|
self.stdout.flush()
|
|
|
|
# Update each unique field value in bulk
|
|
changed = model.objects.filter(name=value).update(**{field.name: naturalized_value})
|
|
|
|
if options['verbosity'] >= 2:
|
|
self.stdout.write(f" ({changed})")
|
|
count += changed
|
|
|
|
# Print the total count of alterations for the field
|
|
if options['verbosity'] >= 2:
|
|
self.stdout.write(self.style.SUCCESS(
|
|
f"{count} {model._meta.verbose_name_plural} updated ({queryset.count()} unique values)"
|
|
))
|
|
elif options['verbosity']:
|
|
self.stdout.write(self.style.SUCCESS(str(count)))
|
|
|
|
# Invalidate cached queries
|
|
invalidate_model(model)
|
|
|
|
if options['verbosity']:
|
|
self.stdout.write(self.style.SUCCESS("Done."))
|