Add support for Tags as AWX Groups with m2m-changed signal handling, including...

Add support for Tags as AWX Groups with m2m-changed signal handling, including host-group association and disassociation - INFRA-10378

netbox_awx_plugin/signals.py

-from django.db.models.signals import post_save, pre_save, post_delete
+import logging
+from django.db.models.signals import post_save, pre_save, post_delete, m2m_changed
 from django.dispatch import receiver
 from django_rq import get_queue
 from dcim.models import Interface, Device, DeviceRole, DeviceType, Site
 from ipam.models import Prefix
 from virtualization.models import VirtualMachine
+from extras.models import Tag
 from netbox.config import get_config
 from netbox.constants import RQ_QUEUE_DEFAULT
 from .models import AWXInventory
 
-import logging
-
-
 logger = logging.getLogger(__name__)
 
-__all__ = [
-    "update_host",
-]
-
-
+# Configure queue name from settings
 queue_name = get_config().QUEUE_MAPPINGS.get("awx", RQ_QUEUE_DEFAULT)
 
-
-@receiver(post_save, sender=Site)
-@receiver(post_save, sender=DeviceRole)
-@receiver(post_save, sender=DeviceType)
-@receiver(post_save, sender=Prefix)
-def save_group_instance(sender, instance, created, **kwargs):
-    rq_queue = get_queue(queue_name)
-    for inventory in AWXInventory.objects.all():
-        if inventory.enabled == True:
-            params = {"inventory": inventory, "sender": sender, "instance": instance}
-            rq_queue.enqueue("netbox_awx_plugin.synchronization.sync_group", **params)
-
-
-@receiver(pre_save, sender=Site)
-@receiver(pre_save, sender=DeviceRole)
-@receiver(pre_save, sender=DeviceType)
-@receiver(pre_save, sender=Prefix)
-def pre_save_group_instance(sender, instance, **kwargs):
-    if sender.objects.filter(id=instance.id).exists():
-        instance.__original_object = sender.objects.get(id=instance.id)
-
-
-@receiver(post_delete, sender=Site)
-@receiver(post_delete, sender=DeviceRole)
-@receiver(post_delete, sender=DeviceType)
-@receiver(post_delete, sender=Prefix)
-def save_group_instance(sender, instance, **kwargs):
+def enqueue_task(task_name, inventory, sender, instance):
+    """
+    Helper function to enqueue AWX tasks in the RQ queue.
+    """
     rq_queue = get_queue(queue_name)
