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 def _pipe_to_proc(proc: 'subprocess.Popen[bytes]',
  22                   fdesc: Union[IO[bytes], gzip.GzipFile]) -> int:
  23     assert proc.stdin is not None
  24     chunk = fdesc.read(2048)
  25     while chunk and proc.poll() is None:
  26         try:
  27             proc.stdin.write(chunk)
  28         except BrokenPipeError as exc:
  29             raise UsageError("Failed to execute SQL file.") from exc
  30         chunk = fdesc.read(2048)
  31
  32     return len(chunk)
  33
  34 def execute_file(dsn: str, fname: Path,
  35                  ignore_errors: bool = False,
  36                  pre_code: Optional[str] = None,
  37                  post_code: Optional[str] = None) -> None:
  38     """ Read an SQL file and run its contents against the given database
  39         using psql. Use `pre_code` and `post_code` to run extra commands
  40         before or after executing the file. The commands are run within the
  41         same session, so they may be used to wrap the file execution in a
  42         transaction.
  43     """
  44     cmd = ['psql']
  45     if not ignore_errors:
  46         cmd.extend(('-v', 'ON_ERROR_STOP=1'))
  47     if not LOG.isEnabledFor(logging.INFO):
  48         cmd.append('--quiet')
  49
  50     with subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE) as proc:
  51         assert proc.stdin is not None
  52         try:
  53             if not LOG.isEnabledFor(logging.INFO):
  54                 proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))
  55
  56             if pre_code:
  57                 proc.stdin.write((pre_code + ';').encode('utf-8'))
  58
  59             if fname.suffix == '.gz':
  60                 with gzip.open(str(fname), 'rb') as fdesc:
  61                     remain = _pipe_to_proc(proc, fdesc)
  62             else:
  63                 with fname.open('rb') as fdesc:
  64                     remain = _pipe_to_proc(proc, fdesc)
  65
  66             if remain == 0 and post_code:
  67                 proc.stdin.write((';' + post_code).encode('utf-8'))
  68         finally:
  69             proc.stdin.close()
  70             ret = proc.wait()
  71
  72     if ret != 0 or remain > 0:
  73         raise UsageError("Failed to execute SQL file.")