Add postgresql schema support.

dbbackup/db/postgresql.py

@@ -1,4 +1,5 @@
 import logging
+from typing import List, Optional
 from urllib.parse import quote
 
 from .base import BaseCommandDBConnector
@@ -37,16 +38,23 @@ class PgDumpConnector(BaseCommandDBConnector):
     restore_cmd = "psql"
     single_transaction = True
     drop = True
+    schemas: Optional[List[str]] = []
 
     def _create_dump(self):
         cmd = f"{self.dump_cmd} "
         cmd = cmd + create_postgres_uri(self)
 
         for table in self.exclude:
             cmd += f" --exclude-table-data={table}"
+
         if self.drop:
             cmd += " --clean"
 
+        if self.schemas:
+            # First schema is not prefixed with -n
+            # when using join function so add it manually.
+            cmd += " -n " + " -n ".join(self.schemas)
+
         cmd = f"{self.dump_prefix} {cmd} {self.dump_suffix}"
         stdout, stderr = self.run_command(cmd, env=self.dump_env)
         return stdout
@@ -57,8 +65,13 @@ def _restore_dump(self, dump):
 
         # without this, psql terminates with an exit value of 0 regardless of errors
         cmd += " --set ON_ERROR_STOP=on"
+
+        if self.schemas:
+            cmd += " -n " + " -n ".join(self.schemas)
+
         if self.single_transaction:
             cmd += " --single-transaction"
+
         cmd += " {}".format(self.settings["NAME"])
         cmd = f"{self.restore_prefix} {cmd} {self.restore_suffix}"
         stdout, stderr = self.run_command(cmd, stdin=dump, env=self.restore_env)
@@ -77,10 +90,13 @@ def _enable_postgis(self):
         cmd = f'{self.psql_cmd} -c "CREATE EXTENSION IF NOT EXISTS postgis;"'
         cmd += " --username={}".format(self.settings["ADMIN_USER"])
         cmd += " --no-password"
+
         if self.settings.get("HOST"):
             cmd += " --host={}".format(self.settings["HOST"])
+
         if self.settings.get("PORT"):
             cmd += " --port={}".format(self.settings["PORT"])
+
         return self.run_command(cmd)
 
     def _restore_dump(self, dump):
@@ -108,8 +124,12 @@ def _create_dump(self):
         cmd += " --format=custom"
         for table in self.exclude:
             cmd += f" --exclude-table-data={table}"
+
+        if self.schemas:
+            cmd += " -n " + " -n ".join(self.schemas)
+
         cmd = f"{self.dump_prefix} {cmd} {self.dump_suffix}"
-        stdout, stderr = self.run_command(cmd, env=self.dump_env)
+        stdout, _ = self.run_command(cmd, env=self.dump_env)
         return stdout
 
     def _restore_dump(self, dump):
@@ -118,8 +138,13 @@ def _restore_dump(self, dump):
 
         if self.single_transaction:
             cmd += " --single-transaction"
+
         if self.drop:
             cmd += " --clean"
+
+        if self.schemas:
+            cmd += " -n " + " -n ".join(self.schemas)
+
         cmd = f"{self.restore_prefix} {cmd} {self.restore_suffix}"
         stdout, stderr = self.run_command(cmd, stdin=dump, env=self.restore_env)
         return stdout, stderr

dbbackup/management/commands/dbbackup.py

@@ -11,7 +11,7 @@
 
 
 class Command(BaseDbBackupCommand):
-    help = "Backup a database, encrypt and/or compress and write to " "storage." ""
+    help = "Backup a database, encrypt and/or compress."
     content_type = "db"
 
     option_list = BaseDbBackupCommand.option_list + (
@@ -60,6 +60,13 @@ class Command(BaseDbBackupCommand):
         make_option(
             "-x", "--exclude-tables", default=None, help="Exclude tables from backup"
         ),
+        make_option(
+            "-n",
+            "--schema",
+            action="append",
+            default=[],
+            help="Specify schema(s) to backup. Can be used multiple times.",
+        ),
     )
 
     @utils.email_uncaught_exception
@@ -78,6 +85,7 @@ def handle(self, **options):
         self.path = options.get("output_path")
         self.exclude_tables = options.get("exclude_tables")
         self.storage = get_storage()
+        self.schemas = options.get("schema")
 
         self.database = options.get("database") or ""
 
