NodeTasksBase

Author: lexxai

src/cluster_tasks/scenarios/scenario_base.py

@@ -1,9 +1,8 @@
-import asyncio
 import logging
-import time
 from abc import ABC, abstractmethod
 
 from ext_api.proxmox_api import ProxmoxAPI
+from cluster_tasks.tasks.node_tasks_sync import NodeTasksSync
 
 logger = logging.getLogger("CT.{__name__}")
 
@@ -12,26 +11,7 @@ class ScenarioBase(ABC):
     def __init__(self, api: ProxmoxAPI):
         self.name = self.__class__.__name__
         self.api: ProxmoxAPI = api
-
-    @staticmethod
-    def get_status(api: ProxmoxAPI, node: str, upid: str):
-        result = api.nodes(node).tasks(upid).status.get(filter_keys="status")
-        return result
-
-    def wait_task_done(self, api: ProxmoxAPI, node: str, upid: str):
-        while self.get_status(api, node, upid) != "stopped":
-            logger.info("Waiting for task to finish...")
-            time.sleep(2)
-
-    @staticmethod
-    async def async_get_status(api: ProxmoxAPI, node: str, upid: str):
-        result = await api.nodes(node).tasks(upid).status.get(filter_keys="status")
-        return result
-
-    async def async_wait_task_done(self, api: ProxmoxAPI, node: str, upid: str):
-        while await self.async_get_status(api, node, upid) != "stopped":
-            logger.info("Waiting for task to finish...")
-            await asyncio.sleep(2)
+        self.node_tasks = NodeTasksSync(api)
 
     @abstractmethod
     def run(self):

src/cluster_tasks/scenarios/scenario_clone_template_vm.py

@@ -15,30 +15,33 @@ def __init__(self, api: ProxmoxAPI):
 
     def configure(self, config):
         self.node = config.get("node")
-        self.vmid = config.get("vmid")
+        self.vmid = int(config.get("vmid", 0))
         self.newid = config.get("newid")
         self.name = config.get("name")
         self.full = int(config.get("full", True))
 
     def run(self):
         print(f"Running Scenario Template: {self.template_name} at {self.source_node}")
         # Perform the specific API logic for this scenario
-
-        # data = {"newid": int(self.newid), "name": self.name, "full": self.full}
-        with self.api as api:
-            present = (
-                api.nodes(self.node)
-                .qemu(self.newid)
-                .status.current.get(filter_keys="vmid")
-            )
-            logger.info(f"result: {present} {self.newid} {self.newid==present}")
-            if present and int(present) == int(self.newid):
-                logger.info(f"VM {self.newid} already exists - Deleteting...")
-                upid = api.nodes(self.node).qemu(self.newid).delete()
-                logger.info(f"result: {upid}")
-                self.wait_task_done(api, self.node, upid)
-            data = {"newid": int(self.newid), "name": self.name, "full": self.full}
-            upid = api.nodes(self.node).qemu(self.vmid).clone.create(data=data)
-            logger.info(f"result: {upid}, data: {data}")
-            self.wait_task_done(api, self.node, upid)
-            logger.info(f"VM {self.newid} created")
+        try:
+            # Open a connection session of the Proxmox API
+            with self.api:
+                # Check is VM already exists
+                present = self.node_tasks.vm_status(self.node, self.newid)
+                if present:
+                    # If VM already exists, delete it
+                    logger.info(f"VM {self.newid} already exists - Deleting...")
+                    is_deleted = self.node_tasks.vm_delete(self.node, self.newid)
+                    if is_deleted:
+                        logger.info(f"VM {self.newid} deleted successfully")
+                    else:
+                        raise Exception(f"Failed to delete VM {self.newid}")
+                # Clone the VM from the template
+                data = {"newid": int(self.newid), "name": self.name, "full": self.full}
+                is_created = self.node_tasks.vm_clone(self.node, self.vmid, data)
+                if is_created:
+                    logger.info(f"VM {self.newid} cloned successfully")
+                else:
+                    raise Exception(f"Failed to clone VM {self.newid}")
+        except Exception as e:
+            logger.error(f"Failed to run scenario: {e}")

src/cluster_tasks/scenarios_configs.yaml

@@ -1,3 +1,6 @@
+API:
+  backend: "https"
+  type: "sync"
 Scenarios:
   CloneTemplateVM:
     file: "clone_template_vm"

src/cluster_tasks/tasks/__init__.py

@@ -0,0 +1,24 @@
+from datetime import timedelta
+
+from ext_api.proxmox_api import ProxmoxAPI
+
+
+class BaseTasks:
+    TIMEOUT = 60
+    LOOP_SLEEP = 2
+
+    def __init__(self, api: ProxmoxAPI):
+        self._api: ProxmoxAPI = api
+
+    @property
+    def api(self):
+        return self._api
+
+    @api.setter
+    def api(self, value):
+        self._api = value
+
+    @staticmethod
+    def format_duration(seconds: float | int) -> str:
+        """Format a duration (in seconds) as HH:MM:SS."""
+        return str(timedelta(seconds=seconds)).split(".")[0]

