nominatim/db/utils.py

   1 """
   2 Helper functions for handling DB accesses.
   3 """
   4 import subprocess
   5 import logging
   6 import gzip
   7
   8 from .connection import get_pg_env
   9 from ..errors import UsageError
  10
  11 LOG = logging.getLogger()
  12
  13 def _pipe_to_proc(proc, fdesc):
  14     chunk = fdesc.read(2048)
  15     while chunk and proc.poll() is None:
  16         try:
  17             proc.stdin.write(chunk)
  18         except BrokenPipeError as exc:
  19             raise UsageError("Failed to execute SQL file.") from exc
  20         chunk = fdesc.read(2048)
  21
  22     return len(chunk)
  23
  24 def execute_file(dsn, fname, ignore_errors=False, pre_code=None, post_code=None):
  25     """ Read an SQL file and run its contents against the given database
  26         using psql. Use `pre_code` and `post_code` to run extra commands
  27         before or after executing the file. The commands are run within the
  28         same session, so they may be used to wrap the file execution in a
  29         transaction.
  30     """
  31     cmd = ['psql']
  32     if not ignore_errors:
  33         cmd.extend(('-v', 'ON_ERROR_STOP=1'))
  34     if not LOG.isEnabledFor(logging.INFO):
  35         cmd.append('--quiet')
  36     proc = subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE)
  37
  38     try:
  39         if not LOG.isEnabledFor(logging.INFO):
  40             proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))
  41
  42         if pre_code:
  43             proc.stdin.write((pre_code + ';').encode('utf-8'))
  44
  45         if fname.suffix == '.gz':
  46             with gzip.open(str(fname), 'rb') as fdesc:
  47                 remain = _pipe_to_proc(proc, fdesc)
  48         else:
  49             with fname.open('rb') as fdesc:
  50                 remain = _pipe_to_proc(proc, fdesc)
  51
  52         if remain == 0 and post_code:
  53             proc.stdin.write((';' + post_code).encode('utf-8'))
  54     finally:
  55         proc.stdin.close()
  56         ret = proc.wait()
  57
  58     if ret != 0 or remain > 0:
  59         raise UsageError("Failed to execute SQL file.")