]> git.openstreetmap.org Git - nominatim.git/blob - nominatim/db/utils.py
Merge pull request #2381 from lonvia/reorganise-abbreviations
[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 import io
8
9 from nominatim.db.connection import get_pg_env
10 from nominatim.errors import UsageError
11
12 LOG = logging.getLogger()
13
14 def _pipe_to_proc(proc, fdesc):
15     chunk = fdesc.read(2048)
16     while chunk and proc.poll() is None:
17         try:
18             proc.stdin.write(chunk)
19         except BrokenPipeError as exc:
20             raise UsageError("Failed to execute SQL file.") from exc
21         chunk = fdesc.read(2048)
22
23     return len(chunk)
24
25 def execute_file(dsn, fname, ignore_errors=False, pre_code=None, post_code=None):
26     """ Read an SQL file and run its contents against the given database
27         using psql. Use `pre_code` and `post_code` to run extra commands
28         before or after executing the file. The commands are run within the
29         same session, so they may be used to wrap the file execution in a
30         transaction.
31     """
32     cmd = ['psql']
33     if not ignore_errors:
34         cmd.extend(('-v', 'ON_ERROR_STOP=1'))
35     if not LOG.isEnabledFor(logging.INFO):
36         cmd.append('--quiet')
37     proc = subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE)
38
39     try:
40         if not LOG.isEnabledFor(logging.INFO):
41             proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))
42
43         if pre_code:
44             proc.stdin.write((pre_code + ';').encode('utf-8'))
45
46         if fname.suffix == '.gz':
47             with gzip.open(str(fname), 'rb') as fdesc:
48                 remain = _pipe_to_proc(proc, fdesc)
49         else:
50             with fname.open('rb') as fdesc:
51                 remain = _pipe_to_proc(proc, fdesc)
52
53         if remain == 0 and post_code:
54             proc.stdin.write((';' + post_code).encode('utf-8'))
55     finally:
56         proc.stdin.close()
57         ret = proc.wait()
58
59     if ret != 0 or remain > 0:
60         raise UsageError("Failed to execute SQL file.")
61
62
63 # List of characters that need to be quoted for the copy command.
64 _SQL_TRANSLATION = {ord(u'\\') : u'\\\\',
65                     ord(u'\t') : u'\\t',
66                     ord(u'\n') : u'\\n'}
67
68 class CopyBuffer:
69     """ Data collector for the copy_from command.
70     """
71
72     def __init__(self):
73         self.buffer = io.StringIO()
74
75
76     def __enter__(self):
77         return self
78
79
80     def __exit__(self, exc_type, exc_value, traceback):
81         if self.buffer is not None:
82             self.buffer.close()
83
84
85     def add(self, *data):
86         """ Add another row of data to the copy buffer.
87         """
88         first = True
89         for column in data:
90             if first:
91                 first = False
92             else:
93                 self.buffer.write('\t')
94             if column is None:
95                 self.buffer.write('\\N')
96             else:
97                 self.buffer.write(str(column).translate(_SQL_TRANSLATION))
98         self.buffer.write('\n')
99
100
101     def copy_out(self, cur, table, columns=None):
102         """ Copy all collected data into the given table.
103         """
104         if self.buffer.tell() > 0:
105             self.buffer.seek(0)
106             cur.copy_from(self.buffer, table, columns=columns)