From 7b1fa59a29ec59535976a53bd10528fe80f6e18c Mon Sep 17 00:00:00 2001
From: Gavin D'souza <gavin18d@gmail.com>
Date: Mon, 9 Nov 2020 18:51:33 +0530
Subject: [PATCH] feat: Restore partial backups via bench partial-restore

---
 frappe/commands/site.py       | 17 ++++++++++++++++-
 frappe/database/db_manager.py |  1 -
 frappe/installer.py           | 23 +++++++++++++++++++++++
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/frappe/commands/site.py b/frappe/commands/site.py
index 5305502b17..fd247b4182 100755
--- a/frappe/commands/site.py
+++ b/frappe/commands/site.py
@@ -91,6 +91,20 @@ def restore(context, sql_file_path, mariadb_root_username=None, mariadb_root_pas
 	success_message = "Site {0} has been restored{1}".format(site, " with files" if (with_public_files or with_private_files) else "")
 	click.secho(success_message, fg="green")
 
+@click.command('partial-restore')
+@click.argument('sql-file-path')
+@click.option("--verbose", "-v", is_flag=True)
+@pass_context
+def partial_restore(context, sql_file_path, verbose):
+	from frappe.installer import partial_restore
+	verbose = context.verbose or verbose
+
+	site = get_site(context)
+	frappe.init(site=site)
+	frappe.connect(site=site)
+	partial_restore(sql_file_path, verbose)
+	frappe.destroy()
+
 @click.command('reinstall')
 @click.option('--admin-password', help='Administrator Password for reinstalled site')
 @click.option('--mariadb-root-username', help='Root username for MariaDB')
@@ -650,5 +664,6 @@ commands = [
 	stop_recording,
 	add_to_hosts,
 	start_ngrok,
-	build_search_index
+	build_search_index,
+	partial_restore
 ]
diff --git a/frappe/database/db_manager.py b/frappe/database/db_manager.py
index 3345fce735..da1f584f57 100644
--- a/frappe/database/db_manager.py
+++ b/frappe/database/db_manager.py
@@ -3,7 +3,6 @@ import frappe
 
 
 class DbManager:
-
 	def __init__(self, db):
 		"""
 		Pass root_conn here for access to all databases.
diff --git a/frappe/installer.py b/frappe/installer.py
index b1420e5d9f..2bbf0421e3 100755
--- a/frappe/installer.py
+++ b/frappe/installer.py
@@ -519,3 +519,26 @@ def is_downgrade(sql_file_path, verbose=False):
 							print("Your site will be downgraded from Frappe {0} to {1}".format(current_version, backup_version))
 
 						return downgrade
+
+
+def partial_restore(sql_file_path, verbose=False):
+	sql_file = extract_sql_from_archive(sql_file_path)
+
+	if frappe.conf.db_type in (None, "mariadb"):
+		from frappe.database.mariadb.setup_db import import_db_from_sql
+	elif frappe.conf.db_type == "postgres":
+		from frappe.database.postgres.setup_db import import_db_from_sql
+		import warnings
+		from click import style
+		warn = style(
+			"Delete the tables you want to restore manually before attempting"
+			" partial restore operation for PostreSQL databases",
+			fg="yellow"
+		)
+		warnings.warn(warn)
+
+	import_db_from_sql(source_sql=sql_file, verbose=verbose)
+
+	# Removing temporarily created file
+	if sql_file != sql_file_path:
+		os.remove(sql_file)