Factorize code between commands and made some fixes

dbbackup/management/commands/_base.py

@@ -6,9 +6,11 @@
 from optparse import make_option
 from shutil import copyfileobj
 
-from django.core.management.base import BaseCommand, LabelCommand
+from django.core.management.base import BaseCommand, LabelCommand, CommandError
 from django.utils import six
 
+from ...storage.base import StorageError
+
 input = raw_input if six.PY2 else input  # @ReservedAssignment
 
 
@@ -32,18 +34,62 @@ def _set_logger_level(self):
         self.logger.setLevel(level)
 
     def _ask_confirmation(self):
-        answer = input("Are you sure you want to continue? [Y/n]")
+        answer = input("Are you sure you want to continue? [Y/n] ")
         if answer.lower().startswith('n'):
             self.logger.info("Quitting")
             sys.exit(0)
 
+    def read_from_storage(self, path):
+        return self.storage.read_file(path)
+
+    def write_to_storage(self, file, path):
+        self.logger.info("Writing file to %s: %s, filename: %s",
+                         self.storage.name, self.storage.backup_dir,
+                         path)
+        self.storage.write_file(file, path)
+
     def read_local_file(self, path):
         """Open file in read mode on local filesystem."""
         return open(path, 'rb')
 
-    # TODO: Define chunk size
     def write_local_file(self, outputfile, path):
         """Write file to the desired path."""
+        self.logger.info("Writing file to %s", path)
         outputfile.seek(0)
         with open(path, 'wb') as fd:
             copyfileobj(outputfile, fd)
+
+    def _get_backup_file(self):
+        if self.path:
+            input_filename = self.path
+            input_file = self.read_local_file(self.path)
+        else:
+            if self.filename:
+                input_filename = self.filename
+            # Fetch the latest backup if filepath not specified
+            else:
+                self.logger.info("Finding latest backup")
+                # database = self.database['NAME'] if self.content_type == 'db' else None
+                try:
+                    input_filename = self.storage.get_latest_backup(
+                        encrypted=self.decrypt,
+                        compressed=self.uncompress,
+                        content_type=self.content_type)
+                        # TODO: Make better filter
+                        # database=database)
+                except StorageError as err:
+                    raise CommandError(err.args[0])
+            input_file = self.read_from_storage(input_filename)
+        return input_filename, input_file
+
+    def _cleanup_old_backups(self):
+        """
+        Cleanup old backups, keeping the number of backups specified by
+        DBBACKUP_CLEANUP_KEEP and any backups that occur on first of the month.
+        """
+        # database = self.database if self.content_type == 'db' else None
+        file_list = self.storage.clean_old_backups(encrypted=self.encrypt,
+                                                   compressed=self.compress,
+                                                   content_type=self.content_type)
+                                                   # TODO: Make better filter
+                                                   # database=database)

dbbackup/management/commands/dbbackup.py

@@ -15,9 +15,10 @@
 
 
 class Command(BaseDbBackupCommand):
