1 # SPDX-License-Identifier: GPL-2.0-only
3 # This file is part of Nominatim. (https://nominatim.org)
5 # Copyright (C) 2022 by the Nominatim developer community.
6 # For a full list of authors see the git log.
8 Helper functions for handling DB accesses.
10 from typing import IO, Optional, Union, Any, Iterable
15 from pathlib import Path
17 from nominatim.db.connection import get_pg_env, Cursor
18 from nominatim.errors import UsageError
20 LOG = logging.getLogger()
22 def _pipe_to_proc(proc: 'subprocess.Popen[bytes]',
23 fdesc: Union[IO[bytes], gzip.GzipFile]) -> int:
24 assert proc.stdin is not None
25 chunk = fdesc.read(2048)
26 while chunk and proc.poll() is None:
28 proc.stdin.write(chunk)
29 except BrokenPipeError as exc:
30 raise UsageError("Failed to execute SQL file.") from exc
31 chunk = fdesc.read(2048)
35 def execute_file(dsn: str, fname: Path,
36 ignore_errors: bool = False,
37 pre_code: Optional[str] = None,
38 post_code: Optional[str] = None) -> None:
39 """ Read an SQL file and run its contents against the given database
40 using psql. Use `pre_code` and `post_code` to run extra commands
41 before or after executing the file. The commands are run within the
42 same session, so they may be used to wrap the file execution in a
47 cmd.extend(('-v', 'ON_ERROR_STOP=1'))
48 if not LOG.isEnabledFor(logging.INFO):
51 with subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE) as proc:
52 assert proc.stdin is not None
54 if not LOG.isEnabledFor(logging.INFO):
55 proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))
58 proc.stdin.write((pre_code + ';').encode('utf-8'))
60 if fname.suffix == '.gz':
61 with gzip.open(str(fname), 'rb') as fdesc:
62 remain = _pipe_to_proc(proc, fdesc)
64 with fname.open('rb') as fdesc:
65 remain = _pipe_to_proc(proc, fdesc)
67 if remain == 0 and post_code:
68 proc.stdin.write((';' + post_code).encode('utf-8'))
73 if ret != 0 or remain > 0:
74 raise UsageError("Failed to execute SQL file.")
77 # List of characters that need to be quoted for the copy command.
78 _SQL_TRANSLATION = {ord('\\'): '\\\\',
84 """ Data collector for the copy_from command.
87 def __init__(self) -> None:
88 self.buffer = io.StringIO()
91 def __enter__(self) -> 'CopyBuffer':
95 def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
96 if self.buffer is not None:
100 def add(self, *data: Any) -> None:
101 """ Add another row of data to the copy buffer.
108 self.buffer.write('\t')
110 self.buffer.write('\\N')
112 self.buffer.write(str(column).translate(_SQL_TRANSLATION))
113 self.buffer.write('\n')
116 def copy_out(self, cur: Cursor, table: str, columns: Optional[Iterable[str]] = None) -> None:
117 """ Copy all collected data into the given table.
119 if self.buffer.tell() > 0:
121 cur.copy_from(self.buffer, table, columns=columns) # type: ignore[no-untyped-call]