Add a new parameter --disconnect to register_peers

To allow links between Receptor nodes to be removed from the database.
This commit is contained in:
Jeff Bradberry
2022-01-10 14:15:58 -05:00
parent 37907ad348
commit 63867518ee

View File

@@ -9,23 +9,40 @@ class Command(BaseCommand):
Register the peers of a receptor node. Register the peers of a receptor node.
""" """
help = "Register or remove links between Receptor nodes."
def add_arguments(self, parser): def add_arguments(self, parser):
parser.add_argument('source', type=str, help="") parser.add_argument('source', type=str, help="Receptor node opening the connections.")
parser.add_argument('--peers', type=str, nargs='+', required=True, help="") 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): def handle(self, **options):
nodes = Instance.objects.in_bulk(field_name='hostname') nodes = Instance.objects.in_bulk(field_name='hostname')
if options['source'] not in nodes: if options['source'] not in nodes:
raise CommandError(f"Host {options['source']} is not a registered instance.") raise CommandError(f"Host {options['source']} is not a registered instance.")
missing_peers = set(options['peers']) - set(nodes) if not options['peers'] and not options['disconnect']:
if missing_peers: raise CommandError("One of the options --peers or --disconnect is required.")
missing = ' '.join(missing_peers)
raise CommandError(f"Peers not currently registered as instances: {missing}")
results = 0 if options['peers']:
for target in options['peers']: missing_peers = set(options['peers']) - set(nodes)
instance, created = InstanceLink.objects.get_or_create(source=nodes[options['source']], target=nodes[target]) if missing_peers:
if created: missing = ' '.join(missing_peers)
results += 1 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.")