Clean up migrations

netbox/extras/migrations/0109_script_model.py

@@ -1,12 +1,10 @@
-# Generated by Django 4.2.9 on 2024-02-05 21:37
 import inspect
 import os
-
-from django.conf import settings
-from django.db import migrations, models
-import django.db.models.deletion
 from importlib.machinery import SourceFileLoader
 
+import django.db.models.deletion
+from django.conf import settings
+from django.db import migrations, models
 
 #
 # Note: This has a couple dependencies on the codebase if doing future modifications:
@@ -14,104 +12,117 @@ from importlib.machinery import SourceFileLoader
 # settings.SCRIPTS_ROOT and settings.REPORTS_ROOT to be in settings
 #
 
-def update_event_rules(apps, schema_editor):
-    Script = apps.get_model('extras', 'Script')
-    ContentType = apps.get_model('contenttypes', 'ContentType')
-    ct = ContentType.objects.filter(app_label='extras', model='script').first()
-
-    EventRule = apps.get_model('extras', 'EventRule')
-    ct_script_module = ContentType.objects.filter(app_label='extras', model='scriptmodule').first()
-    for rule in EventRule.objects.filter(action_object_type=ct_script_module):
-        name = rule.action_parameters.get('script_name')
-        obj, created = Script.objects.get_or_create(
-            module_id=rule.action_object_id,
-            name=name,
-            defaults={'is_valid': False}
-        )
-
-        rule.action_object_type = ct
-        rule.action_object_id = obj.id
-        rule.save()
+ROOT_PATHS = {
+    'scripts': settings.SCRIPTS_ROOT,
+    'reports': settings.REPORTS_ROOT,
+}
 
 
-def full_path(instance):
-    root_path = {
-        'scripts': settings.SCRIPTS_ROOT,
-        'reports': settings.REPORTS_ROOT,
-    }[instance.file_root]
-    return os.path.join(root_path, instance.file_path)
+def get_full_path(scriptmodule):
+    """
+    Return the full path to a ScriptModule's file on disk.
+    """
+    root_path = ROOT_PATHS[scriptmodule.file_root]
+    return os.path.join(root_path, scriptmodule.file_path)
 
 
-def python_name(instance):
-    path, filename = os.path.split(full_path(instance))
-    name = os.path.splitext(filename)[0]
-    if name == '__init__':
-        # File is a package
-        return os.path.basename(path)
-    else:
-        return name
+def get_python_name(scriptmodule):
+    """
+    Return the Python name of a ScriptModule's file on disk.
+    """
+    path, filename = os.path.split(scriptmodule.file_path)
+    return os.path.splitext(filename)[0]
 
 
 def is_script(obj):
+    """
+    Returns True if the passed Python object is a Script or Report.
+    """
     from extras.scripts import Script
     from extras.reports import Report
 
     try:
-        return (issubclass(obj, Report) and obj != Report) or (issubclass(obj, Script) and obj != Script)
+        if issubclass(obj, Report) and obj != Report:
+            return True
+        if issubclass(obj, Script) and obj != Script:
+            return True
     except TypeError:
-        return False
+        pass
+    return False
 
 
-def get_python_module(instance):
-    loader = SourceFileLoader(python_name(instance), full_path(instance))
-    module = loader.load_module()
-    return module
-
-
-def get_module_scripts(apps, instance):
-
-    def _get_name(cls):
+def get_module_scripts(scriptmodule):
+    """
+    Return a dictionary mapping of name and script class inside the passed ScriptModule.
+    """
+    def get_name(cls):
         # For child objects in submodules use the full import path w/o the root module as the name
         return cls.full_name.split(".", maxsplit=1)[1]
 
-    try:
-        module = get_python_module(instance)
-    except Exception as e:
-        module = None
-        return
+    loader = SourceFileLoader(get_python_name(scriptmodule), get_full_path(scriptmodule))
+    module = loader.load_module()
 
     scripts = {}
     ordered = getattr(module, 'script_order', [])
 
     for cls in ordered:
-        scripts[_get_name(cls)] = cls
+        scripts[get_name(cls)] = cls
     for name, cls in inspect.getmembers(module, is_script):
         if cls not in ordered:
-            scripts[_get_name(cls)] = cls
+            scripts[get_name(cls)] = cls
 
     return scripts
 
 
 def update_scripts(apps, schema_editor):
