diff --git a/polling_stations/apps/addressbase/management/commands/create_addressbase_lookup.py b/polling_stations/apps/addressbase/management/commands/create_addressbase_lookup.py index 68701b1c67..be82b36961 100644 --- a/polling_stations/apps/addressbase/management/commands/create_addressbase_lookup.py +++ b/polling_stations/apps/addressbase/management/commands/create_addressbase_lookup.py @@ -1,3 +1,4 @@ +from concurrent.futures import ThreadPoolExecutor from itertools import chain from pathlib import Path @@ -221,14 +222,14 @@ def create_single_lookup(self): ) self.stdout.write(f"Output written to: {destination_path}") - def create_outcode_lookup(self, outcode, cursor): + def create_outcode_lookup(self, outcode): destination_path = self.dst_dir / "outcodes" / f"{outcode.strip()}.csv" query_string = ( f"{self.single_lookup_query_string()} WHERE a.postcode LIKE '{outcode}%'" ) - self.stdout.write(f"creating lookup {destination_path.name}...") - with destination_path.open("w") as destination: + self.stdout.write(f"Processing {outcode.strip()}...") + with destination_path.open("w") as destination, connection.cursor() as cursor: cursor.copy_expert( f""" COPY ({query_string}) @@ -242,9 +243,9 @@ def create_outcode_lookup(self, outcode, cursor): def create_sharded_lookup(self): (self.dst_dir / "outcodes").mkdir(exist_ok=True, parents=True) outcodes = self.get_outcodes() - with connection.cursor() as cursor: + with ThreadPoolExecutor(max_workers=4) as executor: for outcode in outcodes: - self.create_outcode_lookup(outcode, cursor) + executor.submit(self.create_outcode_lookup, outcode) def get_outcodes(self): with connection.cursor() as cursor: