read_predictions_from_db.py

import pymongo
import os
import pandas as pd
import logging


class PredictionDBRead:
    def __init__(self):
        self.url = os.getenv('PREDICTION_DB_URL')
        self.database = "indian_news_predictions_db"
        self.collection = "indian_news_predictions_coll"
        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.__client = None


    def __read(self):
        try:
            db = self.__client[self.database]
            coll = db[self.collection]
            docs = []
            for doc in coll.find():
                docs.append(doc)
            rss_df = pd.DataFrame(docs)
        except Exception as insert_err:
            self.__error = 1
            rss_df = None
        return rss_df

    
    def __close_connection(self):
        if self.__client is not None:
            self.__client.close()
            self.__client = None

    
    def read_predictions_from_db(self):
        rss_df = None
        if self.url is not None:
            if self.__error == 0:
                self.__connect()
            if self.__error == 0:
                rss_df = self.__read()
            if self.__error == 0:
                logging.info("Prediction Data Read Successful")
            if self.__client is not None:
                self.__close_connection()
        return rss_df