]> git.openstreetmap.org Git - nominatim.git/blob - nominatim/db/utils.py
move import-data option to native python
[nominatim.git] / nominatim / db / utils.py
1 """
2 Helper functions for handling DB accesses.
3 """
4 import subprocess
5 import logging
6 import gzip
7
8 from .connection import get_pg_env
9 from ..errors import UsageError
10
11 LOG = logging.getLogger()
12
13 def _pipe_to_proc(proc, fdesc):
14     chunk = fdesc.read(2048)
15     while chunk and proc.poll() is None:
16         try:
17             proc.stdin.write(chunk)
18         except BrokenPipeError as exc:
19             raise UsageError("Failed to execute SQL file.") from exc
20         chunk = fdesc.read(2048)
21
22     return len(chunk)
23
24 def execute_file(dsn, fname, ignore_errors=False):
25     """ Read an SQL file and run its contents against the given database
26         using psql.
27     """
28     cmd = ['psql']
29     if not ignore_errors:
30         cmd.extend(('-v', 'ON_ERROR_STOP=1'))
31     proc = subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE)
32
33     if not LOG.isEnabledFor(logging.INFO):
34         proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))
35
36     if fname.suffix == '.gz':
37         with gzip.open(str(fname), 'rb') as fdesc:
38             remain = _pipe_to_proc(proc, fdesc)
39     else:
40         with fname.open('rb') as fdesc:
41             remain = _pipe_to_proc(proc, fdesc)
42
43     proc.stdin.close()
44
45     ret = proc.wait()
46     if ret != 0 or remain > 0:
47         raise UsageError("Failed to execute SQL file.")