From 0d0453291f385b93d98ebcaa86c5a7e5ee3ff1d6 Mon Sep 17 00:00:00 2001
From: henceiusegentoo <Leandersteveschreiber@gmail.com>
Date: Thu, 21 Sep 2023 19:55:14 +0200
Subject: [PATCH] Added backups

---
 .gitignore                         |  2 ++
 screen.sh => PUSH_TO_PRODUCTION.sh |  3 ++
 api/backup_manager/main.py         | 51 ++++++++++++++++++++++++++++++
 checks/client.py                   | 20 ++++++------
 4 files changed, 66 insertions(+), 10 deletions(-)
 rename screen.sh => PUSH_TO_PRODUCTION.sh (84%)
 create mode 100644 api/backup_manager/main.py

diff --git a/.gitignore b/.gitignore
index 29f5410..e16827a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -181,3 +181,5 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+backups/
\ No newline at end of file
diff --git a/screen.sh b/PUSH_TO_PRODUCTION.sh
similarity index 84%
rename from screen.sh
rename to PUSH_TO_PRODUCTION.sh
index 91c13cb..60c7d38 100755
--- a/screen.sh
+++ b/PUSH_TO_PRODUCTION.sh
@@ -3,6 +3,9 @@
 # Commit to the production branch
 # git commit -am "Auto-trigger - Production server started" && git push origin Production
 
+# backup database
+/usr/local/bin/python /home/nova-api/backup_manager/main.py pre_prodpush
+
 # Kill production server
 fuser -k 2333/tcp
 
diff --git a/api/backup_manager/main.py b/api/backup_manager/main.py
new file mode 100644
index 0000000..15a063e
--- /dev/null
+++ b/api/backup_manager/main.py
@@ -0,0 +1,51 @@
+from dotenv import load_dotenv
+import os
+from motor.motor_asyncio import AsyncIOMotorClient
+from bson import json_util
+import json
+from sys import argv
+import asyncio
+
+load_dotenv()
+
+MONGO_URI = os.getenv("MONGO_URI")
+FILE_DIR = os.path.dirname(os.path.realpath(__file__))
+
+async def main(output_dir):
+    await make_backup(output_dir)
+
+async def make_backup(output_dir):
+    output_dir = os.path.join(FILE_DIR, "..", "backups", output_dir)
+
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    client = AsyncIOMotorClient(MONGO_URI)
+    databases = await client.list_database_names()
+    databases = {db: await client[db].list_collection_names() for db in databases}
+
+    for database in databases:
+        if not os.path.exists(f"{output_dir}/{database}"):
+            os.mkdir(f"{output_dir}/{database}")
+
+        for collection in databases[database]:
+            print(f"Making backup for {database}/{collection}")
+            await make_backup_for_collection(database, collection, output_dir)
+
+async def make_backup_for_collection(database, collection, output_dir):
+    path = f"{output_dir}/{database}/{collection}.json"
+
+    client = AsyncIOMotorClient(MONGO_URI)
+    collection = client[database][collection]
+    documents = await collection.find({}).to_list(length=None)
+
+    with open(path, "w") as f:
+        json.dump(documents, f, default=json_util.default)
+
+if __name__ == "__main__":
+    if len(argv) < 2 or len(argv) > 2:
+        print("Usage: python3 main.py <output_dir>")
+        exit(1)
+
+    output_dir = argv[1]
+    asyncio.run(main(output_dir))
\ No newline at end of file
diff --git a/checks/client.py b/checks/client.py
index 9771c09..6b4fbe4 100644
--- a/checks/client.py
+++ b/checks/client.py
@@ -208,20 +208,20 @@ async def demo():
         else:
             raise ConnectionError('API Server is not running.')
 
-        # print('[lightblue]Checking if function calling works...')
-        # print(await test_function_calling())
+        print('[lightblue]Checking if function calling works...')
+        print(await test_function_calling())
 
-        # print('Checking non-streamed chat completions...')
-        # print(await test_chat_non_stream_gpt4())
+        print('Checking non-streamed chat completions...')
+        print(await test_chat_non_stream_gpt4())
 
-        # print('Checking streamed chat completions...')
-        # print(await test_chat_stream_gpt3())
+        print('Checking streamed chat completions...')
+        print(await test_chat_stream_gpt3())
 
-        # print('[lightblue]Checking if  image generation works...')
-        # print(await test_image_generation())
+        print('[lightblue]Checking if  image generation works...')
+        print(await test_image_generation())
 
-        # print('Checking the models endpoint...')
-        # print(await test_models())
+        print('Checking the models endpoint...')
+        print(await test_models())
 
     except Exception as exc:
         print('[red]Error: ' + str(exc))