-    ScriptModule = apps.get_model('extras', 'ScriptModule')
-    Script = apps.get_model('extras', 'Script')
+    """
+    Create a new Script object for each script inside each existing ScriptModule, and update any related jobs to
+    reference the new Script object.
+    """
     ContentType = apps.get_model('contenttypes', 'ContentType')
+    Script = apps.get_model('extras', 'Script')
+    ScriptModule = apps.get_model('extras', 'ScriptModule')
     Job = apps.get_model('core', 'Job')
-    ct = ContentType.objects.filter(app_label='extras', model='script').first()
-    ct_module = ContentType.objects.filter(app_label='extras', model='scriptmodule').first()
+
+    script_ct = ContentType.objects.get_for_model(Script)
+    scriptmodule_ct = ContentType.objects.get_for_model(ScriptModule)
 
     for module in ScriptModule.objects.all():
-        module_scripts = get_module_scripts(apps, module)
-        if module_scripts:
-            for script in module_scripts.keys():
-                obj = Script.objects.create(
-                    name=script,
-                    module=ScriptModule.objects.get(file_root=module.file_root, file_path=module.file_path),
-                )
+        for script_name in get_module_scripts(module):
+            script = Script.objects.create(
+                name=script_name,
+                module=module,
+            )
 
-                # update all jobs associated with this module/name to point to the new script obj
-                if ct:
-                    Job.objects.filter(object_type=ct_module, object_id=module.id, name=script).update(object_type=ct, object_id=obj.id)
+            # Update all Jobs associated with this ScriptModule & script name to point to the new Script object
+            Job.objects.filter(
+                object_type=scriptmodule_ct,
+                object_id=module.pk,
+                name=script_name
+            ).update(object_type=script_ct, object_id=script.pk)
+
+
+def update_event_rules(apps, schema_editor):
+    """
+    Update any existing EventRules for scripts. Change action_object_type from ScriptModule to Script, and populate
+    the ID of the related Script object.
+    """
+    ContentType = apps.get_model('contenttypes', 'ContentType')
+    Script = apps.get_model('extras', 'Script')
+    ScriptModule = apps.get_model('extras', 'ScriptModule')
+    EventRule = apps.get_model('extras', 'EventRule')
+
+    script_ct = ContentType.objects.get_for_model(Script)
+    scriptmodule_ct = ContentType.objects.get_for_model(ScriptModule)
+
+    for eventrule in EventRule.objects.filter(action_object_type=scriptmodule_ct):
+        name = eventrule.action_parameters.get('script_name')
+        obj, created = Script.objects.get_or_create(
+            module_id=eventrule.action_object_id,
+            name=name,
+            defaults={'is_executable': False}
+        )
+        EventRule.objects.filter(pk=eventrule.pk).update(action_object_type=script_ct, action_object_id=obj.id)
 
 
 class Migration(migrations.Migration):

netbox/extras/migrations/0110_remove_eventrule_action_parameters.py

@@ -1,12 +1,10 @@
-# Generated by Django 4.2.9 on 2024-02-13 23:20
-
 from django.db import migrations
 
 
 class Migration(migrations.Migration):
 
     dependencies = [
-        ('extras', '0109_script_models'),
+        ('extras', '0109_script_model'),
     ]
 
     operations = [

netbox/extras/reports.py

@@ -6,6 +6,7 @@ __all__ = (
 )
 
 
+# Required by extras/migrations/0109_script_models.py
 class Report(BaseScript):
 
     #

netbox/netbox/settings.py

@@ -156,6 +156,7 @@ REMOTE_AUTH_SUPERUSERS = getattr(configuration, 'REMOTE_AUTH_SUPERUSERS', [])
 REMOTE_AUTH_STAFF_GROUPS = getattr(configuration, 'REMOTE_AUTH_STAFF_GROUPS', [])
 REMOTE_AUTH_STAFF_USERS = getattr(configuration, 'REMOTE_AUTH_STAFF_USERS', [])
 REMOTE_AUTH_GROUP_SEPARATOR = getattr(configuration, 'REMOTE_AUTH_GROUP_SEPARATOR', '|')
+# Required by extras/migrations/0109_script_models.py
 REPORTS_ROOT = getattr(configuration, 'REPORTS_ROOT', os.path.join(BASE_DIR, 'reports')).rstrip('/')
 RQ_DEFAULT_TIMEOUT = getattr(configuration, 'RQ_DEFAULT_TIMEOUT', 300)
 RQ_RETRY_INTERVAL = getattr(configuration, 'RQ_RETRY_INTERVAL', 60)