nautobot 2.0.0a3__py3-none-any.whl → 2.0.0b1__py3-none-any.whl

nautobot/extras/utils.py

@@ -1,9 +1,7 @@
 import collections
 import hashlib
 import hmac
-import inspect
 import logging
-import pkgutil
 import re
 import sys
 
@@ -12,18 +10,16 @@ from django.conf import settings
 from django.contrib.contenttypes.models import ContentType
 from django.core.validators import ValidationError
 from django.db.models import Q
-from django.template.defaultfilters import slugify
 from django.template.loader import get_template, TemplateDoesNotExist
 from django.utils.deconstruct import deconstructible
-from taggit.managers import _TaggableManager
 
-# 2.0 TODO: remove `is_taggable` import here; included for now for backwards compatibility with <1.4 code.
-from nautobot.core.models.utils import find_models_with_matching_fields, is_taggable  # noqa: F401
+from nautobot.core.choices import ColorChoices
+from nautobot.core.models.managers import TagsManager
+from nautobot.core.models.utils import find_models_with_matching_fields
 from nautobot.extras.constants import (
     EXTRAS_FEATURES,
     JOB_MAX_GROUPING_LENGTH,
     JOB_MAX_NAME_LENGTH,
-    JOB_MAX_SOURCE_LENGTH,
     JOB_OVERRIDABLE_FIELDS,
 )
 from nautobot.extras.registry import registry
@@ -136,17 +132,20 @@ class FeatureQuery:
         # initialization of the application, before `ExtrasConfig.ready` is called.
         # Calling `populate_model_features_registry` in `ExtrasConfig.ready` would lead to an outdated `model_features`
         # `registry` record being used by `FeatureQuery`.
-        # As the `populate_model_features_registry` can be resource-intensive. The check the conditional check is used to
-        # avoid calling the function multiple times and optimize the performance of the application.
-        # TODO(timizuo): Provide a better solution for this; check https://github.com/nautobot/nautobot/pull/3360/files#r1131373797 comment.
-        if not registry["model_features"].get("relationships"):
-            populate_model_features_registry()
+
+        populate_model_features_registry()
         query = Q()
-        for app_label, models in registry["model_features"][self.feature].items():
+        for app_label, models in self.as_dict():
             query |= Q(app_label=app_label, model__in=models)
 
         return query
 
+    def as_dict(self):
+        """
+        Given an extras feature, return a dict of app_label: [models] for content type lookup
+        """
+        return registry["model_features"][self.feature].items()
+
     def get_choices(self):
         """
         Given an extras feature, return a list of 2-tuple of `(model_label, pk)`
@@ -165,19 +164,19 @@ class FeatureQuery:
 @deconstructible
 class TaggableClassesQuery(FeaturedQueryMixin):
     """
-    Helper class to get ContentType models that implements tags(TaggableManager)
+    Helper class to get ContentType models that implements tags(TagsField)
     """
 
     def list_subclasses(self):
         """
-        Return a list of classes that has implements tags e.g tags = TaggableManager(...)
+        Return a list of classes that has implements tags e.g tags = TagsField(...)
         """
         return [
             _class
             for _class in apps.get_models()
             if (
                 hasattr(_class, "tags")
-                and isinstance(_class.tags, _TaggableManager)
+                and isinstance(_class.tags, TagsManager)
                 and ".tests." not in _class.__module__  # avoid leakage from nautobot.core.tests.test_filters
             )
         ]
@@ -223,7 +222,7 @@ def extras_features(*features):
     return wrapper
 
 
-def populate_model_features_registry():
+def populate_model_features_registry(refresh=False):
     """
     Populate the registry model features with new apps.
 
@@ -243,6 +242,9 @@ def populate_model_features_registry():
     - For each dictionary in lookup_confs, calls lookup_by_field() function to look for all models that have fields with the names given in the dictionary.
     - Groups the results by app and updates the registry model features for each app.
     """
+    if registry.get("populate_model_features_registry_called", False) and not refresh:
+        return
+
     RelationshipAssociation = apps.get_model(app_label="extras", model_name="relationshipassociation")
 
     lookup_confs = [
@@ -267,6 +269,9 @@ def populate_model_features_registry():
         feature_name = lookup_conf["feature_name"]
         registry["model_features"][feature_name] = registry_items
 
+    if not registry.get("populate_model_features_registry_called", False):
+        registry["populate_model_features_registry_called"] = True
+
 
 def generate_signature(request_body, secret):
     """
