Update db_operations/db_write.py

db_operations/db_write.py

@@ -1,6 +1,12 @@
 import pymongo
 import os
-from config import DATABASE, COLLECTION, PREDICTION_DATABASE, PREDICTION_COLLECTION, COLLECT_PREDICTION_DATA
+from config import (PRODUCTION_URL, 
+                    PRODUCTION_DATABASE, 
+                    PRODUCTION_COLLECTION, 
+                    PREDICTION_URL,
+                    PREDICTION_DATABASE, 
+                    PREDICTION_COLLECTION, 
+                    COLLECT_PREDICTION_DATA)
 from logger import get_logger
 
 logger = get_logger()
@@ -10,11 +16,12 @@ class DBWrite:
     Inserts processed news into MongoDB
     """
     def __init__(self, db_type: str="production"):
-        self.url = os.getenv('DB_URL')
         self.db_type = db_type
-        self.database = ARCHIVE_DATABASE
-        self.collection = ARCHIVE_COLLECTION
+        self.url = PREDICTION_URL
+        self.database = PREDICTION_DATABASE
+        self.collection = PREDICTION_COLLECTION
         if self.db_type == "production":
+            self.url = PRODUCTION_URL
             self.database = DATABASE
             self.collection = COLLECTION