Improve robustness of Postgres helper (hopefully)

This commit is contained in:
Olivier 'reivilibre' 2021-11-14 10:15:42 +00:00
parent e0e68ce76d
commit 0400bc83fa
1 changed files with 44 additions and 29 deletions

View File

@ -1,6 +1,8 @@
import json
import os
import pwd
import shlex
import shutil
import subprocess
import sys
@ -32,32 +34,20 @@ 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 = [
"pgdump",
database_to_use
]
if host_to_use is not None:
if use_lz4:
# Add an LZ4 compressor on the remote side.
command += ["|", "lz4", "--compress", "--stdout"]
dump_command += ["|", "lz4", "--compress", "--stdout"]
# Then open an LZ4 decompressor on our side.
lz4_process = subprocess.Popen(
@ -68,6 +58,31 @@ def cli():
)
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!
subprocess.check_call(