@@ -328,49 +333,7 @@ def task_queues_as_choices(task_queues):
     return choices
 
 
-# namedtuple class yielded by the jobs_in_directory generator function, below
-# Example: ("devices", <module "devices">, "Hostname", <class "devices.Hostname">, None)
-# Example: ("devices", None, None, None, "error at line 40")
-JobClassInfo = collections.namedtuple(
-    "JobClassInfo",
-    ["module_name", "module", "job_class_name", "job_class", "error"],
-    defaults=(None, None, None, None),  # all parameters except `module_name` are optional and default to None.
-)
-
-
-def jobs_in_directory(path, module_name=None, reload_modules=True, report_errors=False):
-    """
-    Walk the available Python modules in the given directory, and for each module, walk its Job class members.
-
-    Args:
-        path (str): Directory to import modules from, outside of sys.path
-        module_name (str): Specific module name to select; if unspecified, all modules will be inspected
-        reload_modules (bool): Whether to force reloading of modules even if previously loaded into Python.
-        report_errors (bool): If True, when an error is encountered, yield a JobClassInfo with the given error.
-                              If False (default), log the error but do not yield anything.
-
-    Yields:
-        JobClassInfo: (module_name, module, job_class_name, job_class, error)
-    """
-    from .jobs import is_job  # avoid circular import
-
-    for importer, discovered_module_name, _ in pkgutil.iter_modules([path]):
-        if module_name and discovered_module_name != module_name:
-            continue
-        if reload_modules and discovered_module_name in sys.modules:
-            del sys.modules[discovered_module_name]
-        try:
-            module = importer.find_module(discovered_module_name).load_module(discovered_module_name)
-            # Get all members of the module that are Job subclasses
-            for job_class_name, job_class in inspect.getmembers(module, is_job):
-                yield JobClassInfo(discovered_module_name, module, job_class_name, job_class)
-        except Exception as exc:
-            logger.error(f"Unable to load module {discovered_module_name} from {path}: {exc}")
-            if report_errors:
-                yield JobClassInfo(module_name=discovered_module_name, error=exc)
-
-
-def refresh_job_model_from_job_class(job_model_class, job_source, job_class, *, git_repository=None):
+def refresh_job_model_from_job_class(job_model_class, job_class):
     """
     Create or update a job_model record based on the metadata of the provided job_class.
 
@@ -384,13 +347,6 @@ def refresh_job_model_from_job_class(job_model_class, job_source, job_class, *,
     )  # imported here to prevent circular import problem
 
     # Unrecoverable errors
-    if len(job_source) > JOB_MAX_SOURCE_LENGTH:  # Should NEVER happen
-        logger.error(
-            'Unable to store Jobs from "%s" as Job models because the source exceeds %d characters in length!',
-            job_source,
-            JOB_MAX_SOURCE_LENGTH,
-        )
-        return (None, False)
     if len(job_class.__module__) > JOB_MAX_NAME_LENGTH:
         logger.error(
             'Unable to store Jobs from module "%s" as Job models because the module exceeds %d characters in length!',
@@ -435,8 +391,6 @@ def refresh_job_model_from_job_class(job_model_class, job_source, job_class, *,
     existing_job_names = (
         job_model_class.objects.filter(name__startswith=job_name)
         .exclude(
-            source=job_source[:JOB_MAX_SOURCE_LENGTH],
-            git_repository=git_repository,
             module_name=job_class.__module__[:JOB_MAX_NAME_LENGTH],
             job_class_name=job_class.__name__[:JOB_MAX_NAME_LENGTH],
         )
@@ -456,16 +410,15 @@ def refresh_job_model_from_job_class(job_model_class, job_source, job_class, *,
         )
 
     job_model, created = job_model_class.objects.get_or_create(
-        source=job_source[:JOB_MAX_SOURCE_LENGTH],
-        git_repository=git_repository,
         module_name=job_class.__module__[:JOB_MAX_NAME_LENGTH],
         job_class_name=job_class.__name__[:JOB_MAX_NAME_LENGTH],
         defaults={
-            "slug": slugify(job_name)[:JOB_MAX_NAME_LENGTH],
             "grouping": job_class.grouping[:JOB_MAX_GROUPING_LENGTH],
             "name": job_name,
             "is_job_hook_receiver": issubclass(job_class, JobHookReceiver),
             "is_job_button_receiver": issubclass(job_class, JobButtonReceiver),
+            "read_only": job_class.read_only,
+            "supports_dryrun": job_class.supports_dryrun,
             "installed": True,
             "enabled": False,
         },
@@ -474,12 +427,19 @@ def refresh_job_model_from_job_class(job_model_class, job_source, job_class, *,
     if job_name != default_job_name:
         job_model.name_override = True
 
+    if created and job_model.module_name.startswith("nautobot."):
+        # System jobs should be enabled by default when first created
+        job_model.enabled = True
+
     for field_name in JOB_OVERRIDABLE_FIELDS:
         # Was this field directly inherited from the job before, or was it overridden in the database?
         if not getattr(job_model, f"{field_name}_override", False):
             # It was inherited and not overridden
             setattr(job_model, field_name, getattr(job_class, field_name))
 
+    # set supports_dryrun to match the property on the job class
+    job_model.supports_dryrun = job_class.supports_dryrun
+
     if not created:
         # Mark it as installed regardless
         job_model.installed = True
@@ -487,12 +447,10 @@ def refresh_job_model_from_job_class(job_model_class, job_source, job_class, *,
     job_model.save()
 
     logger.info(
-        '%s Job "%s: %s" from <%s%s: %s>',
+        '%s Job "%s: %s" from <%s>',
         "Created" if created else "Refreshed",
         job_model.grouping,
         job_model.name,
-        job_source,
-        f" {git_repository.name}" if git_repository is not None else "",
         job_class.__name__,
     )
 
@@ -508,110 +466,126 @@ def remove_prefix_from_cf_key(field_name):
     return field_name[3:]
 
 
-def check_if_key_is_graphql_safe(model_name, key):
+def check_if_key_is_graphql_safe(model_name, key, field_name="key"):
     """
     Helper method to check if a key field is Python/GraphQL safe.
