From a8f4659552e936dd050e387df635170cb0738a05 Mon Sep 17 00:00:00 2001
From: palmtree5 <3577255+palmtree5@users.noreply.github.com>
Date: Fri, 16 Mar 2018 13:37:12 -0800
Subject: [PATCH] [V3 Instance Setup] Storage swapping (#1421)

* [V3 Instance Setup] start work on storage swapping

* This should do the trick for Mongo -> JSON

* Fix typo

* Fix a few more typos

* resolve the data path

* Upsert the imported data

* need a list of the documents

* to_list is a coro
---
 redbot/setup.py | 181 +++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 156 insertions(+), 25 deletions(-)

diff --git a/redbot/setup.py b/redbot/setup.py
index acc15613f..f519d3bf0 100644
--- a/redbot/setup.py
+++ b/redbot/setup.py
@@ -1,6 +1,8 @@
 #!/usr/bin/env python
 
 import argparse
+import asyncio
+import json
 import os
 import shutil
 import sys
@@ -40,6 +42,11 @@ def parse_cli_args():
         help="Interactively delete an instance",
         action="store_true"
     )
+    parser.add_argument(
+        "--edit", "-e",
+        help="Interactively edit an instance",
+        action="store_true"
+    )
     return parser.parse_known_args()
 
 
@@ -59,12 +66,7 @@ def save_config(name, data, remove=False):
     JsonIO(config_file)._save_json(config)
 
 
