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,7 +34,7 @@ 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)
@@ -47,3 +49,38 @@ def write_local_file(self, outputfile, 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.storage.read_file(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,
+                                                   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)
 

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,32 @@
 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')
         try:
             self.media_storage = get_storage_class()()
             self.storage = BaseStorage.storage_factory()
@@ -69,7 +70,7 @@ def backup_mediafiles(self):
         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:
@@ -92,13 +93,3 @@ def _create_tar(self, name):
         # Close the TAR for writing
         tar_file.close()
         return fileobj
-
-    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.
-        """
-        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)

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)

dbbackup/tests/commands/test_base.py

@@ -6,6 +6,7 @@
 from django.test import TestCase
 from django.utils import six
 from dbbackup.management.commands._base import BaseDbBackupCommand
+from dbbackup.tests.utils import (FakeStorage, DEV_NULL, HANDLED_FILES)
 
 
 class BaseDbBackupCommandSetLoggerLevelTest(TestCase):
@@ -27,6 +28,11 @@ def test_quiet(self):
         self.command._set_logger_level()
         self.assertGreater(self.command.logger.level, 50)
 
+
+class BaseDbBackupCommandMethodsTest(TestCase):
+    def setUp(self):
+        self.command = BaseDbBackupCommand()
+
     def test_read_local_file(self):
         # setUp
         self.command.path = '/tmp/foo.bak'
@@ -63,3 +69,37 @@ def test_ask_confirmation(self):
         with patch('dbbackup.management.commands._base.input', return_value='No'):
             with self.assertRaises(SystemExit):
                 self.command._ask_confirmation()
+
+
+class BaseDbBackupCommandCleanupOldBackupsTest(TestCase):
+    def setUp(self):
+        HANDLED_FILES.clean()
+        self.command = BaseDbBackupCommand()
+        self.command.stdout = DEV_NULL
+        self.command.encrypt = False
+        self.command.compress = False
+        self.command.servername = 'foo-server'
+        self.command.storage = FakeStorage()
+        HANDLED_FILES['written_files'] = [(f, None) for f in [
+            '2015-02-06-042810.tar',
+            '2015-02-07-042810.tar',
+            '2015-02-08-042810.tar',
+            'foodb-2015-02-06-042810.dump',
+            'foodb-2015-02-07-042810.dump',
+            'foodb-2015-02-08-042810.dump',
+        ]]
+
+    @patch('dbbackup.settings.CLEANUP_KEEP', 1)
+    def test_clean_db(self):
+        self.command.content_type = 'db'
+        self.command.database = 'foodb'
+        self.command._cleanup_old_backups()
+        self.assertEqual(2, len(HANDLED_FILES['deleted_files']))
+        self.assertNotIn('foodb-2015-02-08-042810.dump', HANDLED_FILES['deleted_files'])
+
+    @patch('dbbackup.settings.CLEANUP_KEEP_MEDIA', 1)
+    def test_clean_media(self):
+        self.command.content_type = 'media'
+        self.command._cleanup_old_backups()
+        self.assertEqual(2, len(HANDLED_FILES['deleted_files']))
+        self.assertNotIn('2015-02-08-042810.tar', HANDLED_FILES['deleted_files'])

dbbackup/tests/commands/test_dbrestore.py

@@ -42,7 +42,7 @@ def tearDown(self):
     def test_no_filename(self, *args):
         # Prepare backup
         HANDLED_FILES['written_files'].append(
-            (utils.filename_generate('foo'), get_dump()))
+            (utils.filename_generate(TEST_DATABASE), get_dump()))
         # Check
         self.command.path = None
         self.command.filename = None