From 899dfe56885fc35fa9454f8929d988b9705e9b7e Mon Sep 17 00:00:00 2001 From: tomerlf1 Date: Wed, 13 Dec 2023 11:38:23 +0200 Subject: [PATCH] task wrapper decorator --- gce_rescue/gce.py | 7 +++++++ gce_rescue/messages.py | 2 +- gce_rescue/tasks/actions.py | 4 ++-- gce_rescue/tasks/disks.py | 6 +++++- gce_rescue/tasks/metadata.py | 3 +++ gce_rescue/tasks/operations.py | 4 ++++ gce_rescue/utils.py | 22 +++++++++++++++++++++- 7 files changed, 43 insertions(+), 5 deletions(-) diff --git a/gce_rescue/gce.py b/gce_rescue/gce.py index aac89d2..f1295d5 100644 --- a/gce_rescue/gce.py +++ b/gce_rescue/gce.py @@ -14,6 +14,7 @@ """ Initilization Instance() with VM information. """ import sys +import logging from googleapiclient.discovery import Resource from googleapiclient.errors import HttpError @@ -26,6 +27,7 @@ from gce_rescue.tasks.pre_validations import Validations from gce_rescue.config import get_config +_logger = logging.getLogger(__name__) def get_instance_info( compute: Resource, @@ -50,6 +52,9 @@ def guess_guest(data: Dict) -> str: Default: projects/debian-cloud/global/images/family/debian-11""" guests = get_config('source_guests') + if not data.get('disks'): + _logger.error(f'Unable to get disks for vm. Check whether a boot disk is attached to your vm.') + raise Exception("No boot disk was found for vm") for disk in data['disks']: if disk['boot']: if 'architecture' in disk: @@ -58,6 +63,8 @@ def guess_guest(data: Dict) -> str: arch = 'x86_64' guest_default = guests[arch][0] guest_name = guest_default.split('/')[-1] + if not disk.get('licenses'): + return guest_default for lic in disk['licenses']: if guest_name in lic: guest_default = guests[arch][1] diff --git a/gce_rescue/messages.py b/gce_rescue/messages.py index 26d070a..dd60ac1 100644 --- a/gce_rescue/messages.py +++ b/gce_rescue/messages.py @@ -17,7 +17,7 @@ from gce_rescue.gce import Instance def tip_connect_ssh(vm: Instance) -> str: - return (f'└── Your instance is READY! You can now connect your instance ' + return (f'└── Your instance is READY! You can now connect to your instance ' f' {vm.name} via:\n 1. CLI. (add --tunnel-through-iap if necessary)\n' f' $ gcloud compute ssh {vm.name} --zone={vm.zone} ' f'--project={vm.project} --ssh-flag="-o StrictHostKeyChecking=no"\n OR\n' diff --git a/gce_rescue/tasks/actions.py b/gce_rescue/tasks/actions.py index 7d3c453..c6fa815 100644 --- a/gce_rescue/tasks/actions.py +++ b/gce_rescue/tasks/actions.py @@ -133,11 +133,11 @@ def call_tasks(vm: Instance, action: str) -> None: tracker = Tracker(total_tasks) tracker.start() - for task in tasks: + for task_index, task in enumerate(tasks, 1): execute = task['name'] args = task['args'][0] - execute(**args) + execute(**args, task_index=task_index, total_tasks=total_tasks) tracker.advance(step = 1) if async_backup_thread: diff --git a/gce_rescue/tasks/disks.py b/gce_rescue/tasks/disks.py index 657fc9c..85c141b 100644 --- a/gce_rescue/tasks/disks.py +++ b/gce_rescue/tasks/disks.py @@ -22,7 +22,7 @@ from gce_rescue.tasks.keeper import wait_for_operation from gce_rescue.tasks.backup import create_snapshot -from gce_rescue.utils import ThreadHandler as Handler +from gce_rescue.utils import ThreadHandler as Handler, tasks_wrapper from googleapiclient.errors import HttpError _logger = logging.getLogger(__name__) @@ -130,6 +130,7 @@ def list_disk(vm, project_data: Dict, label_filter: str) -> Dict: return result['items'] +@tasks_wrapper def attach_disk( vm, disk_name: str, @@ -189,6 +190,7 @@ def take_snapshot(vm, join_snapshot=None) -> None: snapshot_thread.join() +@tasks_wrapper def create_rescue_disk(vm) -> None: device_name = vm.disks['device_name'] # task1 = multitasks.Handler( @@ -222,6 +224,8 @@ def list_snapshot(vm) -> str: return '' return snapshot_name + +@tasks_wrapper def restore_original_disk(vm) -> None: """ Restore tasks to the original disk """ device_name = vm.disks['device_name'] diff --git a/gce_rescue/tasks/metadata.py b/gce_rescue/tasks/metadata.py index 8ac1821..c83d864 100644 --- a/gce_rescue/tasks/metadata.py +++ b/gce_rescue/tasks/metadata.py @@ -16,11 +16,13 @@ from gce_rescue.config import get_config from gce_rescue.tasks.keeper import wait_for_operation, wait_for_os_boot +from gce_rescue.utils import tasks_wrapper from typing import Dict import logging _logger = logging.getLogger(__name__) +@tasks_wrapper def set_metadata(vm) -> Dict: """Configure Instance custom metadata. https://cloud.google.com/compute/docs/reference/rest/v1/instances/setMetadata @@ -53,6 +55,7 @@ def set_metadata(vm) -> Dict: return result +@tasks_wrapper def restore_metadata_items(vm, remove_rescue_mode: bool = False) -> Dict: """Restore original metadata.items after the instance is running again.""" diff --git a/gce_rescue/tasks/operations.py b/gce_rescue/tasks/operations.py index 8ad0ffe..a725618 100644 --- a/gce_rescue/tasks/operations.py +++ b/gce_rescue/tasks/operations.py @@ -16,10 +16,13 @@ from gce_rescue.gce import Instance from gce_rescue.tasks.keeper import wait_for_operation +from gce_rescue.utils import tasks_wrapper import logging _logger = logging.getLogger(__name__) + +@tasks_wrapper def start_instance(vm: Instance) -> str: """Start instance.""" @@ -38,6 +41,7 @@ def start_instance(vm: Instance) -> str: return vm.status +@tasks_wrapper def stop_instance(vm: Instance) -> str: """Stop instance.""" _logger.info(f'Stopping {vm.name}...') diff --git a/gce_rescue/utils.py b/gce_rescue/utils.py index 156124f..4e41640 100644 --- a/gce_rescue/utils.py +++ b/gce_rescue/utils.py @@ -17,6 +17,10 @@ """ List of classes and functions to be used across the code. """ from time import sleep +from datetime import datetime + +from googleapiclient.errors import HttpError + import logging import multiprocessing from threading import Thread @@ -140,4 +144,20 @@ def read_input(msg: str) -> None: if input_answer.upper() != 'Y': print(f'got input: "{input_answer}". Aborting') sys.exit(1) - \ No newline at end of file + + +def tasks_wrapper(task_func): + def inner(*args, **kwargs): + _logger.info(f"task {task_func.__name__} started at: {datetime.now().strftime('%H:%M:%S')}") + task_index = kwargs.pop('task_index') if kwargs.get('task_index') else None + total_tasks = kwargs.pop('total_tasks') if kwargs.get('total_tasks') else None + try: + task_func(*args, **kwargs) + except HttpError as e: + _logger.error(f'HttpError caught on task {task_func.__name__} with error: {e}') + _logger.info(f"task {task_func.__name__} ended at: {datetime.now().strftime('%H:%M:%S')}") + if task_index and total_tasks: + _logger.info(f'Progress: {task_index}/{total_tasks} tasks completed') + print(f'finished {task_func.__name__} {task_index}/{total_tasks} tasks completed') + + return inner \ No newline at end of file