import psycopg2 import psycopg2.extras from loguru import logger import os import configparser import json import argparse import importlib import sys 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('./config/dbconfig.ini') DB_USER = config["database"]["user"] DB_PASS = config["database"]["pass"] DB_HOST = config["database"]["host"] DB_NAME = config["database"]["name"] parser = argparse.ArgumentParser(description="hv2cli.py") 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 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()