File size: 2,878 Bytes
c2c5fc6 164d548 82e716c c2c5fc6 82e716c c2c5fc6 f6a4997 164d548 f6a4997 164d548 7d8b789 f6a4997 c2c5fc6 22b2737 e3d1ff1 c2c5fc6 22b2737 f6a4997 22b2737 c2c5fc6 22b2737 c2c5fc6 f6a4997 c2c5fc6 f6a4997 c2c5fc6 f6a4997 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import pymongo
import os
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()
class DBWrite:
"""
Inserts processed news into MongoDB
"""
def __init__(self, db_type: str="production"):
self.db_type = db_type
self.url = PREDICTION_URL
self.database = PREDICTION_DATABASE
self.collection = PREDICTION_COLLECTION
if self.db_type == "production":
self.url = PRODUCTION_URL
self.database = PRODUCTION_DATABASE
self.collection = PRODUCTION_COLLECTION
self.__client = None
self.__error = 0
def __connect(self):
try:
self.__client = pymongo.MongoClient(self.url)
_ = self.__client.list_database_names()
except Exception as conn_exception:
self.__error = 1
self.__close_connection()
self.__client = None
raise
def __insert(self, documents):
try:
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)
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:
self.__client.close()
self.__client = None
def insert_news_into_db(self, documents: list):
logger.warning(f'Entering insert_news_into_db() : {self.db_type}')
if self.url is not None:
if self.__error == 0:
self.__connect()
if self.__error == 0:
self.__insert(documents=documents)
if self.__error == 0:
logger.warning(f"Insertion Successful: {self.db_type}. {len(documents)} documents inserted.")
if self.__client is not None:
self.__close_connection()
logger.warning(f'Exiting insert_news_into_db(): {self.db_type}') |