src/cluster_tasks/tasks/base_tasks.py

@@ -0,0 +1,88 @@
+import asyncio
+import logging
+
+from cluster_tasks.tasks.base_tasks import BaseTasks
+from cluster_tasks.tasks.node_tasks_base import NodeTasksBase
+
+logger = logging.getLogger("CT.{__name__}")
+
+
+class NodeTasksAsync(NodeTasksBase):
+    """
+    NodeTasks is a class for managing tasks related to virtual machines
+    on a Proxmox node. This includes operations like checking task status,
+    waiting for tasks to complete, and performing VM operations (e.g., clone, delete).
+
+    Inherits from `NodeTasksBase`, which provides common functionality for API interaction.
+
+    Methods:
+        vm_status(node: str, vm_id: int) -> int:
+            Retrieves the current status of a virtual machine by its ID.
+
+        vm_delete(node: str, vm_id: int, wait: bool = True) -> str | bool | None:
+            Deletes a virtual machine and optionally waits for the task to finish.
+
+        vm_clone(node: str, vm_id: int, data: dict, wait: bool = True) -> str | bool | None:
+            Clones a virtual machine and optionally waits for the task to finish.
+    """
+
+    async def vm_status(self, node: str, vm_id: int) -> int:
+        """
+        Retrieves the current status of a virtual machine.
+
+        Args:
+            node (str): The name of the Proxmox node.
+            vm_id (int): The ID of the virtual machine.
+
+        Returns:
+            int: The status of the virtual machine (e.g., running, stopped).
+        """
+        status_vm = (
+            await self.api.nodes(node)
+            .qemu(vm_id)
+            .status.current.get(filter_keys="vmid")
+        )
+        return int(status_vm) if status_vm else 0
+
+    async def vm_delete(
+        self, node: str, vm_id: int, wait: bool = True
+    ) -> str | bool | None:
+        """
+        Deletes a virtual machine and optionally waits for the deletion task to complete.
+
+        Args:
+            node (str): The name of the Proxmox node.
+            vm_id (int): The ID of the virtual machine.
+            wait (bool): Whether to wait for the task to complete (default is True).
+
+        Returns:
+            str | bool | None: The task UPID if `wait` is False;
+                               `True` if task is completed successfully,
+                               `False` if task timed out.
+        """
+        upid = await self.api.nodes(node).qemu(vm_id).delete()
+        if wait:
+            return await self.wait_task_done_async(node, upid)
+        return upid
+
+    async def vm_clone(
+        self, node: str, vm_id: int, data: dict, wait: bool = True
+    ) -> str | bool | None:
+        """
+        Clones a virtual machine and optionally waits for the cloning task to complete.
+
+        Args:
+            node (str): The name of the Proxmox node.
+            vm_id (int): The ID of the virtual machine to clone.
+            data (dict): The parameters for cloning the virtual machine.
+            wait (bool): Whether to wait for the task to complete (default is True).
+
+        Returns:
+            str | bool | None: The task UPID if `wait` is False;
+                               `True` if the task finished successfully,
+                               `False` if task timed out.
+        """
+        upid = await self.api.nodes(node).qemu(vm_id).clone.create(data=data)
+        if wait:
+            return await self.wait_task_done_async(node, upid)
+        return upid

src/cluster_tasks/tasks/node_tasks_async.py

