1import os
 2import subprocess
 3
 4
 5class PostgresBackupClient:
 6    def __init__(self, connection):
 7        self.connection = connection
 8
 9    def get_env(self):
10        settings_dict = self.connection.settings_dict
11        options = settings_dict.get("OPTIONS", {})
12        env = {}
13        if options.get("passfile"):
14            env["PGPASSFILE"] = str(options.get("passfile"))
15        if settings_dict.get("PASSWORD"):
16            env["PGPASSWORD"] = str(settings_dict.get("PASSWORD"))
17        if options.get("service"):
18            env["PGSERVICE"] = str(options.get("service"))
19        if options.get("sslmode"):
20            env["PGSSLMODE"] = str(options.get("sslmode"))
21        if options.get("sslrootcert"):
22            env["PGSSLROOTCERT"] = str(options.get("sslrootcert"))
23        if options.get("sslcert"):
24            env["PGSSLCERT"] = str(options.get("sslcert"))
25        if options.get("sslkey"):
26            env["PGSSLKEY"] = str(options.get("sslkey"))
27        return env
28
29    def create_backup(self, backup_path, *, pg_dump="pg_dump"):
30        settings_dict = self.connection.settings_dict
31
32        args = pg_dump.split()
33        options = settings_dict.get("OPTIONS", {})
34
35        host = settings_dict.get("HOST")
36        port = settings_dict.get("PORT")
37        dbname = settings_dict.get("NAME")
38        user = settings_dict.get("USER")
39        service = options.get("service")
40
41        if not dbname and not service:
42            # Connect to the default 'postgres' db.
43            dbname = "postgres"
44        if user:
45            args += ["-U", user]
46        if host:
47            args += ["-h", host]
48        if port:
49            args += ["-p", str(port)]
50
51        args += ["-Fc"]
52        # args += ["-f", backup_path]
53
54        if dbname:
55            args += [dbname]
56
57        # Using stdin/stdout let's us use executables from within a docker container too
58        args += ["|", "gzip", ">", str(backup_path)]
59
60        cmd = " ".join(args)
61
62        subprocess.run(
63            cmd, env={**os.environ, **self.get_env()}, check=True, shell=True
64        )
65
66    def restore_backup(self, backup_path, *, pg_restore="pg_restore"):
67        settings_dict = self.connection.settings_dict
68
69        args = pg_restore.split()
70        options = settings_dict.get("OPTIONS", {})
71
72        host = settings_dict.get("HOST")
73        port = settings_dict.get("PORT")
74        dbname = settings_dict.get("NAME")
75        user = settings_dict.get("USER")
76        service = options.get("service")
77
78        if not dbname and not service:
79            # Connect to the default 'postgres' db.
80            dbname = "postgres"
81        if user:
82            args += ["-U", user]
83        if host:
84            args += ["-h", host]
85        if port:
86            args += ["-p", str(port)]
87
88        args += ["--clean"]  # Drop existing tables
89        args += ["-d", dbname]
90
91        # Using stdin/stdout let's us use executables from within a docker container too
92        args = ["gunzip", "<", str(backup_path), "|"] + args
93
94        cmd = " ".join(args)
95
96        subprocess.run(
97            cmd, env={**os.environ, **self.get_env()}, check=True, shell=True
98        )