-def basic_setup():
-    """
-    Creates the data storage folder.
-    :return:
-    """
-
+def get_data_dir():
     default_data_dir = Path(appdir.user_data_dir)
 
     print("Hello! Before we begin the full configuration process we need to"
@@ -96,10 +98,10 @@ def basic_setup():
     if not confirm("Please confirm (y/n):"):
         print("Please start the process over.")
         sys.exit(0)
+    return default_data_dir
 
-    default_dirs = deepcopy(basic_config_default)
-    default_dirs['DATA_PATH'] = str(default_data_dir.resolve())
 
+def get_storage_type():
     storage_dict = {
         1: "JSON",
         2: "MongoDB"
@@ -118,15 +120,10 @@ def basic_setup():
         else:
             if storage not in storage_dict:
                 storage = None
+    return storage
 
-    default_dirs['STORAGE_TYPE'] = storage_dict.get(storage, 1)
-
-    if storage_dict.get(storage, 1) == "MongoDB":
-        from redbot.core.drivers.red_mongo import get_config_details
-        default_dirs['STORAGE_DETAILS'] = get_config_details()
-    else:
-        default_dirs['STORAGE_DETAILS'] = {}
 
+def get_name():
     name = ""
     while len(name) == 0:
         print()
@@ -135,7 +132,35 @@ def basic_setup():
         name = input("> ")
         if " " in name:
             name = ""
+    return name
 
+
+def basic_setup():
+    """
+    Creates the data storage folder.
+    :return:
+    """
+
+    default_data_dir = get_data_dir()
+
+    default_dirs = deepcopy(basic_config_default)
+    default_dirs['DATA_PATH'] = str(default_data_dir.resolve())
+
+    storage = get_storage_type()
+
+    storage_dict = {
+        1: "JSON",
+        2: "MongoDB"
+    }
+    default_dirs['STORAGE_TYPE'] = storage_dict.get(storage, 1)
+
+    if storage_dict.get(storage, 1) == "MongoDB":
+        from redbot.core.drivers.red_mongo import get_config_details
+        default_dirs['STORAGE_DETAILS'] = get_config_details()
+    else:
+        default_dirs['STORAGE_DETAILS'] = {}
+
+    name = get_name()
     save_config(name, default_dirs)
 
     print()
@@ -143,6 +168,116 @@ def basic_setup():
           " continue your setup process and to run the bot.")
 
 
+async def edit_instance():
+    instance_list = load_existing_config()
+    if not instance_list:
+        print("No instances have been set up!")
+        return
+
+    print(
+        "You have chosen to edit an instance. The following "
+        "is a list of instances that currently exist:\n"
+    )
+    for instance in instance_list.keys():
+        print("{}\n".format(instance))
+    print("Please select one of the above by entering its name")
+    selected = input("> ")
+
+    if selected not in instance_list.keys():
+        print("That isn't a valid instance!")
+        return
+    instance_data = instance_list[selected]
+    default_dirs = deepcopy(basic_config_default)
+
+    current_data_dir = Path(instance_data["DATA_PATH"])
+    print(
+        "You have selected '{}' as the instance to modify.".format(selected)
+    )
+    if not confirm("Please confirm (y/n):"):
+        print("Ok, we will not continue then.")
+        return
+
+    print("Ok, we will continue on.")
+    print()
+    if confirm("Would you like to change the instance name? (y/n)"):
+        name = get_name()
+    else:
+        name = selected
+
+    if confirm("Would you like to change the data location? (y/n)"):
+        default_data_dir = get_data_dir()
+        default_dirs["DATA_PATH"] = str(default_data_dir.resolve())
+    else:
+        default_dirs["DATA_PATH"] = str(current_data_dir.resolve())
+
+    if confirm("Would you like to change the storage type? (y/n):"):
+        storage = get_storage_type()
+
+        storage_dict = {
+            1: "JSON",
+            2: "MongoDB"
+        }
+        default_dirs["STORAGE_TYPE"] = storage_dict[storage]
+        if storage_dict.get(storage, 1) == "MongoDB":
+            from redbot.core.drivers.red_mongo import get_config_details, Mongo
+            storage_details = get_config_details()
+            default_dirs["STORAGE_DETAILS"] = storage_details
+
+            if instance_data["STORAGE_TYPE"] == "JSON":
+                if confirm("Would you like to import your data? (y/n)"):
+                    core_data_file = list(current_data_dir.glob("core/settings.json"))[0]
+                    m = Mongo("Core", **storage_details)
+                    with core_data_file.open(mode="r") as f:
+                        core_data = json.loads(f.read())
+                    m.unique_cog_identifier = 0
+                    collection = m.get_collection()
+                    await collection.update_one(
+                        {'_id': m.unique_cog_identifier},
+                        update={"$set": core_data["0"]},
+                        upsert=True
+                    )
+                    for p in current_data_dir.glob("cogs/**/settings.json"):
+                        cog_m = Mongo(p.parent.stem, **storage_details)
+                        cog_c = cog_m.get_collection()
+                        with p.open(mode="r") as f:
+                            cog_data = json.loads(f.read())
+                        for ident in list(cog_data.keys()):
+                            cog_m.unique_cog_identifier = int(ident)
+                            await cog_c.update_one(
+                                {"_id": cog_m.unique_cog_identifier},
+                                update={"$set": cog_data[ident]},
+                                upsert=True
+                            )
+        else:
+            default_dirs["STORAGE_DETAILS"] = {}
+            if instance_data["STORAGE_TYPE"] == "MongoDB":
+                from redbot.core.drivers.red_mongo import Mongo
+                from redbot.core.drivers.red_json import JSON
+                m = Mongo("Core", **instance_data["STORAGE_DETAILS"])
+                db = m.db
+                collection_names = await db.collection_names(include_system_collections=False)
+                for c_name in collection_names:
+                    if c_name == "Core":
+                        c_data_path = current_data_dir / "core"
+                    else:
+                        c_data_path = current_data_dir / "cogs/{}".format(c_name)
+                    output = {}
+                    docs = await db[c_name].find().to_list(None)
+                    for item in docs:
+                        item_id = str(item.pop("_id"))
+                        output[item_id] = item
+                    target = JSON(c_name, data_path_override=c_data_path)
+                    await target.jsonIO._threadsafe_save_json(output)
+
+    if name != selected:
+        save_config(selected, {}, remove=True)
+    save_config(name, default_dirs)
+
+    print(
+        "Your basic configuration has been edited"
+    )
+
+
 def remove_instance():
     instance_list = load_existing_config()
     if not instance_list:
@@ -162,12 +297,8 @@ def remove_instance():
         print("That isn't a valid instance!")
         return
     instance_data = instance_list[selected]
-    print(
-        "Would you like to make a backup of "
-        "the data for this instance (y/n)?"
-    )
-    yesno = input("> ")
-    if yesno.lower() == "y":
+
+    if confirm("Would you like to make a backup of the data for this instance? (y/n)"):
         if instance_data["STORAGE_TYPE"] == "MongoDB":
             raise NotImplementedError(
                 "Support for removing instances with MongoDB as the storage "
@@ -197,7 +328,7 @@ def remove_instance():
             save_config(selected, {}, remove=True)
             print("The instance has been removed")
             return
-    elif yesno.lower() == "n":
+    else:
         pth = Path(instance_data["DATA_PATH"])
         print("Removing the instance...")
         try:
@@ -207,9 +338,6 @@ def remove_instance():
         save_config(selected, {}, remove=True)
         print("The instance has been removed")
         return
-    else:
-        print("That's not a valid option!")
-        return
 
 
 def main():
@@ -218,6 +346,9 @@ def main():
             remove_instance()
         except NotImplementedError as e:
             print(str(e))
+    elif args.edit:
+        loop = asyncio.get_event_loop()
+        loop.run_until_complete(edit_instance())
     else:
         basic_setup()