]> git.openstreetmap.org Git - nominatim.git/blob - nominatim/db/utils.py
Merge remote-tracking branch 'upstream/master'
[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 nominatim.db.connection import get_pg_env
9 from nominatim.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, pre_code=None, post_code=None):
25     """ Read an SQL file and run its contents against the given database
26         using psql. Use `pre_code` and `post_code` to run extra commands
27         before or after executing the file. The commands are run within the
28         same session, so they may be used to wrap the file execution in a
29         transaction.
30     """
31     cmd = ['psql']
32     if not ignore_errors:
33         cmd.extend(('-v', 'ON_ERROR_STOP=1'))
34     if not LOG.isEnabledFor(logging.INFO):
35         cmd.append('--quiet')
36     proc = subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE)
37
38     try:
39         if not LOG.isEnabledFor(logging.INFO):
40             proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))
41
42         if pre_code:
43             proc.stdin.write((pre_code + ';').encode('utf-8'))
44
45         if fname.suffix == '.gz':
46             with gzip.open(str(fname), 'rb') as fdesc:
47                 remain = _pipe_to_proc(proc, fdesc)
48         else:
49             with fname.open('rb') as fdesc:
50                 remain = _pipe_to_proc(proc, fdesc)
51
52         if remain == 0 and post_code:
53             proc.stdin.write((';' + post_code).encode('utf-8'))
54     finally:
55         proc.stdin.close()
56         ret = proc.wait()
57
58     if ret != 0 or remain > 0:
59         raise UsageError("Failed to execute SQL file.")