feat: allow passing of postgres options (#556)

---------

Co-authored-by: Adam Ford <adam.ford@neosnetworks.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

README.rst

@@ -102,7 +102,11 @@ different server. You may also specify an explicit local file to backup from.
     -p PASSPHRASE, --passphrase=PASSPHRASE
                           Passphrase for decrypt file
     -z, --uncompress      Uncompress gzip data before restoring
-
+    -n SCHEMA, --schema SCHEMA
+                        Specify schema(s) to restore. Can be used multiple times.
+    -r, --no-drop         Don't clean (drop) the database. This only works with mongodb and postgresql.
+    --pg-options PG_OPTIONS
+                        Additional pg_restore options, e.g. '--if-exists --no-owner'. Use quotes.
 
 mediabackup
 -----------

dbbackup/db/postgresql.py

@@ -113,6 +113,7 @@ class PgDumpBinaryConnector(PgDumpConnector):
     single_transaction = True
     drop = True
     if_exists = False
+    pg_options = None
 
     def _create_dump(self):
         cmd = f"{self.dump_cmd} "
@@ -129,22 +130,64 @@ def _create_dump(self):
         stdout, _ = self.run_command(cmd, env=self.dump_env)
         return stdout
 
-    def _restore_dump(self, dump):
+    def _restore_dump(self, dump: str):
+        """
+        Restore a PostgreSQL dump using subprocess with argument list.
+
+        Assumes that restore_prefix, restore_cmd, pg_options, and restore_suffix
+        are either None, strings (single args), or lists of strings.
+
+        Builds the command as a list.
+        """
+
         dbname = create_postgres_uri(self)
-        cmd = f"{self.restore_cmd} {dbname}"
+        cmd = []
+
+        # Flatten optional values
+        if self.restore_prefix:
+            cmd.extend(
+                self.restore_prefix
+                if isinstance(self.restore_prefix, list)
+                else [self.restore_prefix]
+            )
+
+        if self.restore_cmd:
+            cmd.extend(
+                self.restore_cmd
+                if isinstance(self.restore_cmd, list)
+                else [self.restore_cmd]
+            )
+
+        if self.pg_options:
+            cmd.extend(
+                self.pg_options
+                if isinstance(self.pg_options, list)
+                else [self.pg_options]
+            )
+
+        cmd.extend([dbname])
 
         if self.single_transaction:
-            cmd += " --single-transaction"
+            cmd.extend(["--single-transaction"])
 
         if self.drop:
-            cmd += " --clean"
+            cmd.extend(["--clean"])
 
         if self.schemas:
-            cmd += " -n " + " -n ".join(self.schemas)
+            for schema in self.schemas:
+                cmd.extend(["-n", schema])
 
         if self.if_exists:
-            cmd += " --if-exists"
+            cmd.extend(["--if-exists"])
 
-        cmd = f"{self.restore_prefix} {cmd} {self.restore_suffix}"
-        stdout, stderr = self.run_command(cmd, stdin=dump, env=self.restore_env)
-        return stdout, stderr
+        if self.restore_suffix:
+            cmd.extend(
+                self.restore_suffix
+                if isinstance(self.restore_suffix, list)
+                else [self.restore_suffix]
+            )
+
+        cmd_str = " ".join(cmd)
+        stdout, _ = self.run_command(cmd_str, stdin=dump, env=self.dump_env)
+
+        return stdout

dbbackup/management/commands/dbrestore.py

@@ -16,6 +16,7 @@ class Command(BaseDbBackupCommand):
     help = "Restore a database backup from storage, encrypted and/or compressed."
     content_type = "db"
     no_drop = False
+    pg_options = ""
 
     option_list = BaseDbBackupCommand.option_list + (
         make_option("-d", "--database", help="Database to restore"),
@@ -60,6 +61,12 @@ class Command(BaseDbBackupCommand):
             default=False,
             help="Don't clean (drop) the database. This only works with mongodb and postgresql.",
         ),
+        make_option(
+            "--pg-options",
+            dest="pg_options",
+            default="",
+            help="Additional pg_restore options, e.g. '--if-exists --no-owner'. Use quotes.",
+        ),
     )
 
     def handle(self, *args, **options):
@@ -83,6 +90,7 @@ def handle(self, *args, **options):
             )
             self.storage = get_storage()
             self.no_drop = options.get("no_drop")
+            self.pg_options = options.get("pg_options", "")
             self.schemas = options.get("schema")
             self._restore_backup()
         except StorageError as err:
@@ -141,4 +149,5 @@ def _restore_backup(self):
         if self.schemas:
             self.connector.schemas = self.schemas
         self.connector.drop = not self.no_drop
+        self.connector.pg_options = self.pg_options
         self.connector.restore_dump(input_file)

dbbackup/tests/test_connectors/test_postgresql.py

@@ -223,6 +223,27 @@ def test_create_dump_if_exists(self, mock_run_command):
         self.connector.restore_dump(dump)
         self.assertIn(" --if-exists", mock_run_command.call_args[0][0])
 
+    def test_pg_options(self, mock_run_command):
+        dump = self.connector.create_dump()
+        self.connector.pg_options = "--foo"
+        self.connector.restore_dump(dump)
+        cmd_args = mock_run_command.call_args[0][0]
+        self.assertIn("--foo", cmd_args)
+
+    def test_restore_prefix(self, mock_run_command):
+        dump = self.connector.create_dump()
+        self.connector.restore_prefix = "foo"
+        self.connector.restore_dump(dump)
+        cmd_args = mock_run_command.call_args[0][0]
+        self.assertTrue(cmd_args.startswith("foo "))
+
+    def test_restore_suffix(self, mock_run_command):
+        dump = self.connector.create_dump()
+        self.connector.restore_suffix = "foo"
+        self.connector.restore_dump(dump)
+        cmd_args = mock_run_command.call_args[0][0]
+        self.assertTrue(cmd_args.endswith(" foo"))
+
     @patch(
         "dbbackup.db.postgresql.PgDumpBinaryConnector.run_command",
         return_value=(BytesIO(), BytesIO()),

docs/changelog.rst

@@ -4,6 +4,7 @@ Changelog
 Unreleased
 ----------
 
+* Add generic `--pg-options` to pass custom options to postgres.
 * Add option `--if-exists` for pg_dump command
 * Empty string as HOST for postgres unix domain socket connection is now supported.
 * Support Python 3.13 and Django 5.2