From 7f24bc7c73775dbb162a093bed7c711e59dd9676 Mon Sep 17 00:00:00 2001
From: AlvaroVillanueva <alvaro.villanueva@deimos-space.com>
Date: Fri, 28 Jan 2022 09:24:55 +0000
Subject: [PATCH] EOEPCA-515 New import export function in management_tools

---
 src/handlers/mongo_handler.py | 81 +++++++++++++++++++++++++++++++++++
 src/management_tools.py       | 22 +++++++++-
 2 files changed, 102 insertions(+), 1 deletion(-)

diff --git a/src/handlers/mongo_handler.py b/src/handlers/mongo_handler.py
index 50cf95e..34479d8 100644
--- a/src/handlers/mongo_handler.py
+++ b/src/handlers/mongo_handler.py
@@ -1,6 +1,9 @@
 #!/usr/bin/env python3
 import pymongo
 import logging
+import json
+from bson.objectid import ObjectId
+from bson.json_util import dumps
 
 class Mongo_Handler:
 
@@ -190,3 +193,81 @@ def insert_rpt_in_mongo(self, rpt: str, rpt_limit_uses: int, timestamp: str):
             myres = { "rpt": rpt, "rpt_limit_uses": rpt_limit_uses, "timestamp": timestamp }
             x = col.insert_one(myres)
             return x
+
+                     
+    def export_database(self, res_path, tkn_path):
+        '''
+        Exports the main colection to a json file located in the / of the container or the specified path
+        '''
+        col = self.db['resources']
+        cursor = col.find({})
+        count = col.count_documents({})
+        path = ''
+        if res_path:
+            path = res_path
+        else:
+            path = "/resources.json"
+        with open(str(path), 'w') as file:
+            file.write('[')
+            aux = 0
+            for document in cursor:
+                aux += 1 
+                n = document
+                file.write(dumps(document))
+                if aux != count:
+                    file.write(',')
+            file.write(']')
+        #Same for the token collection
+        col = self.db['rpts']
+        cursor = col.find({})
+        count = col.count_documents({})
+        tknpath = ''
+        if tkn_path:
+            tknpath = tkn_path
+        else:
+            tknpath = "/rpts.json"
+        with open(str(tknpath), 'w') as file:
+            file.write('[')
+            aux = 0
+            for document in cursor:
+                aux += 1 
+                n = document
+                file.write(dumps(document))
+                if aux != count:
+                    file.write(',')
+            file.write(']')
+        
+        return "Exported resources to "+ str(path) + " and rpts to " + str(tknpath)
+
+                   
+    def import_database(self, res_path, tkn_path):
+        '''
+        Imports the exported data of a previous database located in the / of the container or in the specified path
+        '''
+        col = self.db['resources']
+        path = ''
+        if res_path:
+            path = res_path
+        else:
+            path = "/resources.json"
+        with open(str(path)) as f:
+            n = f.readlines()
+            for i in n:
+                for u in json.loads(i):
+                    oid = u["_id"]["$oid"]
+                    u["_id"] = ObjectId(str(oid))
+                    col.insert_one(u)
+        col = self.db['rpts']
+        tknpath = ''
+        if tkn_path:
+            tknpath = tkn_path
+        else:
+            tknpath = "/rpts.json"
+        with open(str(tknpath)) as f:
+            n = f.readlines()
+            for i in n:
+                for u in json.loads(i):
+                    oid = u["_id"]["$oid"]
+                    u["_id"] = ObjectId(str(oid))
+                    col.insert_one(u)
+        return "Inserted "+ str(path) + ' and '+ str(tknpath) + ' files in database'
\ No newline at end of file
diff --git a/src/management_tools.py b/src/management_tools.py
index 172924a..f598eb1 100644
--- a/src/management_tools.py
+++ b/src/management_tools.py
@@ -23,10 +23,18 @@ def remove_resources(user,resource,all):
         return custom_mongo.remove_resources()
     return "No action taken (missing --all flag?)"
 
+def import_database(res_path, rpt_path):
+    print("importing database....")
+    return custom_mongo.import_database(res_path, rpt_path)
+
+def export_database(res_path, rpt_path):
+    print("Exporting database....")
+    return custom_mongo.export_database(res_path, rpt_path)
+
 
 parser = argparse.ArgumentParser(description='Operational management of resources.')
 parser.add_argument('action', metavar='action', type=str,
-                    help='Operation to perform: list/remove')
+                    help='Operation to perform: list/remove import/export')
 parser.add_argument('-u',
                        '--user',
                        help='Filter action by user ID')
@@ -34,6 +42,14 @@ def remove_resources(user,resource,all):
                        '--resource',
                        help='Filter action by resource ID')
 
+parser.add_argument('-p',
+                       '--path',
+                       help='Select directory from where to save/read the resources collection')
+
+parser.add_argument('-t',
+                       '--token',
+                       help='Select directory from where to save/read the token collection')
+
 parser.add_argument('-a',
                        '--all',
                        action='store_true',
@@ -48,6 +64,10 @@ def remove_resources(user,resource,all):
     if args["resource"] is not None:
         args["all"] = False
     result = remove_resources(args['user'],args['resource'],args['all'])
+elif args["action"] == "import": 
+    result = import_database(args['path'], args['token'])
+elif args["action"] == "export": 
+    result = export_database(args['path'], args['token'])
 else:
     print("Allowed actions are 'remove' or 'list'")
     sys.exit(-1)