-    for inventory in AWXInventory.objects.all():
-        if inventory.enabled == True:
-            params = {"inventory": inventory, "sender": sender, "instance": instance}
-            rq_queue.enqueue("netbox_awx_plugin.synchronization.delete_group", **params)
-
-
-@receiver(post_save, sender=Device)
-@receiver(post_save, sender=VirtualMachine)
-def save_device(sender, instance, created, **kwargs):
-    device = sender.objects.get(id=instance.id)
-    if not device.primary_ip4 is None and device.primary_ip4.dns_name:
-        rq_queue = get_queue(queue_name)
-        for inventory in AWXInventory.objects.all():
-            if inventory.enabled == True:
-                params = {"inventory": inventory, "sender": sender, "instance": device}
-                rq_queue.enqueue("netbox_awx_plugin.synchronization.sync_host", **params)
-
-
-@receiver(post_save, sender=Interface)
-def save_interface(sender, instance, **kwargs):
-    save_device(Device, instance.device, **kwargs)
+    task_params = {"inventory": inventory, "sender": sender, "instance": instance}
+    logger.debug(f"Enqueuing task '{task_name}' for {sender.__name__} instance {instance.pk}")
+    rq_queue.enqueue(f"netbox_awx_plugin.synchronization.{task_name}", **task_params)
+
+def process_inventory_task(task_name, sender, instance):
+    """
+    Processes a task for all enabled AWX inventories.
+    """
+    inventories = AWXInventory.objects.filter(enabled=True)
+    if not inventories.exists():
+        logger.info(f"No enabled AWX inventories found. Skipping {task_name}.")
+        return
+
+    for inventory in inventories:
+        enqueue_task(task_name, inventory, sender, instance)
+
+
+### Signal Handlers for Group-Related Models ###
+@receiver([post_save], sender=Site)
+@receiver([post_save], sender=DeviceRole)
+@receiver([post_save], sender=DeviceType)
+@receiver([post_save], sender=Prefix)
+def handle_group_post_save(sender, instance, created, **kwargs):
+    """
+    Handles post-save events for group-related models like Site, DeviceRole, etc.
+    This synchronizes AWX groups when the object is saved.
+    """
+    process_inventory_task("sync_group", sender, instance)
+
+@receiver([pre_save], sender=Site)
+@receiver([pre_save], sender=DeviceRole)
+@receiver([pre_save], sender=DeviceType)
+@receiver([pre_save], sender=Prefix)
+def handle_group_pre_save(sender, instance, **kwargs):
+    """
+    Handles pre-save events for group-related models to store the original instance for comparison.
+    """
+    if instance.pk and sender.objects.filter(pk=instance.pk).exists():
+        instance._original = sender.objects.get(pk=instance.pk)
+        logger.debug(f"Original object stored for {sender.__name__} instance {instance.pk}")
+
+
+@receiver([post_delete], sender=Site)
+@receiver([post_delete], sender=DeviceRole)
+@receiver([post_delete], sender=DeviceType)
+@receiver([post_delete], sender=Prefix)
+def handle_group_post_delete(sender, instance, **kwargs):
+    """
+    Handles post-delete events for group-related models to remove AWX groups when objects are deleted.
+    """
+    process_inventory_task("delete_group", sender, instance)
+
+### Signal Handlers for Device and VirtualMachine ###
+@receiver([post_save], sender=Device)
+@receiver([post_save], sender=VirtualMachine)
+def handle_device_post_save(sender, instance, **kwargs):
+    """
+    Handles post-save events for Devices and Virtual Machines.
+    Synchronizes AWX hosts if the instance meets the IP and DNS requirements.
+    """
+    instance = sender.objects.select_related("primary_ip4").get(pk=instance.pk)
+    if instance.primary_ip4 and instance.primary_ip4.dns_name:
+        logger.debug(f"Device/VM {instance.pk} meets IP/DNS criteria. Triggering host sync.")
+        process_inventory_task("sync_host", sender, instance)
+    else:
+        logger.debug(f"Device/VM {instance.pk} does not meet IP/DNS criteria for AWX sync.")
+
+
+
+@receiver([post_save], sender=Interface)
+def handle_interface_post_save(sender, instance, **kwargs):
+    """
+    Handles post-save events for Interfaces. Delegates synchronization to the associated Device.
+    """
+    if instance.device:
+        logger.debug(f"Interface {instance.pk} saved, delegating sync to associated Device.")
+        handle_device_post_save(Device, instance.device, **kwargs)
+
+### Tag Synchronization ###
+@receiver(post_save, sender=Tag)
+def handle_tag_post_save(sender, instance, created, **kwargs):
+    """
+    Handles post-save events for Tags. Synchronizes AWX groups when a Tag is saved.
+    """
+    process_inventory_task("sync_group", sender, instance)
+
+@receiver(m2m_changed, sender=Device.tags.through)
+@receiver(m2m_changed, sender=VirtualMachine.tags.through)
+def handle_tag_assignment(sender, instance, action, reverse, model, pk_set, **kwargs):
+    """
+    Handles many-to-many change events (m2m_changed) for Device and VirtualMachine tags.
+    Associates hosts with AWX groups when tags are assigned.
+    """
+    if action == "post_add":
+        logger.debug(f"Tags added to {instance}. Syncing host with AWX group.")
+        process_inventory_task("sync_host", Device if Device.tags.through else VirtualMachine, instance)
+    if action == "post_remove":
+        logger.debug(f"Tags removed to {instance}. Syncing host with AWX group.")
+        process_inventory_task("sync_host", Device if Device.tags.through else VirtualMachine, instance)
+    
\ No newline at end of file

netbox_awx_plugin/synchronization.py

@@ -3,6 +3,7 @@ from ipam.models import Prefix
 from dcim.choices import DeviceStatusChoices
 from virtualization.choices import VirtualMachineStatusChoices
 from virtualization.models import VirtualMachine
+from extras.models import Tag
 from rest_framework import serializers
 import json
 import logging
@@ -11,9 +12,16 @@ from .models import AWXInventory
 
 logger = logging.getLogger(__name__)
 
