Ensure register_peers target is ReceptorAddress

command has 3 "targets", all of which should be a
ReceptorAddress object

- peers, disconnect, exact

Add logic to make sure each entry in those lists
are receptor addresses.

When creating InstanceLink objects, make sure
target is ReceptorAddress, not an Instance.

Signed-off-by: Seth Foster <fosterbseth@gmail.com>
This commit is contained in:
Seth Foster
2024-01-16 17:12:01 -05:00
committed by Seth Foster
parent 15a16b3dd1
commit 9531f8377a

View File

@@ -3,7 +3,7 @@ import warnings
from django.core.management.base import BaseCommand, CommandError from django.core.management.base import BaseCommand, CommandError
from django.db import transaction from django.db import transaction
from awx.main.models import Instance, InstanceLink from awx.main.models import Instance, InstanceLink, ReceptorAddress
class Command(BaseCommand): class Command(BaseCommand):
@@ -28,7 +28,9 @@ class Command(BaseCommand):
def handle(self, **options): def handle(self, **options):
# provides a mapping of hostname to Instance objects # provides a mapping of hostname to Instance objects
nodes = Instance.objects.all().prefetch_related('receptor_addresses').in_bulk(field_name='hostname') nodes = Instance.objects.all().in_bulk(field_name='hostname')
# provides a mapping of address to ReceptorAddress objects
addresses = ReceptorAddress.objects.all().in_bulk(field_name='address')
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.")
@@ -39,11 +41,13 @@ class Command(BaseCommand):
if options['exact'] is not None and options['disconnect']: if options['exact'] is not None and options['disconnect']:
raise CommandError("The option --disconnect may not be used with --exact.") raise CommandError("The option --disconnect may not be used with --exact.")
# make sure peers and source instances only have one receptor address # make sure each target has a receptor address
for hostname, node in nodes.items(): peers = options['peers'] or []
if hostname in options.get('peers', []) or hostname == options['source']: disconnect = options['disconnect'] or []
if node.receptor_addresses.count() > 1: exact = options['exact'] or []
raise CommandError(f"Instance {hostname} has more than one receptor address.") for peer in peers + disconnect + exact:
if peer not in addresses:
raise CommandError(f"Peer {peer} does not have a receptor address.")
# No 1-cycles # No 1-cycles
for collection in ('peers', 'disconnect', 'exact'): for collection in ('peers', 'disconnect', 'exact'):
@@ -53,9 +57,12 @@ class Command(BaseCommand):
# No 2-cycles # No 2-cycles
if options['peers'] or options['exact'] is not None: if options['peers'] or options['exact'] is not None:
peers = set(options['peers'] or options['exact']) peers = set(options['peers'] or options['exact'])
incoming = set(InstanceLink.objects.filter(target=nodes[options['source']]).values_list('source__hostname', flat=True)) if options['source'] in addresses:
incoming = set(InstanceLink.objects.filter(target=addresses[options['source']]).values_list('source__hostname', flat=True))
else:
incoming = set()
if peers & incoming: if peers & incoming:
warnings.warn(f"Source node {options['source']} should not link to nodes already peering to it: {peers & incoming}.") raise CommandError(f"Source node {options['source']} should not link to nodes already peering to it: {peers & incoming}.")
if options['peers']: if options['peers']:
missing_peers = set(options['peers']) - set(nodes) missing_peers = set(options['peers']) - set(nodes)
@@ -66,7 +73,7 @@ class Command(BaseCommand):
results = 0 results = 0
for target in options['peers']: for target in options['peers']:
_, created = InstanceLink.objects.update_or_create( _, created = InstanceLink.objects.update_or_create(
source=nodes[options['source']], target=nodes[target].receptor_addresses.get(), defaults={'link_state': InstanceLink.States.ESTABLISHED} source=nodes[options['source']], target=addresses[target], defaults={'link_state': InstanceLink.States.ESTABLISHED}
) )
if created: if created:
results += 1 results += 1
@@ -76,9 +83,9 @@ class Command(BaseCommand):
if options['disconnect']: if options['disconnect']:
results = 0 results = 0
for target in options['disconnect']: for target in options['disconnect']:
if target not in nodes: # Be permissive, the node might have already been de-registered. if target not in addresses: # Be permissive, the node might have already been de-registered.
continue continue
n, _ = InstanceLink.objects.filter(source=nodes[options['source']], target=nodes[target].receptor_addresses.get()).delete() n, _ = InstanceLink.objects.filter(source=nodes[options['source']], target=addresses[target]).delete()
results += n results += n
print(f"{results} peer links removed from the database.") print(f"{results} peer links removed from the database.")
@@ -87,11 +94,11 @@ class Command(BaseCommand):
additions = 0 additions = 0
with transaction.atomic(): with transaction.atomic():
peers = set(options['exact']) peers = set(options['exact'])
links = set(InstanceLink.objects.filter(source=nodes[options['source']]).values_list('target__hostname', flat=True)) links = set(InstanceLink.objects.filter(source=nodes[options['source']]).values_list('target__address', flat=True))
removals, _ = InstanceLink.objects.filter(source=nodes[options['source']], target__instance__hostname__in=links - peers).delete() removals, _ = InstanceLink.objects.filter(source=nodes[options['source']], target__instance__hostname__in=links - peers).delete()
for target in peers - links: for target in peers - links:
_, created = InstanceLink.objects.update_or_create( _, created = InstanceLink.objects.update_or_create(
source=nodes[options['source']], target=nodes[target].receptor_addresses.get(), defaults={'link_state': InstanceLink.States.ESTABLISHED} source=nodes[options['source']], target=addresses[target], defaults={'link_state': InstanceLink.States.ESTABLISHED}
) )
if created: if created:
additions += 1 additions += 1