From 8d2e756c50ddc933fef93b298ce5086fed3c9100 Mon Sep 17 00:00:00 2001 From: Jayasimha Raghavan <87547684+jayasimha-raghavan-unskript@users.noreply.github.com> Date: Wed, 13 Sep 2023 21:21:38 -0700 Subject: [PATCH] Changes to use globals in the config yaml file as inputParamsJson (#867) Co-authored-by: Jayasimha Raghavan --- unskript-ctl/db_utils.py | 40 ++++++ unskript-ctl/unskript-client.py | 135 ++++++++++++++++----- unskript-ctl/unskript_config.yaml.template | 7 ++ 3 files changed, 154 insertions(+), 28 deletions(-) create mode 100644 unskript-ctl/unskript_config.yaml.template diff --git a/unskript-ctl/db_utils.py b/unskript-ctl/db_utils.py index 5a62d3afd..6c5f7c768 100644 --- a/unskript-ctl/db_utils.py +++ b/unskript-ctl/db_utils.py @@ -173,6 +173,45 @@ def delete_pss_record(record_name: str, document_name: str) -> bool: return True +def get_checks_by_uuid(check_uuid_list: list): + """get_checks_by_uuid This function queries the snippets DB for + checks that match the give UUID and returns the checks. + + :type check_uuid_list: list + :param check_uuid_list: List of Check UUIDs + + :rtype: List of checks that match uuids + """ + if not check_uuid_list: + return None + + try: + db = DB(CS_DB_PATH) + except Exception as e: + raise e + tm = transaction.TransactionManager() + connection = db.open(tm) + root = connection.root() + cs = root.get('unskript_cs') + list_checks = [] + if cs is None: + raise Exception("Code Snippets Are missing") + for s in cs: + d = s + + if d.get('metadata').get('action_is_check') is False: + continue + c_uuid = d.get('uuid') + if c_uuid and c_uuid in check_uuid_list: + list_checks.append(d) + + tm.commit() + del root + connection.close() + db.close() + return list_checks + + def get_checks_by_connector(connector_name: str, full_snippet: bool = False): """get_checks_by_connector This function queries the snippets DB for checks of type connect and returns the checks. @@ -251,3 +290,4 @@ def get_creds_by_connector(connector_type: str): retval = (None, None) return retval + diff --git a/unskript-ctl/unskript-client.py b/unskript-ctl/unskript-client.py index c1b00f06a..4f2291825 100755 --- a/unskript-ctl/unskript-client.py +++ b/unskript-ctl/unskript-client.py @@ -18,17 +18,18 @@ import uuid import psutil import subprocess +import yaml +import nbformat +import ZODB +import ZODB.FileStorage +from pathlib import Path from datetime import datetime from argparse import ArgumentParser, REMAINDER from db_utils import * -import yaml -import nbformat from tabulate import tabulate from nbclient import NotebookClient from nbclient.exceptions import CellExecutionError from unskript.legos.utils import CheckOutputStatus -import ZODB -import ZODB.FileStorage from ZODB import DB # This python client can be used to @@ -43,9 +44,9 @@ # installed before using this script. # LIST OF CONSTANTS USED IN THIS FILE - +UNSKRIPT_GLOBALS = {} if os.environ.get('GLOBAL_CONFIG_PATH') is None: - GLOBAL_CONFIG_PATH="/unskript/data/unskript_config.yaml" + GLOBAL_CONFIG_PATH="/unskript/data/actions/unskript_config.yaml" CREDENTIAL_DIR="/.local/share/jupyter/metadata/credential-save" ZODB_DB_PATH="/var/unskript/snippets.db" @@ -70,16 +71,18 @@ def load_or_create_global_configuration(): and sets os.env variables which we shall use it in the subsequent functions. :rpath: None """ - global_content = {} + global UNSKRIPT_GLOBALS if os.path.exists(GLOBAL_CONFIG_PATH) is True: # READ EXISTING FILE AND SET ENV VARIABLES with open(GLOBAL_CONFIG_PATH, 'r') as f: - global_content = yaml.safe_load(f) + UNSKRIPT_GLOBALS = yaml.safe_load(f) - - for k, v in global_content.items(): - k = k.upper() - os.environ[k] = json.dumps(v) + if UNSKRIPT_GLOBALS.get('globals'): + for k, v in UNSKRIPT_GLOBALS.get('globals').items(): + os.environ[k] = json.dumps(v) + else: + _f_path = Path(GLOBAL_CONFIG_PATH) + _f_path.touch() def insert_first_and_last_cell(nb: nbformat.NotebookNode) -> nbformat.NotebookNode: @@ -104,6 +107,7 @@ def insert_first_and_last_cell(nb: nbformat.NotebookNode) -> nbformat.NotebookNo runbook_params = json.loads(os.environ.get('ACA_RUNBOOK_PARAMS')) runbook_variables = '' + if runbook_params: for k, v in runbook_params.items(): runbook_variables = runbook_variables + \ @@ -121,6 +125,14 @@ def insert_first_and_last_cell(nb: nbformat.NotebookNode) -> nbformat.NotebookNo paramsJson = json.dumps(paramDict) nbParamsObj = nbparams.NBParams(paramsJson) {runbook_variables} +''' + for k,v in UNSKRIPT_GLOBALS.get('globals').items(): + if isinstance(v,str) is True: + first_cell_content += f'{k} = \"{v}\"' + '\n' + else: + first_cell_content += f'{k} = {v}' + '\n' + + first_cell_content += f''' w = Workflow(env, secret_store_cfg, None, global_vars=globals(), check_uuids={ids})''' # Firstcell content. This is a static content @@ -226,6 +238,7 @@ def run_ipynb(filename: str, status_list_of_dict: list = None): client = NotebookClient(nb=nb, kernel_name="python3") try: + execution = client.execute() except CellExecutionError as e: raise e @@ -241,8 +254,8 @@ def run_ipynb(filename: str, status_list_of_dict: list = None): ids = get_code_cell_action_uuids(nb.dict()) result_table = [["Checks Name", "Result", "Failed Count", "Error"]] if len(outputs) == 0: - raise Exception( - "Unable to execute Runbook. Last cell content is empty") + print("OUTPUT for the Check Action is Empty") + return results = outputs[0] idx = 0 @@ -250,6 +263,7 @@ def run_ipynb(filename: str, status_list_of_dict: list = None): failed_result_available = False failed_result = {} + print(r) for result in r.split('\n'): if result == '': continue @@ -365,6 +379,37 @@ def run_checks(filter: str): print_run_summary(status_of_runs) +def run_suites(suite_name: str): + """run_suites This function takes the suite_name as an argument + and queries the DB and with given UUIDs and creates a temporary + runnable runbooks and run them, updates the audit results. + + :type suite_name: string + :param suite_name: Suite Name string, suite name should match the + content under `suites` section of the unskript_config.yaml + + :rtype: None + """ + if suite_name in ("", None): + raise Exception("Run Suite needs suite_name to be specified.") + + runbooks = [] + + check_list = [] + if UNSKRIPT_GLOBALS.get('suites') and UNSKRIPT_GLOBALS.get('suites').get(suite_name): + check_list = get_checks_by_uuid(UNSKRIPT_GLOBALS.get('suites').get(suite_name)) + + if len(check_list) > 0: + runbooks.append(create_jit_runbook(check_list)) + + status_of_runs = [] + for rb in runbooks: + run_ipynb(rb, status_of_runs) + + update_audit_trail(status_of_runs) + print_run_summary(status_of_runs) + + def print_run_summary(status_list_of_dict): """print_run_summary This function is used to just print the Run Summary. :type status_list_of_dict: list @@ -484,6 +529,29 @@ def update_current_execution(status, id: str, content: dict): if os.path.exists(failed_runbook) is not True: print(f"ERROR Unable to create failed runbook at {failed_runbook}") +def replace_input_with_globals(inputSchema: str): + if not inputSchema: + return None + + input_json_start_line = ''' +task.configure(inputParamsJson=\'\'\'{ + ''' + input_json_end_line = '''}\'\'\') + ''' + input_json_line = '' + try: + schema = inputSchema[0] + if schema.get('properties'): + for key in schema.get('properties').keys(): + if key in UNSKRIPT_GLOBALS.get('globals').keys(): + input_json_line += f"\"{key}\": \"{key}\" ," + except Exception as e: + print(f"EXCEPTION {e}") + pass + + retval = input_json_start_line + input_json_line.rstrip(',') + '\n' + input_json_end_line + return retval + def create_jit_runbook(check_list: list): """create_jit_runbook This function creates Just In Time runbook @@ -499,20 +567,26 @@ def create_jit_runbook(check_list: list): :rtype: None """ nb = nbformat.v4.new_notebook() - failed_notebook = os.environ.get( - 'EXECUTION_DIR', '/unskript').strip('"') + '/workspace/' + str(uuid.uuid4()) + '.ipynb' + failed_notebook = os.environ.get('EXECUTION_DIR', '/unskript/data').strip('"') + '/workspace/' + str(uuid.uuid4()) + '.ipynb' for check in check_list: s_connector = check.get('metadata').get('action_type') s_connector = s_connector.replace('LEGO', 'CONNECTOR') cred_name, cred_id = get_creds_by_connector(s_connector) + # No point proceeding further if the Credential is incomplete + if cred_name is None or cred_id is None: + print('\x1B[1;20;46m' + f"~~ Skipping {check.get('name')} {cred_name} {cred_id} ~~" + '\x1B[0m') + continue + + task_lines = ''' task.configure(printOutput=True) task.configure(credentialsJson=\'\'\'{ \"credential_name\":''' + f" \"{cred_name}\"" + ''', - \"credential_type\":''' + f" \"{s_connector}\"" + ''', - \"credential_id\":''' + f" \"{cred_id}\"" + '''}\'\'\') + \"credential_type\":''' + f" \"{s_connector}\"" + '''}\'\'\') ''' - + input_json = replace_input_with_globals(check.get('inputschema')) + if input_json: + task_lines += input_json try: c = check.get('code') idx = c.index("task = Task(Workflow())") @@ -526,7 +600,7 @@ def create_jit_runbook(check_list: list): check['code'] = [] for line in c[:]: check['code'].append(str(line + "\n")) - + check['metadata']['action_uuid'] = check['uuid'] check['metadata']['name'] = check['name'] cc = nbformat.v4.new_code_cell(check.get('code')) @@ -698,7 +772,7 @@ def display_failed_checks(connector: str = ''): def display_failed_logs(exec_id: str = None): - output = os.environ.get('EXECUTION_DIR', '/unskript/execution').strip( + output = os.environ.get('EXECUTION_DIR', '/unskript/data/execution').strip( '"') + '/workspace/' + f"{exec_id}_output.ipynb" if not os.path.exists(output): print( @@ -971,18 +1045,18 @@ def create_creds_mapping(): :rtype: None """ - creds_files = os.environ.get('HOME').strip( - '"') + CREDENTIAL_DIR + '/*.json' + creds_files = os.environ.get('HOME').strip('"') + CREDENTIAL_DIR + '/*.json' list_of_creds = glob.glob(creds_files) d = {} for creds in list_of_creds: with open(creds, 'r') as f: c_data = json.load(f) - d[c_data.get('metadata').get('type')] = {"name": c_data.get( - 'metadata').get('name'), "id": c_data.get('metadata').get('id')} - if not c_data.get('metadata').get('connectorData'): - print(f"ERROR: The Credential data for {c_data.get('metadata').get('type')} is empty!" + if c_data.get('metadata').get('connectorData') == "{}": + print(f"WARNING: The Credential data for {c_data.get('metadata').get('type')} is empty!" " Please use unskript-ctl.sh -cc option to create the credential") + continue + d[c_data.get('metadata').get('type')] = {"name": c_data.get('metadata').get('name'), + "id": c_data.get('id')} upsert_pss_record('default_credential_id', d, False) @@ -1428,7 +1502,7 @@ def stop_debug(): if __name__ == "__main__": try: if os.environ.get('EXECUTION_DIR') is None: - os.environ['EXECUTION_DIR'] = '/unskript/execution' + os.environ['EXECUTION_DIR'] = '/unskript/data/execution' create_creds_mapping() load_or_create_global_configuration() @@ -1449,6 +1523,8 @@ def stop_debug(): help='Run the given runbook FILENAME [-RUNBOOK_PARM1 VALUE1] etc..') parser.add_argument('-rc', '--run-checks', type=str, help='Run all available checks [all | connector | failed]') + # parser.add_argument('-rs', '--run-suites', type=str, + # help='Run Health Check Suites (as defined in the unskript_config.yaml file)') parser.add_argument('-df', '--display-failed-checks', help='Display Failed Checks [all | connector]') parser.add_argument('-lc', '--list-checks', type=str, @@ -1483,6 +1559,8 @@ def stop_debug(): parse_runbook_param(args.run_runbook) elif args.run_checks not in ('', None): run_checks(args.run_checks) + # elif args.run_suites not in ('', None): + # run_suites(args.run_suites) elif args.display_failed_checks not in ('', None): display_failed_checks(args.display_failed_checks) elif args.list_checks not in ('', None): @@ -1504,3 +1582,4 @@ def stop_debug(): stop_debug() else: parser.print_help() + diff --git a/unskript-ctl/unskript_config.yaml.template b/unskript-ctl/unskript_config.yaml.template new file mode 100644 index 000000000..cb3711147 --- /dev/null +++ b/unskript-ctl/unskript_config.yaml.template @@ -0,0 +1,7 @@ +globals: + namespace: "awesome-ops" + threshold: "string" + region: "us-west-2" + services: ["calendar", "audit"] + +