port database setup function to python

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

diff --git a/nominatim/db/utils.py b/nominatim/db/utils.py
index 1a104e5128c6e2ae44a27b72ff37838eba234970..575f301082e044e2d04a3a31fbcc69c8dfd2f8ac 100644 (file)
--- a/nominatim/db/utils.py
+++ b/nominatim/db/utils.py
@@ -3,12 +3,24 @@ Helper functions for handling DB accesses.
 """
 import subprocess
 import logging
+import gzip
 
 from .connection import get_pg_env
 from ..errors import UsageError
 
 LOG = logging.getLogger()
 
+def _pipe_to_proc(proc, fdesc):
+    chunk = fdesc.read(2048)
+    while chunk and proc.poll() is None:
+        try:
+            proc.stdin.write(chunk)
+        except BrokenPipeError as exc:
+            raise UsageError("Failed to execute SQL file.") from exc
+        chunk = fdesc.read(2048)
+
+    return len(chunk)
+
 def execute_file(dsn, fname, ignore_errors=False):
     """ Read an SQL file and run its contents against the given database
         using psql.
@@ -21,15 +33,15 @@ 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'))
 
-    with fname.open('rb') as fdesc:
-        chunk = fdesc.read(2048)
-        while chunk and proc.poll() is None:
-            proc.stdin.write(chunk)
-            chunk = fdesc.read(2048)
+    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()
 
     ret = proc.wait()
-    print(ret, chunk)
-    if ret != 0 or chunk:
+    if ret != 0 or remain > 0:
         raise UsageError("Failed to execute SQL file.")