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