@@ -0,0 +1,88 @@
+import asyncio
+import logging
+import time
+
+from cluster_tasks.tasks.base_tasks import BaseTasks
+
+logger = logging.getLogger("CT.{__name__}")
+
+
+class NodeTasksBase(BaseTasks):
+    """
+    Base class with shared logic for both sync and async task handling.
+    Contains methods for both sync and async API calls.
+    Methods:
+        get_status_sync(node: str, upid: str) -> str | None:
+            Retrieves the status of a task by its UPID.
+
+        wait_task_done_sync(node: str, upid: str) -> bool:
+            Waits for a task to complete, checking the status periodically.
+
+        get_status_async(node: str, upid: str) -> str | None:
+            Retrieves the status of a task by its UPID.
+
+        wait_task_done_async(node: str, upid: str) -> bool:
+            Waits for a task to complete, checking the status periodically.
+    """
+
+    def get_status_sync(self, node: str, upid: str) -> str | None:
+        """
+        Synchronously retrieves the status of a task identified by its UPID.
+        """
+        if upid:
+            result = self.api.nodes(node).tasks(upid).status.get(filter_keys="status")
+            return result
+
+    async def get_status_async(self, node: str, upid: str) -> str | None:
+        """
+        Asynchronously retrieves the status of a task identified by its UPID.
+        """
+        if upid:
+            result = (
+                await self.api.nodes(node).tasks(upid).status.get(filter_keys="status")
+            )
+            return result
+
+    def wait_task_done_sync(self, node: str, upid: str) -> bool:
+        """
+        Synchronously waits for a task to complete.
+        """
+        start_time = time.time()
+        while (result := self.get_status_sync(node, upid)) is not None:
+            if result == "stopped":
+                return True
+            duration = time.time() - start_time
+            formatted_duration = self.format_duration(duration)
+            formatted_timeout = self.format_duration(self.TIMEOUT)
+            logger.info(
+                f"Waiting for task to finish... [ {formatted_duration} / {formatted_timeout} ]"
+            )
+            time.sleep(self.LOOP_SLEEP)
+            if time.time() - start_time > self.TIMEOUT:
+                logger.warning(
+                    f"Timeout reached while waiting for task to finish. {upid=}"
+                )
+                break
+        return False
+
+    async def wait_task_done_async(self, node: str, upid: str) -> bool:
+        """
+        Asynchronously waits for a task to complete.
+        """
+        start_time = time.time()
+        while (result := await self.get_status_async(node, upid)) is not None:
+            if result == "stopped":
+                return True
+            duration = time.time() - start_time
+            formatted_duration = self.format_duration(duration)
+            formatted_timeout = self.format_duration(self.TIMEOUT)
+            logger.info(
+                f"Waiting for task to finish... [ {formatted_duration} / {formatted_timeout} ]"
+            )
+            await asyncio.sleep(self.LOOP_SLEEP)
+            if time.time() - start_time > self.TIMEOUT:
+                logger.warning(
+                    f"Timeout reached while waiting for task to finish. {upid=}"
+                )
+                break
+        return False

src/cluster_tasks/tasks/node_tasks_base.py

@@ -0,0 +1,80 @@
+import logging
+import time
+
+from cluster_tasks.tasks.base_tasks import BaseTasks
+from cluster_tasks.tasks.node_tasks_base import NodeTasksBase
+
+# Creating a logger instance specific to the current module
+logger = logging.getLogger("CT.{__name__}")
+"""
+The logger is used for logging messages related to task status monitoring and
+task operations within the NodeTasks class. It is named with the format "CT.{module_name}"
+to reflect the source module where the log entries are generated.
+"""
+
+
+class NodeTasksSync(NodeTasksBase):
+    """
+    NodeTasks class provides functionality for interacting with tasks on a Proxmox node.
+    It includes methods to check task status, wait for tasks to finish, and perform
+    operations on virtual machines (VMs), such as deleting and cloning them.
+
+    Inherits from `NodeTasksBase` for common API functionality.
+    """
+
+    def vm_status(self, node: str, vm_id: int) -> int:
+        """
+        Retrieves the current status of a virtual machine.
+
+        Args:
+            node (str): The name of the Proxmox node.
+            vm_id (int): The ID of the virtual machine.
+
+        Returns:
+            int: The status of the virtual machine (e.g., running, stopped).
+        """
+        status_vm = (
+            self.api.nodes(node).qemu(vm_id).status.current.get(filter_keys="vmid")
+        )
+        return int(status_vm) if status_vm else 0
+
+    def vm_delete(self, node: str, vm_id: int, wait: bool = True) -> str | bool | None:
+        """
+        Deletes a virtual machine and optionally waits for the deletion task to complete.
+
+        Args:
+            node (str): The name of the Proxmox node.
+            vm_id (int): The ID of the virtual machine.
+            wait (bool): Whether to wait for the task to complete (default is True).
+
+        Returns:
+            str | bool | None: The task UPID if `wait` is False;
+                               `True` if task is completed successfully,
+                               `False` if task timed out.
+        """
+        upid = self.api.nodes(node).qemu(vm_id).delete()
+        if wait:
+            return self.wait_task_done_sync(node, upid)
+        return upid
+
+    def vm_clone(
+        self, node: str, vm_id: int, data: dict, wait: bool = True
+    ) -> str | bool | None:
+        """
+        Clones a virtual machine and optionally waits for the cloning task to complete.
+
+        Args:
+            node (str): The name of the Proxmox node.
+            vm_id (int): The ID of the virtual machine to clone.
+            data (dict): The parameters for cloning the virtual machine.
+            wait (bool): Whether to wait for the task to complete (default is True).
+
+        Returns:
+            str | bool | None: The task UPID if `wait` is False;
+                               `True` if the task finished successfully,
+                               `False` if task timed out.
+        """
+        upid = self.api.nodes(node).qemu(vm_id).clone.create(data=data)
+        if wait:
+            return self.wait_task_done_sync(node, upid)
+        return upid