db-manager/src/mongo.py

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

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: Database URI: {db_host}:{db_port}", "DBUG")
client = MongoClient(f"mongodb://{db_user}:{db_pass}@{db_host}:{db_port}")
db = client[db_name]

log.out("mongo.py: Mongo module loaded", "DBUG")

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', "EROR")
    return 0

def metaUpdateTime(target):
    col = db["meta"]
    log.out(f'mongo.metaUpdateTime: Updating updated time for {target}', "DBUG")
    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 createTtlIndex(collection, field, time):
    col = db[collection]
    col.create_index(field, expireAfterSeconds = time)
    log.out(f'mongo.createTtlIndex: Created TTL Index of {field} in {collection} to expire after {time} seconds', "INFO")

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 putBulkPis(data):
    collection = "pis"
    startCount = getLength(collection)
    col = db[collection]
    incrementCounter(collection)
    if startCount > 0:
        log.out(f'mongo.putBulkPid: Dropping {startCount} pis documents', "INFO")
        col.drop()
    col.insert_many(data)
    endCount = getLength(collection)
    log.out(f'mongo.putBulkPis: {endCount} documents inserted', "INFO")
    log.out(f'mongo.putBulkPis: {endCount-startCount} new documents', "INFO")
    log.out('mongo.putBulkPis: Updating meta time', "INFO")
    metaUpdateTime(collection)
    return

def putMany(collection :str, data :list):
  log.out(f"mongo.putMany: Inserting many documents to: {collection}")
  col = db[collection]
  incrementCounter(collection)
  col.insert_many(data)
  metaUpdateTime(collection)


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"})
    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

def putVersion(version):
    collection = "meta"
    col = db[collection]
    res = col.update_one({"target": "versions"}, {"$set":{"target": "versions","dbmanager": version}}, upsert=True)
    return

def putTimetable(data):
    collection = "timetable"
    col = db[collection]
    res = col.insert_many(data)

def dropCollection(collection):
    log.out(f"mongo.dropCollection: Dropping collection '{collection}'")
    col = db[collection]
    res = col.drop()

def deleteTimetableData(query):
    collection = "timetable"
    col = db[collection]
    res = col.delete_one(query)

def getMetaHash(target):
    collection = "meta"
    col = db[collection]
    res = col.find_one({"target": target, "type": "collection"})
    incrementCounter("meta")
    if type(res) is dict:
      if 'updated' in res:
        try:
          log.out(f'mongo.metaGetHash: {target} hash is {res["hash"]}', 'INFO')
          return res["hash"]
        except:
          return None
    return None

def putMetaHash(target :str, hash :str):
  collection = "meta"
  col = db[collection]
  filter = {
    "target": target,
    "type": "collection"
  }
  update = {
    "target": target,
    "type": "collection",
    "hash": hash
  }
  res = col.update_one(filter, {"$set": update}, upsert=True)

def query(collection, query):
  col = db[collection]
  log.out(f"mongo.query: Running query: {query}")
  return col.find_one(query)