Improve robustness of Postgres helper (hopefully)

datman-helper-postgres/datman_helper_postgres/backup.py

@@ -1,6 +1,8 @@
 import json
 import os
 import pwd
+import shlex
+import shutil
 import subprocess
 import sys
 
@@ -32,41 +34,54 @@ def cli():
         if hostname == host_to_use:
             host_to_use = None
 
-    command = []
-
-    if host_to_use is not None:
-        command.append("ssh")
-        if user_to_use is not None:
-            command.append(f"{user_to_use}@{host_to_use}")
-        else:
-            command.append(f"{host_to_use}")
-    elif user_to_use is not None:
-        current_username = pwd.getpwuid(os.getuid()).pw_name
-        if current_username != user_to_use:
-            command.append("sudo")
-            command.append("-u")
-            command.append(user_to_use)
-
-    command.append("pg_dump")
-    command.append(database_to_use)
-
     # Where the output of the dump command should go.
     output_of_dump = sys.stdout
     # The process (if any) that is our LZ4 decompressor.
     lz4_process = None
 
-    if use_lz4 and host_to_use is not None:
+    dump_command = [
-        # Add an LZ4 compressor on the remote side.
+        "pgdump",
-        command += ["|", "lz4", "--compress", "--stdout"]
+        database_to_use
+    ]
 
-        # Then open an LZ4 decompressor on our side.
+    if host_to_use is not None:
-        lz4_process = subprocess.Popen(
+        if use_lz4:
-            ["lz4", "--decompress", "--stdout"],
+            # Add an LZ4 compressor on the remote side.
-            stdin=subprocess.PIPE,
+            dump_command += ["|", "lz4", "--compress", "--stdout"]
-            stdout=sys.stdout,
+
-            stderr=sys.stderr,
+            # Then open an LZ4 decompressor on our side.
-        )
+            lz4_process = subprocess.Popen(
-        output_of_dump = lz4_process.stdin
+                ["lz4", "--decompress", "--stdout"],
+                stdin=subprocess.PIPE,
+                stdout=sys.stdout,
+                stderr=sys.stderr,
+            )
+            output_of_dump = lz4_process.stdin
+
+        # We want to open a bash on the other side with pipefail
+        # so that an issue with the backup gets noticed
+        # (rather than lz4 covering it).
+        command = [
+            "ssh",
+            f"{user_to_use}@{host_to_use}" if user_to_use is not None else f"{host_to_use}",
+            "bash",
+            "-o",
+            "pipefail",
+            "-c",
+            shlex.quote(" ".join(dump_command))
+        ]
+    elif user_to_use is not None:
+        current_username = pwd.getpwuid(os.getuid()).pw_name
+        if current_username != user_to_use:
+            command = [
+                "sudo",
+                "-u",
+                user_to_use
+            ] + dump_command
+        else:
+            command = dump_command
+    else:
+        command = dump_command
 
     # we MUST disable shell here otherwise the local side will do both
     # the compression and decompression which would be silly!