X-Git-Url: https://git.openstreetmap.org./nominatim.git/blobdiff_plain/b93ec2522e25102e8a36ca78423a53b4579c68c2..f6e894a53af83a69f553555cb4a6248d57a58391:/nominatim/db/utils.py diff --git a/nominatim/db/utils.py b/nominatim/db/utils.py index 1a104e51..575f3010 100644 --- a/nominatim/db/utils.py +++ b/nominatim/db/utils.py @@ -3,12 +3,24 @@ Helper functions for handling DB accesses. """ import subprocess import logging +import gzip from .connection import get_pg_env from ..errors import UsageError LOG = logging.getLogger() +def _pipe_to_proc(proc, fdesc): + chunk = fdesc.read(2048) + while chunk and proc.poll() is None: + try: + proc.stdin.write(chunk) + except BrokenPipeError as exc: + raise UsageError("Failed to execute SQL file.") from exc + chunk = fdesc.read(2048) + + return len(chunk) + def execute_file(dsn, fname, ignore_errors=False): """ Read an SQL file and run its contents against the given database using psql. @@ -21,15 +33,15 @@ def execute_file(dsn, fname, ignore_errors=False): if not LOG.isEnabledFor(logging.INFO): proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8')) - with fname.open('rb') as fdesc: - chunk = fdesc.read(2048) - while chunk and proc.poll() is None: - proc.stdin.write(chunk) - chunk = fdesc.read(2048) + if fname.suffix == '.gz': + with gzip.open(str(fname), 'rb') as fdesc: + remain = _pipe_to_proc(proc, fdesc) + else: + with fname.open('rb') as fdesc: + remain = _pipe_to_proc(proc, fdesc) proc.stdin.close() ret = proc.wait() - print(ret, chunk) - if ret != 0 or chunk: + if ret != 0 or remain > 0: raise UsageError("Failed to execute SQL file.")