db.py

import os
import json
import logging
import firebase_admin
from firebase_admin import credentials, firestore

public_creds_path = 'public_creds.json'

def load_credentials():
    try:
        with open(public_creds_path) as f:
            credentials_dict = json.load(f)
        secret = {
            'private_key_id': os.environ.get('private_key_id'),
            'private_key': os.environ.get('private_key').replace(r'\n', '\n')
        }
        credentials_dict.update(secret)
        return credentials_dict
    except Exception as e:            
        logging.error(f'Error while loading credentials: {e}')
        return None
    
def connect_to_db(credentials_dict):
    try:
        cred = credentials.Certificate(credentials_dict)
        if not firebase_admin._apps:
            firebase_admin.initialize_app(cred)
        logging.info('Established connection to db!')
        return firestore.client()
    except Exception as e:
        logging.error(f'Error while connecting to db: {e}')
        return None
    
def write_to_db(db, payload):
    
    try:
        collection_ref = db.collection('SyntheticCorrelations')
        doc_ref = collection_ref.document('UserInputs')
        doc = doc_ref.get()

        if doc.exists:
            doc_ref.update({
                'Data': firestore.ArrayUnion([payload])
            })
        else:
            doc_ref.set({
                'Data': [payload]
            })
        logging.info(f'Sent payload to db: {payload}')
        return True
    except Exception as e:
        logging.error(f'Error while sending payload to db: {e}')
        return False