db-manager/src/mongo.py

import os
from pymongo import MongoClient
import time
import urllib.parse
import logger as log

log.out("mongo.py: Fetching configuration", "INFO")
db_host = os.getenv('OWL_DB_HOST', 'localhost')
db_port = os.getenv('OWL_DB_PORT', 27017)
db_user = urllib.parse.quote_plus(os.getenv('OWL_DB_USER', "owl"))
db_pass = urllib.parse.quote_plus(os.getenv('OWL_DB_PASS', "twittwoo"))
db_name = os.getenv('OWL_DB_NAME', "owlboard")

log.out(f"mongo.py: Connecting to database at {db_host}:{db_port}", "INFO")
client = MongoClient(f"mongodb://{db_user}:{db_pass}@{db_host}:{db_port}")
db = client[db_name]

def metaCheckTime(target):
    col = db["meta"]
    res = col.find_one({"target": target, "type": "collection"})
    incrementCounter("meta")
    if type(res) is dict:
      if 'updated' in res:
        log.out(f'mongo.metaUpdateTime: {target} last updated at {res["updated"]}', "INFO")
        return res["updated"]
    log.out(f'mongo.metaUpdatetime: {target} does not exist', "INFO")
    return 0

def metaUpdateTime(target):
    col = db["meta"]
    log.out(f'mongo.metaUpdateTime: Updating updated time for {target}', "INFO")
    res = col.update_one({"target": target, "type":"collection"}, {"$set":{"updated": int(time.time()),"target":target, "type":"collection"}}, upsert=True)
    incrementCounter("meta")

def getLength(collection):
    col = db[collection]
    incrementCounter(collection)
    return col.count_documents({})

def createSingleIndex(collection, field):
    col = db[collection]
    col.create_index(field)
    incrementCounter(collection)
    log.out(f'mongo.createSingleIndex: Created index of {field} in {collection}', "INFO")
    return

def putBulkCorpus(data):
    collection = "corpus"
    startCount = getLength(collection)
    col = db[collection]
    incrementCounter(collection)
    if startCount > 0:
        log.out(f'mongo.putBulkCorpus: Dropping {startCount} CORPUS documents', "INFO")
        col.drop()
    col.insert_many(data)
    endCount = getLength(collection)
    log.out(f'mongo.putBulkCorpus: {endCount} documents inserted', "INFO")
    log.out(f'mongo.putBulkCorpus: {endCount - startCount} new documents', "INFO")
    log.out('mongo.putBulkCorpus: Building collection indexes',"INFO")
    createSingleIndex(collection, "NLC")
    createSingleIndex(collection, "3ALPHA")
    log.out('mongo.putBulkCorpus: Updating meta time',"INFO")
    metaUpdateTime(collection)
    return

def putBulkStations(data):
    collection = "stations"
    startCount = getLength(collection)
    col = db[collection]
    incrementCounter("stations")
    if startCount > 0:
        log.out(f'mongo.putBulkStations: Dropping {startCount} station documents', "INFO")
        col.drop()
    col.insert_many(data)
    endCount = getLength(collection)
    log.out(f'mongo.putBulkStations: {endCount} documents inserted', "INFO")
    log.out(f'mongo.putBulkStations: {endCount - startCount} new documents', "INFO")
    log.out('mongo.putBulkStations: Building collection indexes',"INFO")
    createSingleIndex(collection, "3ALPHA")
    createSingleIndex(collection, "STANOX")
    createSingleIndex(collection, "TIPLOC")
    log.out('mongo.putBulkStations: Updating meta time',"INFO")
    metaUpdateTime(collection)
    return

def incrementCounter(target):
    collection = "meta"
    col = db[collection]
    log.out(f'mongo.incrementCounter: Incrementing counter for {target}', "INFO")
    col.update_one({"target": "counters","type": "count"}, {"$inc":{target: 1}})
    return

def metaCounters():
    collection = "meta"
    col = db[collection]
    res = col.find_one({"target": "counters","type": "count"})
    log.out(f'mongo.metaCounters: Query returned `{res}`', "DEBG")
    if type(res) is dict:
      if 'since' in res:
        log.out('mongo.metaCounters: counters already exists, skipping', "INFO")
        incrementCounter(collection)
        return
    log.out('mongo.metaCounters: counters does not exist, creating', "INFO")
    col.update_one({"target": "counters","type": "count"}, {"$set":{"target": "counters","type": "count","since": int(time.time())}}, upsert=True)
    incrementCounter(collection)
    return