Explorar el Código

Refactor update method and add variable for advanced update queries

ogert hace 4 años
padre
commit
f95102a0ab
Se han modificado 1 ficheros con 16 adiciones y 13 borrados
  1. 16 13
      cdplib/db_handlers/MongodbHandler.py

+ 16 - 13
cdplib/db_handlers/MongodbHandler.py

@@ -469,31 +469,34 @@ class MongodbHandler:
             self._log.warning('No data was found for the query')
             return None
 
-    def update_data_in_collection(self, update_label:str, update_value: str, collection_name:str, query_label: str = None, query_value: str = None, create_if_not_exist: bool = True, find_query: dict = None, update_many: bool = False):
+    def update_data_in_collection(self, update_label:str, update_value: str, collection_name:str, query_label: str = None, query_value: str = None, create_if_not_exist: bool = True, find_query: dict = None, update_query:dict = None, update_many: bool = False):
 
         if isinstance(update_value, pd.DataFrame):
             update_value = simplejson.loads(update_value.to_json(orient="records",
                                                  date_format="iso"))
+
+        if update_query is None:
+            update_query = {"$set": {update_label: update_value}}
+
+        if find_query is None:
+            if query_label and query_value:
+                find_query = {query_label:query_value}
+                
+                                             
         try:
             if update_many:
-                if query_label and query_value:
-                    self._database[collection_name].update_many({query_label:query_value}, {"$set": {update_label: update_value}}, upsert=create_if_not_exist)
-
-                elif find_query:
-                    self._database[collection_name].update_many(find_query, {"$set": {update_label: update_value}}, upsert=create_if_not_exist)
+                if find_query is not None:
+                    self._database[collection_name].update_many(find_query, update_query, upsert=create_if_not_exist)
 
                 else:
-                    self._database[collection_name].update_many({}, {"$set": {update_label: update_value}}, upsert=create_if_not_exist)
+                    self._database[collection_name].update_many({}, update_query, upsert=create_if_not_exist)
 
             else:
-                if query_label and query_value:
-                    self._database[collection_name].update_one({query_label:query_value}, {"$set": {update_label: update_value}}, upsert=create_if_not_exist)
-
-                elif find_query:
-                    self._database[collection_name].update_one(find_query, {"$set": {update_label: update_value}}, upsert=create_if_not_exist)
+                if find_query is not None:
+                    self._database[collection_name].update_one(find_query, update_query, upsert=create_if_not_exist)
 
                 else:
-                    self._database[collection_name].update_one({}, {"$set": {update_label: update_value}}, upsert=create_if_not_exist)
+                    self._database[collection_name].update_one({}, update_query, upsert=create_if_not_exist)
 
             self._log.info(('Data for label: {} has been updated').format(update_label))