-group_prefixes = {Site: "site_", DeviceRole: "devicerole_", DeviceType: "devicetype_", Prefix: "prefix_"}
+group_prefixes = {Site: "site_", DeviceRole: "devicerole_", DeviceType: "devicetype_", Prefix: "prefix_", Tag:"tag_"}
 
 
+class TagSerializer(serializers.BaseSerializer):
+    def to_representation(self, instance):
+        return {
+            "name": "{}{}".format(group_prefixes[Tag], instance.slug.replace("-", "_")),
+            "variables": json.dumps({"netbox_tag_name": instance.name}),
+        }
+
 class SiteSerializer(serializers.BaseSerializer):
     def to_representation(self, instance):
         return {
@@ -149,45 +157,83 @@ serializers = {
     Prefix: PrefixSerializer,
     Device: DeviceSerializer,
     VirtualMachine: VMSerializer,
+    Tag: TagSerializer
 }
 
 
 def sync_host(inventory, sender, instance):
     serializer = serializers[sender](instance)
-    host = inventory.get_host(serializer.data["name"])
+    host = inventory.get_host(serializer.data["name"]) 
     if host is None:
+        # If the host doesn't exist, create it.
         inventory.create_host(serializer.data)
         host = inventory.get_host(serializer.data["name"])
     else:
+        # If the host exists, update it if necessary.
         host_serializer = AWXHostSerializer(data=host)
         if not host_serializer.is_valid() or host_serializer.validated_data != serializer.data:
             inventory.update_host(host["id"], serializer.data)
-
+    
+    # Sync group associations
+    current_groups = host["summary_fields"]["groups"]["results"]
+    
+    # Handle associations for each type (Site, DeviceRole, DeviceType, and Tags)
     if instance.site:
-        sync_host_group_association(inventory, host, Site, instance.site)
+        sync_host_group_association(inventory, host, Site, instance.site, current_groups)
     if instance.role:
-        sync_host_group_association(inventory, host, DeviceRole, instance.role)
-    if sender == Device:
-        sync_host_group_association(inventory, host, DeviceType, instance.device_type)
+        sync_host_group_association(inventory, host, DeviceRole, instance.role, current_groups)
+    if isinstance(instance, Device) and instance.device_type:
+        sync_host_group_association(inventory, host, DeviceType, instance.device_type, current_groups)
+    if tags := instance.tags.all():
+        for tag in tags:
+            sync_host_group_association(inventory, host, Tag, tag, current_groups)
 
-    # print(Prefix.objects.filter(prefix__net_contains_or_equals=device.primary_ip4.address))
+    # Disassociate groups that are no longer relevant
+    disassociate_removed_groups(inventory, host, instance, current_groups)
 
 
-def sync_host_group_association(inventory, host, object_type, instance):
+def sync_host_group_association(inventory, host, object_type, instance, current_groups):
+    """
+    Associates the host with the appropriate AWX group.
+    """
     serializer = serializers[object_type](instance)
     group_name = serializer.data["name"]
-    groups = [
-        group
-        for group in host["summary_fields"]["groups"]["results"]
-        if group["name"].startswith(group_prefixes[object_type])
+    
+    # Filter out the groups that match the current object_type (e.g., Site, Tag)
+    relevant_groups = [
+        group for group in current_groups if group["name"].startswith(group_prefixes[object_type])
     ]
-    for group in groups:
-        if group["name"] != group_name:
-            inventory.disassociate_host_group(host["id"], group["id"])
-    if group_name not in [group["name"] for group in groups]:
+    
+    # Check if the host is already in the desired group
+    if group_name not in [group["name"] for group in relevant_groups]:
         group = inventory.get_group(group_name)
-        inventory.associate_host_group(host["id"], group["id"])
-
+        if group:
+            inventory.associate_host_group(host["id"], group["id"])
+            logger.info(f"Host {host['name']} associated with group {group_name}.")
+
+
+def disassociate_removed_groups(inventory, host, instance, current_groups):
+    """
+    Disassociates the host from AWX groups that are no longer relevant.
+    """
+    # Collect the current groups that should remain (based on the instance's attributes)
+    valid_group_names = set()
+    
+    if instance.site:
+        valid_group_names.add(f"{group_prefixes[Site]}{instance.site.slug.replace('-', '_')}")
+    if instance.role:
+        valid_group_names.add(f"{group_prefixes[DeviceRole]}{instance.role.slug.replace('-', '_')}")
+    if isinstance(instance, Device) and instance.device_type:
+        valid_group_names.add(f"{group_prefixes[DeviceType]}{instance.device_type.slug.replace('-', '_')}")
+    if tags := instance.tags.all():
+        valid_group_names.update(
+            f"{group_prefixes[Tag]}{tag.slug.replace('-', '_')}" for tag in tags
+        )
+    # Disassociate from groups that are no longer valid
+    for group in current_groups:
+        if group["name"] not in valid_group_names:
+            inventory.disassociate_host_group(host["id"], group["id"])
+            logger.info(f"Host {host['name']} disassociated from group {group['name']}.")
 
 def delete_host(inventory, sender, instance):
     serializer = serializers[sender](instance)
@@ -224,6 +270,9 @@ def sync_all(job):
         sync_group(inventory, DeviceRole, device_role)
     for device_type in DeviceType.objects.all():
         sync_group(inventory, DeviceType, device_type)
+    for tag in Tag.objects.all():
+        sync_group(inventory, Tag, tag)
+
     for device in Device.objects.all():
         if not device.primary_ip4 is None and device.primary_ip4.dns_name:
             sync_host(inventory, Device, device)