-    help = """
-    Backup a database, encrypt and/or compress and write to storage.
-    """
+    help = """Backup a database, encrypt and/or compress and write to
+    storage."""
+    content_type = 'db'
+
     option_list = BaseDbBackupCommand.option_list + (
         make_option("-c", "--clean", dest='clean', action="store_true",
                     default=False, help="Clean up old backup files"),
@@ -41,7 +42,6 @@ def handle(self, **options):
         self.verbosity = int(options.get('verbosity'))
         self.quiet = options.get('quiet')
         self.clean = options.get('clean')
-        self.clean_keep = dbbackup_settings.CLEANUP_KEEP
         self.database = options.get('database')
         self.servername = options.get('servername')
         self.compress = options.get('compress')
@@ -56,10 +56,7 @@ def handle(self, **options):
             try:
                 self._save_new_backup(database)
                 if self.clean:
-                    self.storage.clean_old_backups(self.encrypt,
-                                                   self.compress,
-                                                   content_type='db',
-                                                   database=self.database)
+                    self._cleanup_old_backups()
             except StorageError as err:
                 raise CommandError(err)
 
@@ -82,10 +79,6 @@ def _save_new_backup(self, database):
             self.logger.info("Backup size: %s", utils.handle_size(outputfile))
         # Store backup
         if self.path is None:
-            self.logger.info("Writing file to %s: %s, filename: %s",
-                             self.storage.name, self.storage.backup_dir,
-                             filename)
-            self.storage.write_file(outputfile, filename)
+            self.write_to_storage(outputfile, filename)
         else:
-            self.logger.info("Writing file to %s", self.path)
             self.write_local_file(outputfile, self.path)

dbbackup/management/commands/dbrestore.py

@@ -17,9 +17,10 @@
 
 
 class Command(BaseDbBackupCommand):
-    help = """
-    Restore a 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 + (
         make_option("-d", "--database", help="Database to restore"),
         make_option("-i", "--input-filename", help="Specify filename to backup from"),
@@ -28,7 +29,7 @@ class Command(BaseDbBackupCommand):
 
         make_option("-c", "--decrypt", help="Decrypt data before restoring", default=False, action='store_true'),
         make_option("-p", "--passphrase", help="Passphrase for decrypt file", default=None),
-        make_option("-z", "--uncompress", help="Uncompress gzip data before restoring", action='store_true'),
+        make_option("-z", "--uncompress", help="Uncompress gzip data before restoring", action='store_true', default=False),
     )
 
     def handle(self, *args, **options):
@@ -62,25 +63,6 @@ def _get_database(self, options):
             database_key = list(settings.DATABASES.keys())[0]
         return settings.DATABASES[database_key]
 
-    def _get_backup_file(self):
-        if self.path:
-            input_filename = self.path
-            input_file = self.read_local_file(self.path)
-        else:
-            if self.filename:
-                input_filename = self.filename
-            # Fetch the latest backup if filepath not specified
-            else:
-                self.logger.info("Finding latest backup")
-                try:
-                    input_filename = self.storage.get_latest_backup(encrypted=self.decrypt,
-                                                                    compressed=self.uncompress,
-                                                                    content_type='db')
-                except StorageError as err:
-                    raise CommandError(err.args[0])
-            input_file = self.storage.read_file(input_filename)
-        return input_filename, input_file
-
     def _restore_backup(self):
         """Restore the specified database."""
         self.logger.info("Restoring backup for database: %s", self.database['NAME'])

dbbackup/management/commands/mediabackup.py

@@ -13,31 +13,34 @@
 from ._base import BaseDbBackupCommand
 from ... import utils
 from ...storage.base import BaseStorage, StorageError
-from ... import settings
 
 
 class Command(BaseDbBackupCommand):
-    help = """
-    Backup media files, gather all in a tarball and encrypt or compress.
-    """
-
+    help = """Backup media files, gather all in a tarball and encrypt or
+    compress."""
     content_type = "media"
 
     option_list = BaseDbBackupCommand.option_list + (
         make_option("-c", "--clean", help="Clean up old backup files", action="store_true",
                     default=False),
         make_option("-s", "--servername", help="Specify server name to include in backup filename"),
-        make_option("-e", "--encrypt", help="Encrypt the backup files", action="store_true",
-                    default=False),
         make_option("-z", "--compress", help="Do not compress the archive", action="store_true",
                     default=False),
+        make_option("-e", "--encrypt", help="Encrypt the backup files", action="store_true",
+                    default=False),
+        make_option("-o", "--output-filename", default=None,
+                    help="Specify filename on storage"),
+        make_option("-O", "--output-path", default=None,
+                    help="Specify where to store on local filesystem",)
     )
 
     @utils.email_uncaught_exception
     def handle(self, *args, **options):
         self.encrypt = options.get('encrypt', False)
         self.compress = options.get('compress', False)
-        self.servername = options.get('servername') or settings.HOSTNAME
+        self.servername = options.get('servername')
+        self.filename = options.get('output_filename')
+        self.path = options.get('output_path')
         try:
             self.media_storage = get_storage_class()()
             self.storage = BaseStorage.storage_factory()
@@ -48,28 +51,8 @@ def handle(self, *args, **options):
         except StorageError as err:
             raise CommandError(err)
 
-    def backup_mediafiles(self):
-        """
-        Create backup file and write it to storage.
-        """
-        # Create file name
-        extension = "tar%s" % ('.gz' if self.compress else '')
-        filename = utils.filename_generate(extension,
-                                           servername=self.servername,
-                                           content_type=self.content_type)
-
-        outputfile = self._create_tar(filename)
-
-        if self.encrypt:
-            encrypted_file = utils.encrypt_file(outputfile, filename)
-            outputfile, filename = encrypted_file
-
-        self.logger.debug("Backup size: %s", utils.handle_size(outputfile))
-        self.logger.info("Writing file to %s" % filename)
-        self.storage.write_file(outputfile, filename)
-
     def _explore_storage(self):
-        """Generator of a all files contained in media storage."""
+        """Generator of all files contained in media storage."""
         path = ''
         dirs = [path]
         while dirs:
@@ -93,12 +76,24 @@ def _create_tar(self, name):
         tar_file.close()
         return fileobj
 
-    def _cleanup_old_backups(self):
+    def backup_mediafiles(self):
         """
-        Cleanup old backups, keeping the number of backups specified by
-        DBBACKUP_CLEANUP_KEEP and any backups that occur on first of the month.
+        Create backup file and write it to storage.
         """
-        self.logger.info("Cleaning Old Backups for media files")
-        file_list = self.storage.clean_old_backups(encrypted=self.encrypt,
-                                                   compressed=self.compress,
-                                                   keep_number=settings.CLEANUP_KEEP_MEDIA)
+        # Create file name
+        extension = "tar%s" % ('.gz' if self.compress else '')
+        filename = utils.filename_generate(extension,
+                                           servername=self.servername,
+                                           content_type=self.content_type)
+
+        outputfile = self._create_tar(filename)
+
+        if self.encrypt:
+            encrypted_file = utils.encrypt_file(outputfile, filename)
+            outputfile, filename = encrypted_file
+
+        self.logger.debug("Backup size: %s", utils.handle_size(outputfile))
+        if self.path is None:
+            self.write_to_storage(outputfile, filename)
+        else:
+            self.storage.write_file(outputfile, filename)

dbbackup/management/commands/mediarestore.py

@@ -13,13 +13,18 @@
 
 
 class Command(BaseDbBackupCommand):
+    help = """Restore a media backup from storage, encrypted and/or
+    compressed."""
+    content_type = 'media'
+
     option_list = BaseDbBackupCommand.option_list + (
         make_option("-i", "--input-filename", help="Specify filename to backup from"),
         make_option("-I", "--input-path", help="Specify path on local filesystem to backup from"),
 
-        make_option("-c", "--decrypt", help="Decrypt data before restoring", default=False, action='store_true'),
+        make_option("-e", "--decrypt", help="Decrypt data before restoring", default=False, action='store_true'),
         make_option("-p", "--passphrase", help="Passphrase for decrypt file", default=None),
         make_option("-z", "--uncompress", help="Uncompress gzip data before restoring", action='store_true'),
+        make_option("-r", "--replace", help="Replace existing files", action='store_true'),
     )
 
     def handle(self, *args, **options):
@@ -31,47 +36,35 @@ def handle(self, *args, **options):
         self.servername = options.get('servername')
         self.decrypt = options.get('decrypt')
         self.uncompress = options.get('uncompress')
+        self.replace = options.get('replace')
         self.passphrase = options.get('passphrase')
         self.interactive = options.get('interactive')
         self.storage = BaseStorage.storage_factory()
         self.media_storage = get_storage_class()()
         self._restore_backup()
 
-    def _get_backup_file(self):
-        if self.path:
-            input_filename = self.path
-            input_file = self.read_local_file(self.path)
-        else:
-            if self.filename:
-                input_filename = self.filename
-            # Fetch the latest backup if filepath not specified
-            else:
-                self.logger.info("Finding latest backup")
-                try:
-                    input_filename = self.storage.get_latest_backup(
-                        encrypted=self.decrypt,
-                        compressed=self.uncompress,
-                        content_type='media')
-                except StorageError as err:
-                    raise CommandError(err.args[0])
-            input_file = self.storage.read_file(input_filename)
-        return input_filename, input_file
-
     def _upload_file(self, name, media_file):
+        if self.media_storage.exists(name):
+            if self.replace:
+                self.media_storage.delete(name)
+                self.logger.info("%s deleted", name)
+            else:
+                return
         self.media_storage.save(name, media_file)
+        self.logger.info("%s uploaded", name)
 
     def _restore_backup(self):
         self.logger.info("Restoring backup for media files")
         input_filename, input_file = self._get_backup_file()
-        self.logger.info("Restoring: %s" % input_filename)
+        self.logger.info("Restoring: %s", input_filename)
 
         if self.decrypt:
             unencrypted_file, input_filename = utils.unencrypt_file(input_file, input_filename,
                                                                     self.passphrase)
             input_file.close()
             input_file = unencrypted_file
 
-        self.logger.info("Restore tempfile created: %s", utils.handle_size(input_file))
+        self.logger.info("Backup size: %s", utils.handle_size(input_file))
         if self.interactive:
             self._ask_confirmation()
 

dbbackup/storage/base.py

@@ -200,7 +200,7 @@ def clean_old_backups(self, encrypted=None, compressed=None,
         """
         if keep_number is None:
             keep_number = settings.CLEANUP_KEEP if content_type == 'db' \
-                else settings.MEDIA_FILENAME_TEMPLATE
+                else settings.CLEANUP_KEEP_MEDIA
         files = self.list_backups(encrypted=encrypted, compressed=compressed,
                                   content_type=content_type, database=database)
         files = sorted(files, key=utils.filename_to_date, reverse=True)