-    Used in CustomField for now, should be used in ComputedField and Relationship as well.
+    Used in CustomField, ComputedField and Relationship models.
     """
     graphql_safe_pattern = re.compile("[_A-Za-z][_0-9A-Za-z]*")
     if not graphql_safe_pattern.fullmatch(key):
         raise ValidationError(
             {
-                "key": "This key is not Python/GraphQL safe. Please do not start the key with a digit and do not use hyphens or whitespace"
+                f"{field_name}": f"This {field_name} is not Python/GraphQL safe. Please do not start the {field_name} with a digit and do not use hyphens or whitespace"
             }
         )
 
 
+def fixup_null_statuses(*, model, model_contenttype, status_model):
+    """For instances of model that have an invalid NULL status field, create and use a special status_model instance."""
+    instances_to_fixup = model.objects.filter(status__isnull=True)
+    if instances_to_fixup.exists():
+        null_status, _ = status_model.objects.get_or_create(
+            name="NULL",
+            defaults={
+                "color": ColorChoices.COLOR_BLACK,
+                "description": "Created by Nautobot to replace invalid null references",
+            },
+        )
+        null_status.content_types.add(model_contenttype)
+        updated_count = instances_to_fixup.update(status=null_status)
+        print(f"    Found and fixed {updated_count} instances of {model.__name__} that had null 'status' fields.")
+
+
 def migrate_role_data(
-    model,
-    role_model=None,
-    is_choice_field=False,
-    role_choiceset=None,
-    legacy_role="legacy_role",
-    new_role="new_role",
+    model_to_migrate,
+    *,
+    from_role_field_name,
+    from_role_model=None,
+    from_role_choiceset=None,
+    to_role_field_name,
+    to_role_model=None,
+    to_role_choiceset=None,
     is_m2m_field=False,
 ):
     """
