nominatim/db/utils.py

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