desec-stack/api/desecapi/management/commands/sync-from-pdns.py

from django.core.management import BaseCommand, CommandError
from django.db import transaction

from desecapi import pdns
from desecapi.models import Domain, RRset, RR, RR_SET_TYPES_AUTOMATIC


class Command(BaseCommand):
    help = "Import authoritative data from pdns, making the local database consistent with pdns."

    def add_arguments(self, parser):
        parser.add_argument(
            "domain-name",
            nargs="*",
            help="Domain name to import. If omitted, will import all domains that are known locally.",
        )

    def handle(self, *args, **options):
        domains = Domain.objects.all()

        if options["domain-name"]:
            domains = domains.filter(name__in=options["domain-name"])
            domain_names = domains.values_list("name", flat=True)

            for domain_name in options["domain-name"]:
                if domain_name not in domain_names:
                    raise CommandError("{} is not a known domain".format(domain_name))

        for domain in domains:
            self.stdout.write("%s ..." % domain.name, ending="")
            try:
                self._sync_domain(domain)
                self.stdout.write(" synced")
            except Exception as e:
                self.stdout.write(" failed")
                msg = "Error while processing {}: {}".format(domain.name, e)
                raise CommandError(msg)

    @staticmethod
    @transaction.atomic
    def _sync_domain(domain):
        domain.rrset_set.all().delete()
        rrsets = []
        rrs = []
        for rrset_data in pdns.get_rrset_datas(domain):
            if rrset_data["type"] in RR_SET_TYPES_AUTOMATIC:
                continue
            records = rrset_data.pop("records")
            rrset = RRset(**rrset_data)
            rrsets.append(rrset)
            rrs.extend([RR(rrset=rrset, content=record) for record in records])
        RRset.objects.bulk_create(rrsets)
        RR.objects.bulk_create(rrs)