-    Migrate legacy_role's `role_model` equivalent record into new_role.
+    Update all `model_to_migrate` with a value for `to_role_field` based on `from_role_field` values.
 
     Args:
-        model(Model): Model with role field to alter
-        role_model(Model): Role Model (optional)
-        is_choice_field(bool): True if Model's role field is a choice field not an FK field else False
-        role_choiceset(ChoiceSet): Role ChoiceSet (optional)
-        legacy_role(str): Models role field legacy name; This is the current role field name
-        new_role(str): Models role field new name; This is the name role field should be updated to
-        is_m2m_field(bool): True if the role field is a ManyToManyField, else False
-    """
-
-    # Retrieve the queryset of `model` instances that have a value for the `legacy_role` field
-    queryset = model.objects.filter(**{legacy_role + "__isnull": False}).only(*["pk", legacy_role])
-    instances_to_update_data = get_instances_to_pk_and_new_role(
-        queryset=queryset,
-        role_model=role_model,
-        role_choiceset=role_choiceset,
-        legacy_role=legacy_role,
-        new_role=new_role,
-        is_m2m_field=is_m2m_field,
-    )
-
-    if instances_to_update_data:
-        if is_m2m_field:
-            # Update the new_role field using the set() method for ManyToManyFields
-            for data in instances_to_update_data:
-                instance = model.objects.get(pk=data["id"])
-                getattr(instance, new_role).set(data[new_role])
+        model_to_migrate (Model): Model with role fields to alter
+        from_role_field_name (str): Name of the field on `model_to_migrate` to use as source data
+        from_role_model (Model): If `from_role_field` is a ForeignKey or M2M field, the corresponding model for it
+        from_role_choiceset (ChoiceSet): If `from_role_field` is a choices field, the corresponding ChoiceSet for it
+        to_role_field_name (str): Name of the field on `model_to_migrate` to update based on the `from_role_field`
+        to_role_model (Model): If `to_role_field` is a ForeignKey or M2M field, the corresponding model for it
+        to_role_choiceset (ChoiceSet): If `to_role_field` is a choices field, the corresponding ChoiceSet for it
+        is_m2m_field (bool): True if the role fields are both ManyToManyFields, else False
+    """
+    if from_role_model is not None:
+        assert from_role_choiceset is None
+        if to_role_model is not None:
+            assert to_role_choiceset is None
+            # Mapping "from" model instances to corresponding "to" model instances
+            roles_translation_mapping = {
+                # Use .filter().first(), not .get() because "to" role might not exist, especially on reverse migrations
+                from_role: to_role_model.objects.filter(name=from_role.name).first()
+                for from_role in from_role_model.objects.all()
+            }
         else:
