Source code for personalcapital.database

import structlog

from pymongo import MongoClient

log = structlog.getLogger()


class Database(object):
    DB_PORT = 27017
    DB_HOST = 'localhost'
    DATE_FORMAT = '%Y-%m-%d'

    def __init__(self, host=DB_HOST, port=DB_PORT):
        self.host = host
        self.port = port
        self.conn = None

    def __enter__(self):
        self.conn = MongoClient(host=self.host, port=self.port)
        return self.conn

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.conn.close()


def most_recent_transaction():
    with Database() as db:
        txns = db.banking.transactions
        most_recent = txns.find(
            {},
            {"transactionDate": 1}
        ).sort("transactionDate", -1).limit(1)
    return next(most_recent)


def insert_transactions(txns):
    if len(txns) > 0:
        with Database() as db:
            log.msg("Inserting {} transactions.".format(len(txns)))
            db.banking.transactions.insert_many(txns)
    else:
        log.msg("Not inserting new transactions, none given.")


def has_transactions(txns):
    txn_map = {
        txn["userTransactionId"]: txn
        for txn in txns
    }
    with Database() as db:
        matches = db.banking.transactions.find(
            {
                "userTransactionId": { "$in": list(txn_map.keys()) }
            },
            {"userTransactionId": 1}
        )
    return [match["userTransactionId"] in txn_map for match in matches]