@@ -103,22 +111,32 @@ def _save_new_backup(self, database):
         Save a new backup file.
         """
         self.logger.info("Backing Up Database: %s", database["NAME"])
-        # Get backup and name
+        # Get backup, schema and name
         filename = self.connector.generate_filename(self.servername)
+
+        if self.schemas:
+            self.connector.schemas = self.schemas
+
         outputfile = self.connector.create_dump()
+
         # Apply trans
         if self.compress:
             compressed_file, filename = utils.compress_file(outputfile, filename)
             outputfile = compressed_file
+
         if self.encrypt:
             encrypted_file, filename = utils.encrypt_file(outputfile, filename)
             outputfile = encrypted_file
+
         # Set file name
         filename = self.filename or filename
-        self.logger.debug("Backup size: %s", utils.handle_size(outputfile))
+        self.logger.info("Backup tempfile created: %s", utils.handle_size(outputfile))
+
         # Store backup
         outputfile.seek(0)
+
         if self.path is None:
             self.write_to_storage(outputfile, filename)
+
         else:
             self.write_local_file(outputfile, self.path)

dbbackup/management/commands/dbrestore.py

@@ -13,8 +13,7 @@
 
 
 class Command(BaseDbBackupCommand):
-    help = """Restore a database backup from storage, encrypted and/or
-    compressed."""
+    help = "Restore a database backup from storage, encrypted and/or compressed."
     content_type = "db"
 
     option_list = BaseDbBackupCommand.option_list + (
@@ -46,6 +45,13 @@
             default=False,
             help="Uncompress gzip data before restoring",
         ),
+        make_option(
+            "-n",
+            "--schema",
+            action="append",
+            default=[],
+            help="Specify schema(s) to restore. Can be used multiple times.",
+        ),
     )
 
     def handle(self, *args, **options):
@@ -68,6 +74,7 @@
                 self.input_database_name
             )
             self.storage = get_storage()
+            self.schemas = options.get("schema")
             self._restore_backup()
         except StorageError as err:
             raise CommandError(err) from err
@@ -91,11 +98,16 @@
         input_filename, input_file = self._get_backup_file(
             database=self.input_database_name, servername=self.servername
         )
+
         self.logger.info(
             "Restoring backup for database '%s' and server '%s'",
             self.database_name,
             self.servername,
         )
+
+        if self.schemas:
+            self.logger.info(f"Restoring schemas: {self.schemas}")
+
         self.logger.info(f"Restoring: {input_filename}")
 
         if self.decrypt:
@@ -117,4 +129,8 @@
 
         input_file.seek(0)
         self.connector = get_connector(self.database_name)
+
+        if self.schemas:
+            self.connector.schemas = self.schemas
+
         self.connector.restore_dump(input_file)

dbbackup/tests/commands/test_dbbackup.py

@@ -3,9 +3,9 @@
 """
 
 import os
+from unittest.mock import patch
 
 from django.test import TestCase
-from mock import patch
 
 from dbbackup.db.base import get_connector
 from dbbackup.management.commands.dbbackup import Command as DbbackupCommand
@@ -27,6 +27,7 @@ def setUp(self):
         self.command.stdout = DEV_NULL
         self.command.filename = None
         self.command.path = None
+        self.command.schemas = []
 
     def tearDown(self):
         clean_gpg_keys()
@@ -50,6 +51,12 @@ def test_path(self):
         # tearDown
         os.remove(self.command.path)
 
+    def test_schema(self):
+        self.command.schemas = ["public"]
+        result = self.command._save_new_backup(TEST_DATABASE)
+
+        self.assertIsNone(result)
+
     @patch("dbbackup.settings.DATABASES", ["db-from-settings"])
     def test_get_database_keys(self):
         with self.subTest("use --database from CLI"):
@@ -76,6 +83,7 @@ def setUp(self):
         self.command.filename = None
         self.command.path = None
         self.command.connector = get_connector("default")
+        self.command.schemas = []
 
     def tearDown(self):
         clean_gpg_keys()

dbbackup/tests/commands/test_dbrestore.py

@@ -4,12 +4,12 @@
 
 from shutil import copyfileobj
 from tempfile import mktemp
+from unittest.mock import patch
 
 from django.conf import settings
 from django.core.files import File
 from django.core.management.base import CommandError
 from django.test import TestCase
-from mock import patch
 
 from dbbackup import utils
 from dbbackup.db.base import get_connector
@@ -47,6 +47,7 @@ def setUp(self):
         self.command.input_database_name = None
         self.command.database_name = "default"
         self.command.connector = get_connector("default")
+        self.command.schemas = []
         HANDLED_FILES.clean()
 
     def tearDown(self):
@@ -147,6 +148,7 @@ def setUp(self):
         self.command.database_name = "mongo"
         self.command.input_database_name = None
         self.command.servername = HOSTNAME
+        self.command.schemas = []
         HANDLED_FILES.clean()
         add_private_gpg()
 

dbbackup/tests/test_connectors/test_postgresql.py

@@ -1,7 +1,7 @@
 from io import BytesIO
+from unittest.mock import patch
 
 from django.test import TestCase
-from mock import patch
 
 from dbbackup.db.exceptions import DumpError
 from dbbackup.db.postgresql import (
@@ -132,6 +132,41 @@ def test_restore_dump_user(self, mock_dump_cmd, mock_restore_cmd):
             "postgresql://foo@hostname/dbname", mock_restore_cmd.call_args[0][0]
         )
 
