From 08ffd39674873c1895adb2c5742a6d12032607f6 Mon Sep 17 00:00:00 2001 From: Jonas Braathen Date: Tue, 3 Nov 2026 03:52:41 +0100 Subject: [PATCH] fix: accept filename argument in import_valiases --- dpdapi/management/commands/import_valiases.py | 72 +++++++++---------- 1 file changed, 36 insertions(+), 36 deletions(-) diff --git a/dpdapi/management/commands/import_valiases.py b/dpdapi/management/commands/import_valiases.py index f5a3630..d1fa82d 100644 --- a/dpdapi/management/commands/import_valiases.py +++ b/dpdapi/management/commands/import_valiases.py @@ -1,60 +1,60 @@ +from pathlib import Path + from django.core.management.base import BaseCommand -from django.db.utils import IntegrityError -import os +from django.db import IntegrityError, transaction from dpdapi.models import Alias, Domain class Command(BaseCommand): - help = 'Import aliases for a single domain at a time, which is the domain of the first source address.' + help = "Import aliases for a single domain at a time, which is the domain of the first source address." - def handle(self, *args, **options): - if len(args) == 1: - if os.path.exists(args[0]): - self.import_file(args[0]) - else: - self.stdout.write("{} does not exist".format(args[0])) - else: + def add_arguments(self, parser): + parser.add_argument("filename", nargs="?") + + def handle(self, *args, filename=None, **options): + if filename is None: self.stdout.write("Missing argument") + elif not Path(filename).exists(): + self.stdout.write(f"{filename} does not exist") + else: + self.import_file(filename) def import_file(self, filename): new_aliases = [] - counts = {'imported': 0, 'existing': 0} - # Open file - with open(filename) as f: - # Parse - for l in f: - if l.startswith('*: '): + counts = {"imported": 0, "existing": 0} + with Path(filename).open() as f: + for line in f: + if line.startswith("*: "): continue # Skip catch all - source, destinations = l.strip().split(': ') - for destination in destinations.split(','): - if destination.startswith('"|') or '-admin@' in destination: + source, destinations = line.strip().split(": ") + for destination in destinations.split(","): + if destination.startswith('"|') or "-admin@" in destination: continue # Skip mailman-aliases new_aliases.append(Alias(source=source, destination=destination)) - if len(new_aliases) == 0: - self.stdout.write('No aliases found') + if not new_aliases: + self.stdout.write("No aliases found") return - # Lookup domain - domain_name = new_aliases[0].source.split('@')[1] # Assumes that every alias is in the same domain - d, created = Domain.objects.get_or_create(name=domain_name) + # Assumes that every alias is in the same domain + domain_name = new_aliases[0].source.split("@")[1] + domain, created = Domain.objects.get_or_create(name=domain_name) if created: - self.stdout.write('Created domain {}.'.format(d.name)) + self.stdout.write(f"Created domain {domain.name}.") # Try to create, log if exists. Slow but works - for a in new_aliases: - a.domain = d + for alias in new_aliases: + alias.domain = domain try: - a.save() - counts['imported'] += 1 # Woop! + with transaction.atomic(): + alias.save() + counts["imported"] += 1 except IntegrityError: - counts['existing'] += 1 # Exists - self.stdout.write('Could not import \'{}={}\', already exists.'.format(a.source, a.destination)) + counts["existing"] += 1 + self.stdout.write(f"Could not import '{alias.source}={alias.destination}', already exists.") - # Print num imported, existing, total - self.stdout.write('Imported {}, skipped {} existing, total {}'.format( - counts['imported'], - counts['existing'], - len(new_aliases))) + self.stdout.write( + f"Imported {counts['imported']}, skipped {counts['existing']} existing, total {len(new_aliases)}" + )