-            # Update the new_role field using the set() method for Foreign Key fields
-            instances_to_update_data = [model(**data) for data in instances_to_update_data]
-            model.objects.bulk_update(instances_to_update_data, fields=[new_role], batch_size=1000)
-
-
-def get_instances_to_pk_and_new_role(queryset, role_model, role_choiceset, legacy_role, new_role, is_m2m_field):
-    """
-    Return a list of dictionaries containing the primary key and the new role value for each instance in the queryset.
-
-    The new role value is determined by the following logic:
-    - If `role_model` is not None, the value of the `legacy_role` field for each instance is used to
-      query the `role_model` for a matching role. If a match is found, it is used as the new role value.
-      If the `legacy_role` field is a ManyToManyField, all role names are obtained and used to query
-      the `role_model`.
-    - If `role_model` is None, the `legacy_role` field is assumed to be a ChoiceField, and the
-      `role_choiceset` is used to look up the value corresponding to the label of the `legacy_role` field.
-      This value is used as the new role value.
-
-    Args:
-        queryset (QuerySet): A queryset of instances to get the data for.
-        role_model (Model): A model class to use to look up the new role value.
-        role_choiceset (ChoiceSet): A `ChoiceSet` containing the choices for the new role field.
-        legacy_role (str): The name of the field on each instance containing the legacy role value.
-        new_role (str): The name of the field on each instance that will be updated with the new role value.
-        is_m2m_field (bool): Whether the `legacy_role` field is a ManyToManyField.
-    """
-    data = []
-
-    for item in queryset:
-        role_equivalent = None
-        if role_model:
-            # Get the value for the legacy role field e.g. Device.role, IPAddress.role
-            legacy_role_field = getattr(item, legacy_role)
-            if is_m2m_field:
-                # In the case of a m2m field, `legacy_role_field` we would need to get all role names e.g ConfigContext.roles.values_list("name", flat=True)
-                role_names = legacy_role_field.values_list("name", flat=True)
-                role_equivalent = role_model.objects.filter(Q(name__in=role_names) | Q(slug__in=role_names))
-            else:
-                if not isinstance(legacy_role_field, str):
-                    legacy_role_field = legacy_role_field.name
-                try:
-                    role_equivalent = role_model.objects.get(Q(name=legacy_role_field) | Q(slug=legacy_role_field))
-                except role_model.DoesNotExist:
-                    logger.error(f"Role with name {legacy_role_field} not found")
+            assert to_role_choiceset is not None
+            # Mapping "from" model instances to corresponding "to" choices
+            # We need to use `label` to look up the from_role instance, but `value` is what we set for the to_role_field
+            inverted_to_role_choiceset = {label: value for value, label in to_role_choiceset.CHOICES}
+            roles_translation_mapping = {
+                from_role: inverted_to_role_choiceset.get(from_role.name, None)
+                for from_role in from_role_model.objects.all()
+            }
+    else:
+        assert from_role_choiceset is not None
+        if to_role_model is not None:
+            assert to_role_choiceset is None
+            # Mapping "from" choices to corresponding "to" model instances
+            roles_translation_mapping = {
+                # Use .filter().first(), not .get() because "to" role might not exist, especially on reverse migrations
+                from_role_value: to_role_model.objects.filter(name=from_role_label).first()
+                for from_role_value, from_role_label in from_role_choiceset.CHOICES
+            }
         else:
-            # ChoiceSet.CHOICES has to be inverted to obtain the value using its label
-            # i.e {"vrf": "VRF"} --> {"VRF": "vrf"}
-            inverted_role_choiceset = {label: value for value, label in role_choiceset.CHOICES}
-            role_equivalent = inverted_role_choiceset.get(legacy_role_field)
-        model_data = {"id": item.pk, new_role: role_equivalent}
-        data.append(model_data)
-
-    return data
+            assert to_role_choiceset is not None
+            # Mapping "from" choices to corresponding "to" choices; we don't currently use this case, but it should work
+            # We need to use `label` to look up the from_role instance, but `value` is what we set for the to_role_field
+            inverted_to_role_choiceset = {label: value for value, label in to_role_choiceset.CHOICES}
+            roles_translation_mapping = {
+                from_role_value: inverted_to_role_choiceset.get(from_role_label, None)
+                for from_role_value, from_role_label in from_role_choiceset.CHOICES
+            }
+
+    if not is_m2m_field:
+        # Bulk updates of a single field are easy enough...
+        for from_role_value, to_role_value in roles_translation_mapping.items():
+            if to_role_value is not None:
+                updated_count = model_to_migrate.objects.filter(**{from_role_field_name: from_role_value}).update(
+                    **{to_role_field_name: to_role_value}
+                )
+                logger.info(
+                    'Updated %d %s records to reference %s "%s"',
+                    updated_count,
+                    model_to_migrate._meta.label,
+                    to_role_field_name,
+                    to_role_value.name if to_role_model else to_role_value,
+                )
+    else:
+        # ...but we have to update each instance's M2M field independently?
+        for instance in model_to_migrate.objects.all():
+            to_role_set = {
+                roles_translation_mapping[from_role_value]
+                for from_role_value in getattr(instance, from_role_field_name).all()
+            }
+            # Discard any null values
+            to_role_set.discard(None)
+            getattr(instance, to_role_field_name).set(to_role_set)
+        logger.info(
+            "Updated %d %s record %s M2M fields",
+            model_to_migrate.objects.count(),
+            model_to_migrate._meta.label,
+            to_role_field_name,
+        )