nominatim/db/utils.py

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