Add postgresql schema support

.vscode/settings.json

@@ -0,0 +1,4 @@
+{
+    "python.linting.flake8Enabled": true,
+    "python.linting.enabled": true
+}

dbbackup/db/postgresql.py

@@ -36,6 +36,8 @@ class PgDumpConnector(BaseCommandDBConnector):
     restore_cmd = 'psql'
     single_transaction = True
     drop = True
+    schema = None
+    no_owner = False
 
     def _create_dump(self):
         cmd = '{} '.format(self.dump_cmd)
@@ -45,7 +47,8 @@ def _create_dump(self):
             cmd += ' --exclude-table-data={}'.format(table)
         if self.drop:
             cmd += ' --clean'
-
+        if self.schema:
+            cmd += ' -n {}'.format(self.schema)
         cmd = '{} {} {}'.format(self.dump_prefix, cmd, self.dump_suffix)
         stdout, stderr = self.run_command(cmd, env=self.dump_env)
         return stdout
@@ -56,8 +59,12 @@ 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.schema:
+            cmd += ' -n {}'.format(self.schema)
         if self.single_transaction:
             cmd += ' --single-transaction'
+        if self.no_owner:
+            cmd += ' --no-owner'
         cmd += ' {}'.format(self.settings['NAME'])
         cmd = '{} {} {}'.format(self.restore_prefix, cmd, self.restore_suffix)
         stdout, stderr = self.run_command(cmd, stdin=dump, env=self.restore_env)
@@ -103,6 +110,8 @@ def _create_dump(self):
         cmd = '{} '.format(self.dump_cmd)
         cmd = cmd + create_postgres_uri(self)
 
+        if self.schema:
+            cmd += ' -n {}'.format(self.schema)
         cmd += ' --format=custom'
         for table in self.exclude:
             cmd += ' --exclude-table-data={}'.format(table)
@@ -114,10 +123,14 @@ def _restore_dump(self, dump):
         dbname = create_postgres_uri(self)
         cmd = '{} {}'.format(self.restore_cmd, dbname)
 
+        if self.schema:
+            cmd += ' -n {}'.format(self.schema)
         if self.single_transaction:
             cmd += ' --single-transaction'
         if self.drop:
             cmd += ' --clean'
+        if self.no_owner:
+            cmd += ' --no-owner'
         cmd = '{} {} {}'.format(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

@@ -34,7 +34,9 @@ class Command(BaseDbBackupCommand):
         make_option("-O", "--output-path", default=None,
                     help="Specify where to store on local filesystem"),
         make_option("-x", "--exclude-tables", default=None,
-                    help="Exclude tables from backup")
+                    help="Exclude tables from backup"),
+        make_option("-n", "--schema", default=None,
+                    help="Specify schema to limit backup to given schema only")
     )
 
     @utils.email_uncaught_exception
@@ -52,6 +54,7 @@ def handle(self, **options):
         self.filename = options.get('output_filename')
         self.path = options.get('output_path')
         self.exclude_tables = options.get("exclude_tables")
+        self.schema = options.get("schema")
         self.storage = get_storage()
 
         self.database = options.get('database') or ''
@@ -76,6 +79,8 @@ def _save_new_backup(self, database):
         self.logger.info("Backing Up Database: %s", database['NAME'])
         # Get backup and name
         filename = self.connector.generate_filename(self.servername)
+        if self.schema:
+            self.connector.schema = self.schema or None
         outputfile = self.connector.create_dump()
         # Apply trans
         if self.compress:

dbbackup/management/commands/dbrestore.py

@@ -30,7 +30,11 @@ class Command(BaseDbBackupCommand):
                     help="Decrypt data before restoring"),
         make_option("-p", "--passphrase", help="Passphrase for decrypt file", default=None),
         make_option("-z", "--uncompress", action='store_true', default=False,
-                    help="Uncompress gzip data before restoring")
+                    help="Uncompress gzip data before restoring"),
+        make_option("-n", "--schema", default=None,
+                    help="Restore backup to given schema only"),
+        make_option("--no-owner", action='store_true', dest="no-owner",
+                    help="Don't try to set the ownership of the objects to the original owner")
     )
 
     def handle(self, *args, **options):
@@ -48,6 +52,8 @@ def handle(self, *args, **options):
             self.uncompress = options.get('uncompress')
             self.passphrase = options.get('passphrase')
             self.interactive = options.get('interactive')
+            self.schema = options.get("schema")
+            self.no_owner = options.get("no-owner")
             self.database_name, self.database = self._get_database(options)
             self.storage = get_storage()
             self._restore_backup()
@@ -73,6 +79,8 @@ def _restore_backup(self):
                                                            servername=self.servername)
         self.logger.info("Restoring backup for database '%s' and server '%s'",
                          self.database_name, self.servername)
+        if self.schema:
+            self.logger.info("Restoring to schema: %s" % self.schema)
         self.logger.info("Restoring: %s" % input_filename)
 
         if self.decrypt:
@@ -91,4 +99,8 @@ def _restore_backup(self):
 
         input_file.seek(0)
         self.connector = get_connector(self.database_name)
+        if self.schema:
+            self.connector.schema = self.schema
+        if self.no_owner:
+            self.connector.no_owner = True
         self.connector.restore_dump(input_file)