src/nominatim_db/db/utils.py

   1 # SPDX-License-Identifier: GPL-3.0-or-later
   2 #
   3 # This file is part of Nominatim. (https://nominatim.org)
   4 #
   5 # Copyright (C) 2024 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 from typing import IO, Optional, Union
  11 import subprocess
  12 import logging
  13 import gzip
  14 from pathlib import Path
  15
  16 from .connection import get_pg_env
  17 from ..errors import UsageError
  18
  19 LOG = logging.getLogger()
  20
  21
  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:
  27         try:
  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)
  32
  33     return len(chunk)
  34
  35
  36 def execute_file(dsn: str, fname: Path,
  37                  ignore_errors: bool = False,
  38                  pre_code: Optional[str] = None,
  39                  post_code: Optional[str] = None) -> None:
  40     """ Read an SQL file and run its contents against the given database
  41         using psql. Use `pre_code` and `post_code` to run extra commands
  42         before or after executing the file. The commands are run within the
  43         same session, so they may be used to wrap the file execution in a
  44         transaction.
  45     """
  46     cmd = ['psql']
  47     if not ignore_errors:
  48         cmd.extend(('-v', 'ON_ERROR_STOP=1'))
  49     if not LOG.isEnabledFor(logging.INFO):
  50         cmd.append('--quiet')
  51
  52     with subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE) as proc:
  53         assert proc.stdin is not None
  54         try:
  55             if not LOG.isEnabledFor(logging.INFO):
  56                 proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))
  57
  58             if pre_code:
  59                 proc.stdin.write((pre_code + ';').encode('utf-8'))
  60
  61             if fname.suffix == '.gz':
  62                 with gzip.open(str(fname), 'rb') as fdesc:
  63                     remain = _pipe_to_proc(proc, fdesc)
  64             else:
  65                 with fname.open('rb') as fdesc:
  66                     remain = _pipe_to_proc(proc, fdesc)
  67
  68             if remain == 0 and post_code:
  69                 proc.stdin.write((';' + post_code).encode('utf-8'))
  70         finally:
  71             proc.stdin.close()
  72             ret = proc.wait()
  73
  74     if ret != 0 or remain > 0:
  75         raise UsageError("Failed to execute SQL file.")