+    def test_create_dump_schema(self, mock_dump_cmd):
+        # Without
+        self.connector.create_dump()
+        self.assertNotIn(" -n ", mock_dump_cmd.call_args[0][0])
+        # With
+        self.connector.schemas = ["public"]
+        self.connector.create_dump()
+        self.assertIn(" -n public", mock_dump_cmd.call_args[0][0])
+        # With several
+        self.connector.schemas = ["public", "foo"]
+        self.connector.create_dump()
+        self.assertIn(" -n public", mock_dump_cmd.call_args[0][0])
+        self.assertIn(" -n foo", mock_dump_cmd.call_args[0][0])
+
+    @patch(
+        "dbbackup.db.postgresql.PgDumpConnector.run_command",
+        return_value=(BytesIO(), BytesIO()),
+    )
+    def test_restore_dump_schema(self, mock_dump_cmd, mock_restore_cmd):
+        # Without
+        dump = self.connector.create_dump()
+        self.connector.restore_dump(dump)
+        self.assertNotIn(" -n ", mock_restore_cmd.call_args[0][0])
+        # With
+        self.connector.schemas = ["public"]
+        dump = self.connector.create_dump()
+        self.connector.restore_dump(dump)
+        self.assertIn(" -n public", mock_restore_cmd.call_args[0][0])
+        # With several
+        self.connector.schemas = ["public", "foo"]
+        dump = self.connector.create_dump()
+        self.connector.restore_dump(dump)
+        self.assertIn(" -n public", mock_restore_cmd.call_args[0][0])
+        self.assertIn(" -n foo", mock_restore_cmd.call_args[0][0])
+
 
 @patch(
     "dbbackup.db.postgresql.PgDumpBinaryConnector.run_command",
@@ -188,6 +223,41 @@ def test_restore_dump(self, mock_dump_cmd, mock_restore_cmd):
         # Test cmd
         self.assertTrue(mock_restore_cmd.called)
 
+    def test_create_dump_schema(self, mock_dump_cmd):
+        # Without
+        self.connector.create_dump()
+        self.assertNotIn(" -n ", mock_dump_cmd.call_args[0][0])
+        # With
+        self.connector.schemas = ["public"]
+        self.connector.create_dump()
+        self.assertIn(" -n public", mock_dump_cmd.call_args[0][0])
+        # With several
+        self.connector.schemas = ["public", "foo"]
+        self.connector.create_dump()
+        self.assertIn(" -n public", mock_dump_cmd.call_args[0][0])
+        self.assertIn(" -n foo", mock_dump_cmd.call_args[0][0])
+
+    @patch(
+        "dbbackup.db.postgresql.PgDumpBinaryConnector.run_command",
+        return_value=(BytesIO(), BytesIO()),
+    )
+    def test_restore_dump_schema(self, mock_dump_cmd, mock_restore_cmd):
+        # Without
+        dump = self.connector.create_dump()
+        self.connector.restore_dump(dump)
+        self.assertNotIn(" -n ", mock_restore_cmd.call_args[0][0])
+        # With
+        self.connector.schemas = ["public"]
+        dump = self.connector.create_dump()
+        self.connector.restore_dump(dump)
+        self.assertIn(" -n public", mock_restore_cmd.call_args[0][0])
+        # With several
+        self.connector.schemas = ["public", "foo"]
+        dump = self.connector.create_dump()
+        self.connector.restore_dump(dump)
+        self.assertIn(" -n public", mock_restore_cmd.call_args[0][0])
+        self.assertIn(" -n foo", mock_restore_cmd.call_args[0][0])
+
 
 @patch(
     "dbbackup.db.postgresql.PgDumpGisConnector.run_command",

docs/changelog.rst

@@ -7,6 +7,10 @@ Unreleased
 * Fix restore of database from S3 storage by reintroducing inputfile.seek(0) to utils.uncompress_file
 * Fix bug where dbbackup management command would not respect settings.py:DBBACKUP_DATABASES
 
+4.2.0 (2022-01-30)
+------------------
+* Add PostgreSQL Schema support by @angryfoxx in https://github.com/jazzband/django-dbbackup/pull/507
+
 4.1.0 (2024-01-14)
 ------------------
 

docs/databases.rst

@@ -198,6 +198,14 @@ ADMIN_PASSWORD
 Password used for launch action with privileges, extension creation for
 example.
 
+SCHEMAS
+~~~~~~~
+
+Specify schemas for database dumps by using a pattern-matching option,
+including both the selected schema and its contained objects.
+If not specified, the default behavior is to dump all non-system schemas in the target database.
+This feature is exclusive to PostgreSQL connectors, and users can choose multiple schemas for a customized dump.
+
 MongoDB
 -------