port wikipedia importance functions to python

[nominatim.git] / nominatim / db / utils.py

diff --git a/nominatim/db/utils.py b/nominatim/db/utils.py
index 575f301082e044e2d04a3a31fbcc69c8dfd2f8ac..0490bbc869fe3f19a4bc87eac2fbed73d4b54d6f 100644 (file)
--- a/nominatim/db/utils.py
+++ b/nominatim/db/utils.py
@@ -21,9 +21,12 @@ def _pipe_to_proc(proc, fdesc):
 
     return len(chunk)
 
-def execute_file(dsn, fname, ignore_errors=False):
+def execute_file(dsn, fname, ignore_errors=False, pre_code=None, post_code=None):
     """ Read an SQL file and run its contents against the given database
-        using psql.
+        using psql. Use `pre_code` and `post_code` to run extra commands
+        before or after executing the file. The commands are run within the
+        same session, so they may be used to wrap the file execution in a
+        transaction.
     """
     cmd = ['psql']
     if not ignore_errors:
@@ -33,6 +36,9 @@ def execute_file(dsn, fname, ignore_errors=False):
     if not LOG.isEnabledFor(logging.INFO):
         proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))
 
+    if pre_code:
+        proc.stdin.write((pre_code + ';').encode('utf-8'))
+
     if fname.suffix == '.gz':
         with gzip.open(str(fname), 'rb') as fdesc:
             remain = _pipe_to_proc(proc, fdesc)
@@ -40,6 +46,9 @@ def execute_file(dsn, fname, ignore_errors=False):
         with fname.open('rb') as fdesc:
             remain = _pipe_to_proc(proc, fdesc)
 
+    if remain == 0 and post_code:
+        proc.stdin.write((';' + post_code).encode('utf-8'))
+
     proc.stdin.close()
 
     ret = proc.wait()