hv2-all-in-one/cli/hv2cli.py

import psycopg2
import psycopg2.extras
from loguru import logger
import os
import configparser
import json
import argparse
import importlib
import sys


parser = argparse.ArgumentParser(description="hv2cli.py")
parser.add_argument('--config', '-c',
                    help="Config file, default is ./config/dbconfig.ini",
                    required=False,
                    default="./config/dbconfig.ini")
parser.add_argument('--operation', '-o',
                    help='Operation to perform.',
                    required=True)
parser.add_argument('--params', '-p',
                    help='JSON string with parameter for the selected operation, default: {}',                    
                    required=False,
                    default="{}")
parser.add_argument('--verbosity', '-v',
                    help='Minimal log level for output: DEBUG, INFO, WARNING, ..., default: DEBUG',
                    required=False,
                    default="DEBUG")                    
parser.add_argument('--nocolorize', '-n',
                    help='disable colored output (for cron)',
                    required=False,
                    action='store_true',
                    default=False)


args = parser.parse_args()
operation = args.operation
params = json.loads(args.params)
logLevel = args.verbosity
noColorize = args.nocolorize

DB_USER = ""
DB_PASS = ""
DB_HOST = ""
DB_NAME = ""
try:
    DB_USER = os.environ["DB_USER"]
    DB_PASS = os.environ["DB_PASS"]
    DB_HOST = os.environ["DB_HOST"]
    DB_NAME = os.environ["DB_NAME"]
except KeyError:
    config = configparser.ConfigParser()
    config.read(args.config)
    DB_USER = config["database"]["user"]
    DB_PASS = config["database"]["pass"]
    DB_HOST = config["database"]["host"]
    DB_NAME = config["database"]["name"]


logger.remove()
logger.add(sys.stderr, colorize=(not noColorize), level=logLevel)


dbh = None
try:
    opMod = importlib.import_module(operation)

    dbh = psycopg2.connect(user = DB_USER, password = DB_PASS,
                           host = DB_HOST, database = DB_NAME,
                           sslmode = 'require')
    dbh.autocommit = False

    with dbh:
        opMod.perform(dbh, params)
except psycopg2.Error as err:
    raise Exception("Error when working on the database: {}".format(err))
except Exception as err:
    raise err
finally:
    if dbh:
        dbh.close()