# Copyright (C) 2005-2024 Splunk Inc. All Rights Reserved. import os import json import time from threading import Thread from queue import Queue import splunk.rest as rest from itsi.upgrade.itsi_migration import ItsiMigration from itsi.itsi_utils import ITOAInterfaceUtils from ITOA.itoa_common import get_conf_stanza_single_entry from .constants import NEW_VERSION from itsi.upgrade.constants import PREP_TIMEOUT, TRANSFORM_TIMEOUT from itsi.upgrade.precheck.migration_precheck_resolver import MigrationPreCheckResolver from itsi.upgrade.itsi_migration_log import PrefixLogger class MigrationQueueOperation(object): KV_STORE_MIGRATION_QUEUE_COLLECTION_URI = '/servicesNS/nobody/SA-ITOA/storage/collections/data/itsi_migration_queue' ITOA_INTERFACE_URI = '/servicesNS/nobody/SA-ITOA/itoa_interface/' def __init__(self, session_key, logger): ''' Constructor @type: string @param session_key: @type: object @param logger: itsi_upgrade_queue logger @rtype: object @return: None ''' self._session_key = session_key self.logger = logger self.ui_logger = PrefixLogger(prefix='UI', logger=self.logger) def _run_migration_pre_checks_for_versions(self, old_version, new_version, skip_pre_checks=[]): """ Run migration prechecks given the version being migrated from and the version being migrated to. An optional list of precheck IDs to be skipped can be provided. :param old_version: the version being migrated from :param new_version: the version being migrated to :param skip_pre_checks: the list of precheck IDs to be skipped :return: None """ is_precheck_good = True pre_check_results = [] if old_version: # Resolve prechecks to run given old and new versions pre_check_resolver = MigrationPreCheckResolver() resolved_pre_checks_specs = pre_check_resolver.resolve_pre_check_specs(old_version, new_version) # Save initialized default precheck status to migration status pre_check_results = {} for pre_checks_spec in resolved_pre_checks_specs: for pre_check in pre_checks_spec['pre_checks']: pre_check_results[pre_check['id']] = pre_check pre_check_results[pre_check['id']]['status'] = 'Enqueued' pre_check_results[pre_check['id']]['recommendation'] = [] pre_check_results[pre_check['id']]['type'] = [] entry = ITOAInterfaceUtils.get_migration_status_from_kv( self._session_key) ITOAInterfaceUtils.append_data_to_migration_status_kv(self._session_key, entry, precheck_results=pre_check_results, precheck_start_timestamp=time.time()) self.logger.info("Resolved prechecks (old version: %s -> new version: %s) are: %s", old_version, new_version, pre_check_results) # Instantiate prechecks pre_checks_instances = [ pre_checks_spec['type'](self._session_key, self.logger, pre_checks_spec['pre_checks'], skip_pre_checks) for pre_checks_spec in resolved_pre_checks_specs ] # Run prechecks and save results to migration status for pre_check_instance in pre_checks_instances: pre_check_name = pre_check_instance.__class__.__name__ self.ui_logger.info("Running prechecker: %s", pre_check_name) prechecks_in_pre_check_instance = pre_check_instance.pre_checks # First set precheck status to in progress for precheck in prechecks_in_pre_check_instance: entry = ITOAInterfaceUtils.get_migration_status_from_kv( self._session_key) existing_precheck_results = entry.get('precheck_results') existing_precheck_results[precheck['id']]['status'] = 'In Progress' ITOAInterfaceUtils.append_data_to_migration_status_kv(self._session_key, entry, precheck_results=existing_precheck_results) # Then run the prechecks in the instance pre_check_instance.run() entry = ITOAInterfaceUtils.get_migration_status_from_kv( self._session_key) existing_precheck_results = entry.get('precheck_results') precheck_ids = pre_check_instance.pre_check_timestamps.keys() # First update pre_check_results with timestamps and set default status to Failed for precheck_id in precheck_ids: pre_check_results[precheck_id].update(pre_check_instance.pre_check_timestamps[precheck_id]) pre_check_results[precheck_id]['status'] = 'Failed' # Get actual status of prechecks and save to migration _tatus in kv for pre_check_result in pre_check_instance.pre_check_results: id = pre_check_result['id'] recommendation = pre_check_result['recommendation'] passed = pre_check_result['passed'] message_type = pre_check_result['type'] if passed: pre_check_results[id]['status'] = 'Completed' pre_check_results[id]['recommendation'].append(recommendation) pre_check_results[id]['type'].append(message_type) existing_precheck_results.update(pre_check_results) ITOAInterfaceUtils.append_data_to_migration_status_kv(self._session_key, entry, precheck_results=existing_precheck_results) self.ui_logger.info("Completed running prechecker: %s with results: %s", pre_check_name, pre_check_instance.pre_check_results) # Validate precheck results is_precheck_good = self.validate_prechecks(pre_check_results.values(), skip_pre_checks) return is_precheck_good, pre_check_results def _run_migration_pre_checks(self, queue, skip_pre_checks=[]): """ Run migration prechecks that are applicable for the version being migrated from and the version being migrated to. :type: queue :param queue: store the pre-check result in the queue in form of dict :type: list :param skip_pre_checks: ids of prechecks to skip :return: None """ results_dict = queue.get() old_version, id_ = ITOAInterfaceUtils.get_version_from_kv(self._session_key) results_dict["is_precheck_good"], results_dict["pre_check_results"] = self._run_migration_pre_checks_for_versions(old_version, NEW_VERSION, skip_pre_checks) queue.put(results_dict) def _get_kv_object_count(self): objects_to_be_counted = ['glass_table', 'service', 'kpi_threshold_template', 'entity'] object_count = {} for object_to_be_counted in objects_to_be_counted: getargs = {'output_mode': 'json'} uri = MigrationQueueOperation.ITOA_INTERFACE_URI + object_to_be_counted + '/count' try: resp, content = rest.simpleRequest(uri, sessionKey=self._session_key, getargs=getargs) except Exception as e: self.logger.exception(e) raise Exception(e) if resp.status != 200 and resp.status != 201: object_count[object_to_be_counted] = '' json_data = json.loads(content) object_count[object_to_be_counted] = json_data.get('count', '') return object_count def run_migration_task(self, skip_local_failure, queue): results_dict = queue.get() worker = ItsiMigration(self._session_key) results_dict["is_migration_successful"] = worker.run_migration(skip_local_failure) queue.put(results_dict) def execute_migration_queue(self): ''' This is invoked by a modular input. Trigger migration if migration queue being put into collection itsi_migration_queue @rtype: None @return: None ''' prep_timeout = int(get_conf_stanza_single_entry( self._session_key, 'itsi_settings', 'upgrade_timeouts', 'precheck_timeout').get('content', PREP_TIMEOUT)) transform_timeout = int(get_conf_stanza_single_entry( self._session_key, 'itsi_settings', 'upgrade_timeouts', 'migration_timeout').get('content', TRANSFORM_TIMEOUT)) migration_queue_info = self.get_migration_queue_info() is_migration_triggered = migration_queue_info.get('is_migration_triggered') if is_migration_triggered: entry = ITOAInterfaceUtils.get_migration_status_from_kv( self._session_key) ITOAInterfaceUtils.append_data_to_migration_status_kv(self._session_key, entry, migration_timeout=None, upgrade_timeout=None, precheck_timeout=None, ERROR=None) skip_local_failure = migration_queue_info.get('skip_local_failure') skip_pre_checks = migration_queue_info.get('skip_pre_checks') is_precheck_good = False pre_check_results = [] try: queue = Queue() results_dict = {"is_precheck_good" : False, "pre_check_results" : ""} queue.put(results_dict) precheck_thread = Thread(name='PrecheckThread', target=self._run_migration_pre_checks, args=(queue, skip_pre_checks)) precheck_thread.start() precheck_thread.join(timeout=prep_timeout) if precheck_thread.is_alive(): precheck_timeout_dict = {} precheck_timeout_dict["message"] = "Could not complete upgrade because the upgrade \ prechecks took too long to complete. Restart the upgrade and try again." precheck_timeout_dict["timeout"] = prep_timeout entry = ITOAInterfaceUtils.get_migration_status_from_kv( self._session_key) ITOAInterfaceUtils.append_data_to_migration_status_kv(self._session_key, entry, skip_local_failure=skip_local_failure, precheck_timeout=precheck_timeout_dict, end_timestamp=time.time()) self.clear_migration_queue() self.logger.exception("Timeout occurred at precheck stage") # Need to terminate the process for terminate the execution of function run by the thread os._exit(0) else: results_dict = queue.get() is_precheck_good = results_dict.get("is_precheck_good") pre_check_results = results_dict.get("pre_check_results") self.logger.info("Precheck is good: %s and results obtained are: %s", is_precheck_good, pre_check_results) self.clear_migration_queue() entry = ITOAInterfaceUtils.get_migration_status_from_kv( self._session_key) ITOAInterfaceUtils.append_data_to_migration_status_kv(self._session_key, entry, skip_local_failure=skip_local_failure, end_timestamp=time.time()) except Exception as exception: precheck_error_dict = {"message" : "An error occurred while executing prechecks", "type" : "PRECHECK"} self.logger.error(f"An error occurred while executing prechecks {exception}") entry = ITOAInterfaceUtils.get_migration_status_from_kv( self._session_key) ITOAInterfaceUtils.append_data_to_migration_status_kv(self._session_key, entry, skip_local_failure=skip_local_failure, ERROR=precheck_error_dict, end_timestamp=time.time()) self.clear_migration_queue() raise Exception(exception) if is_precheck_good: is_migration_successful = False try: precheck_count = self._get_kv_object_count() queue = Queue() results_dict = {"is_migration_successful" : False} queue.put(results_dict) migration_thread = Thread(name='MigrationThread', target=self.run_migration_task, args=(skip_local_failure, queue)) migration_thread.start() migration_thread.join(timeout=transform_timeout) if migration_thread.is_alive(): migration_timeout_dict = {} migration_timeout_dict["message"] = "Could not complete upgrade because the \ migration took too long to complete. Restart the upgrade and try again." migration_timeout_dict["timeout"] = transform_timeout entry = ITOAInterfaceUtils.get_migration_status_from_kv( self._session_key) ITOAInterfaceUtils.append_data_to_migration_status_kv(self._session_key, entry, is_running=False, end_timestamp=time.time(), has_succeeded=False, migration_timeout=migration_timeout_dict) self.logger.exception("Timeout occurred at transform stage") # Need to terminate the process for terminate the execution of function run by the thread os._exit(0) else: results_dict = queue.get() is_migration_successful = results_dict.get("is_migration_successful") if is_migration_successful: postcheck_count = self._get_kv_object_count() entry = ITOAInterfaceUtils.get_migration_status_from_kv( self._session_key) ITOAInterfaceUtils.append_data_to_migration_status_kv(self._session_key, entry, precheck_count=precheck_count, postcheck_count=postcheck_count) except Exception as exception: migration_error_dict = {"message" : "An error occurred while doing migration", "type" : "MIGRATION"} self.logger.error(f"An error occurred while doing migration {exception}") entry = ITOAInterfaceUtils.get_migration_status_from_kv( self._session_key) ITOAInterfaceUtils.append_data_to_migration_status_kv(self._session_key, entry, is_running=False, end_timestamp=time.time(), has_succeeded=False, ERROR=migration_error_dict) raise Exception(exception) def get_migration_queue_info(self): ''' Get content from collection itsi_migration_queue if content is not empty, migration has been triggered by user @rtype: dict {bool} whether migration has been triggered {bool} whether to skip local failures ''' try: rsp, content = rest.simpleRequest(MigrationQueueOperation.KV_STORE_MIGRATION_QUEUE_COLLECTION_URI, sessionKey=self._session_key, raiseAllErrors=False) if rsp.status != 200 and rsp.status != 201: return {'is_migration_triggered': False, 'skip_local_failure': None, 'skip_pre_checks': []} json_data = json.loads(content) except Exception as e: self.logger.error(e) if len(json_data) == 0: return {'is_migration_triggered': False, 'skip_local_failure': None, 'skip_pre_checks': []} entry = json_data[0] skip_local_failure = entry.get('skip_local_failure', False) skip_pre_checks = entry.get('skip_pre_checks', []) return {'is_migration_triggered': True, 'skip_local_failure': skip_local_failure, 'skip_pre_checks': skip_pre_checks} def clear_migration_queue(self): ''' Clean up collection itsi_migration_queue @rtype: None @return: None ''' try: rest.simpleRequest(MigrationQueueOperation.KV_STORE_MIGRATION_QUEUE_COLLECTION_URI, sessionKey=self._session_key, raiseAllErrors=False, method='DELETE') except Exception as e: self.logger.exception(e) def validate_prechecks(self, precheck_results, skip_pre_checks): ''' Check if all prechecks are either passed or skipped. @type:list @param: precheck_results @rtype: {bool} @return:whether to skip precheck failures ''' for precheck in precheck_results: if precheck['status'] == 'Failed' and precheck['id'] not in skip_pre_checks: return False return True