Update db_operations/db_write.py

db_operations/db_write.py

@@ -36,6 +36,7 @@ class DBWrite:
             self.__error = 1
             self.__close_connection()
             self.__client = None
+            raise
         
 
     def __insert(self, documents):
@@ -43,13 +44,26 @@ class DBWrite:
 
             db = self.__client[self.database]
             coll = db[self.collection]
+            # if (self.db_type == "production") or (COLLECT_PREDICTION_DATA==0):
+            #     coll.drop()
+            # coll.insert_many(documents=documents)
+
+
             if (self.db_type == "production") or (COLLECT_PREDICTION_DATA==0):
                 coll.drop()
-            coll.insert_many(documents=documents)
+                coll.insert_many(documents=documents)
+            else:
+                for doc in documents:
+                    filter_query = {"url": doc["url"]}
+                    update_query = {"$set": doc}
+                    coll.update_one(filter_query, update_query, upsert=True)
+
+            
         except Exception as insert_err:
             self.__error = 1
             self.__close_connection()
-            
+            logger.critical(f'Error while inserting into DB: {insert_error}')
+            raise
 
     def __close_connection(self):
         if self.__client is not None: