From 63867518eefd27abf9a18f24d07c00a65267a9f8 Mon Sep 17 00:00:00 2001
From: Jeff Bradberry <jeff.bradberry@gmail.com>
Date: Mon, 10 Jan 2022 14:15:58 -0500
Subject: [PATCH] Add a new parameter --disconnect to register_peers

To allow links between Receptor nodes to be removed from the database.
---
 .../management/commands/register_peers.py     | 41 +++++++++++++------
 1 file changed, 29 insertions(+), 12 deletions(-)

diff --git a/awx/main/management/commands/register_peers.py b/awx/main/management/commands/register_peers.py
index 19b40c0c59..dd11efd6d1 100644
--- a/awx/main/management/commands/register_peers.py
+++ b/awx/main/management/commands/register_peers.py
@@ -9,23 +9,40 @@ class Command(BaseCommand):
     Register the peers of a receptor node.
     """
 
+    help = "Register or remove links between Receptor nodes."
+
     def add_arguments(self, parser):
-        parser.add_argument('source', type=str, help="")
-        parser.add_argument('--peers', type=str, nargs='+', required=True, help="")
+        parser.add_argument('source', type=str, help="Receptor node opening the connections.")
+        parser.add_argument('--peers', type=str, nargs='+', required=False, help="Nodes that the source node connects out to.")
+        parser.add_argument('--disconnect', type=str, nargs='+', required=False, help="Nodes that should no longer be connected to by the source node.")
 
     def handle(self, **options):
         nodes = Instance.objects.in_bulk(field_name='hostname')
         if options['source'] not in nodes:
             raise CommandError(f"Host {options['source']} is not a registered instance.")
-        missing_peers = set(options['peers']) - set(nodes)
-        if missing_peers:
-            missing = ' '.join(missing_peers)
-            raise CommandError(f"Peers not currently registered as instances: {missing}")
+        if not options['peers'] and not options['disconnect']:
+            raise CommandError("One of the options --peers or --disconnect is required.")
 
-        results = 0
-        for target in options['peers']:
-            instance, created = InstanceLink.objects.get_or_create(source=nodes[options['source']], target=nodes[target])
-            if created:
-                results += 1
+        if options['peers']:
+            missing_peers = set(options['peers']) - set(nodes)
+            if missing_peers:
+                missing = ' '.join(missing_peers)
+                raise CommandError(f"Peers not currently registered as instances: {missing}")
 
-        print(f"{results} new peer links added to the database.")
+            results = 0
+            for target in options['peers']:
+                instance, created = InstanceLink.objects.get_or_create(source=nodes[options['source']], target=nodes[target])
+                if created:
+                    results += 1
+
+            print(f"{results} new peer links added to the database.")
+
+        if options['disconnect']:
+            results = 0
+            for target in options['disconnect']:
+                if target not in nodes:  # Be permissive, the node might have already been de-registered.
+                    continue
+                n, _ = InstanceLink.objects.filter(source=nodes[options['source']], target=nodes[target]).delete()
+                results += n
+
+            print(f"{results} peer links removed from the database.")