From 163bd9ed75052473d1b2b24faaf2488d69f2d188 Mon Sep 17 00:00:00 2001
From: eberonsh <87321531+eberonsh@users.noreply.github.com>
Date: Fri, 23 Aug 2024 18:40:51 -0300
Subject: [PATCH] Add `no-drop` parameter dbrestore (#400)

Co-authored-by: Emiliano Beronich <emilianoberonich@gmail.com>
Co-authored-by: Mark <16909269+Archmonger@users.noreply.github.com>
---
 dbbackup/management/commands/dbrestore.py | 12 ++++++++++--
 docs/changelog.rst                        |  2 +-
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/dbbackup/management/commands/dbrestore.py b/dbbackup/management/commands/dbrestore.py
index babc6b8..ce0ad7e 100644
--- a/dbbackup/management/commands/dbrestore.py
+++ b/dbbackup/management/commands/dbrestore.py
@@ -15,6 +15,7 @@
 class Command(BaseDbBackupCommand):
     help = "Restore a database backup from storage, encrypted and/or compressed."
     content_type = "db"
+    no_drop = False
 
     option_list = BaseDbBackupCommand.option_list + (
         make_option("-d", "--database", help="Database to restore"),
@@ -52,6 +53,13 @@ class Command(BaseDbBackupCommand):
             default=[],
             help="Specify schema(s) to restore. Can be used multiple times.",
         ),
+        make_option(
+            "-r",
+            "--no-drop",
+            action="store_true",
+            default=False,
+            help="Don't clean (drop) the database. This only works with mongodb and postgresql.",
+        ),
     )
 
     def handle(self, *args, **options):
@@ -74,6 +82,7 @@ def handle(self, *args, **options):
                 self.input_database_name
             )
             self.storage = get_storage()
+            self.no_drop = options.get("no_drop")
             self.schemas = options.get("schema")
             self._restore_backup()
         except StorageError as err:
@@ -129,8 +138,7 @@ def _restore_backup(self):
 
         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)
+        self.connector.drop = not self.no_drop
diff --git a/docs/changelog.rst b/docs/changelog.rst
index 88fbda3..3472452 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -4,7 +4,7 @@ Changelog
 Unreleased
 ----------
 
-* Nothing (yet)!
+* Add `--no-drop` option to `dbrestore` command to prevent dropping tables before restoring data.
 
 4.2.0 (2024-08-22)
 ------------------