102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
import json
|
|
import os
|
|
import pwd
|
|
import shlex
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
|
|
|
|
def cli():
|
|
"""
|
|
Performs a backup of a Postgres database.
|
|
|
|
Parameters:
|
|
database: str — the name of the database to back up.
|
|
|
|
user: optional str — the name of the Linux user to use to connect to Postgres.
|
|
Sudo or SSH will be used to make this happen, if it's specified,
|
|
unless it's a local user that is already the current user.
|
|
|
|
host: optional str — if specified, the backup will be made using SSH
|
|
(unless this host is the same as the one named)
|
|
"""
|
|
request_info = json.load(sys.stdin)
|
|
assert isinstance(request_info, dict)
|
|
|
|
database_to_use = request_info["database"]
|
|
user_to_use = request_info.get("user")
|
|
host_to_use = request_info.get("host")
|
|
use_lz4 = request_info.get("use_lz4_for_ssh", True)
|
|
|
|
if host_to_use is not None:
|
|
hostname = subprocess.check_output("hostname").decode().strip()
|
|
if hostname == host_to_use:
|
|
host_to_use = None
|
|
|
|
# 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
|
|
|
|
dump_command = [
|
|
"pg_dump",
|
|
database_to_use
|
|
]
|
|
|
|
if host_to_use is not None:
|
|
if use_lz4:
|
|
# Add an LZ4 compressor on the remote side.
|
|
dump_command += ["|", "lz4", "--compress", "--stdout"]
|
|
|
|
# Then open an LZ4 decompressor on our side.
|
|
lz4_process = subprocess.Popen(
|
|
["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!
|
|
subprocess.check_call(
|
|
command,
|
|
stdin=subprocess.DEVNULL,
|
|
stdout=output_of_dump,
|
|
stderr=sys.stderr,
|
|
shell=False,
|
|
)
|
|
|
|
if lz4_process is not None:
|
|
# must close here, otherwise the decompressor never ends
|
|
lz4_process.stdin.close()
|
|
exit_code = lz4_process.wait()
|
|
if exit_code != 0:
|
|
raise ChildProcessError(f"lz4 not happy: {exit_code}")
|