diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..33fe63f --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,4 @@ +{ + "python.linting.flake8Enabled": true, + "python.linting.enabled": true +} \ No newline at end of file diff --git a/dbbackup/db/postgresql.py b/dbbackup/db/postgresql.py index 89707ca..8d65fbe 100644 --- a/dbbackup/db/postgresql.py +++ b/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 diff --git a/dbbackup/management/commands/dbbackup.py b/dbbackup/management/commands/dbbackup.py index 541e8d8..ccb0adb 100644 --- a/dbbackup/management/commands/dbbackup.py +++ b/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: diff --git a/dbbackup/management/commands/dbrestore.py b/dbbackup/management/commands/dbrestore.py index 4cacfab..5f4c6a3 100644 --- a/dbbackup/management/commands/dbrestore.py +++ b/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)