import subprocess
import logging
import gzip
+import io
-from .connection import get_pg_env
-from ..errors import UsageError
+from nominatim.db.connection import get_pg_env
+from nominatim.errors import UsageError
LOG = logging.getLogger()
cmd = ['psql']
if not ignore_errors:
cmd.extend(('-v', 'ON_ERROR_STOP=1'))
- proc = subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE)
-
if not LOG.isEnabledFor(logging.INFO):
- proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))
+ cmd.append('--quiet')
+ proc = subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE)
- if pre_code:
- proc.stdin.write((pre_code + ';').encode('utf-8'))
+ try:
+ if not LOG.isEnabledFor(logging.INFO):
+ proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))
- if fname.suffix == '.gz':
- with gzip.open(str(fname), 'rb') as fdesc:
- remain = _pipe_to_proc(proc, fdesc)
- else:
- with fname.open('rb') as fdesc:
- remain = _pipe_to_proc(proc, fdesc)
+ if pre_code:
+ proc.stdin.write((pre_code + ';').encode('utf-8'))
- if remain == 0 and post_code:
- proc.stdin.write((';' + post_code).encode('utf-8'))
+ if fname.suffix == '.gz':
+ with gzip.open(str(fname), 'rb') as fdesc:
+ remain = _pipe_to_proc(proc, fdesc)
+ else:
+ with fname.open('rb') as fdesc:
+ remain = _pipe_to_proc(proc, fdesc)
- proc.stdin.close()
+ if remain == 0 and post_code:
+ proc.stdin.write((';' + post_code).encode('utf-8'))
+ finally:
+ proc.stdin.close()
+ ret = proc.wait()
- ret = proc.wait()
if ret != 0 or remain > 0:
raise UsageError("Failed to execute SQL file.")
+
+
+# List of characters that need to be quoted for the copy command.
+_SQL_TRANSLATION = {ord(u'\\'): u'\\\\',
+ ord(u'\t'): u'\\t',
+ ord(u'\n'): u'\\n'}
+
+class CopyBuffer:
+ """ Data collector for the copy_from command.
+ """
+
+ def __init__(self):
+ self.buffer = io.StringIO()
+
+
+ def __enter__(self):
+ return self
+
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ if self.buffer is not None:
+ self.buffer.close()
+
+
+ def add(self, *data):
+ """ Add another row of data to the copy buffer.
+ """
+ first = True
+ for column in data:
+ if first:
+ first = False
+ else:
+ self.buffer.write('\t')
+ if column is None:
+ self.buffer.write('\\N')
+ else:
+ self.buffer.write(str(column).translate(_SQL_TRANSLATION))
+ self.buffer.write('\n')
+
+
+ def copy_out(self, cur, table, columns=None):
+ """ Copy all collected data into the given table.
+ """
+ if self.buffer.tell() > 0:
+ self.buffer.seek(0)
+ cur.copy_from(self.buffer, table, columns=columns)