krkn-lib 5.1.11__py3-none-any.whl → 5.1.13__py3-none-any.whl

krkn_lib/k8s/krkn_kubernetes.py

@@ -167,9 +167,9 @@ class KrknKubernetes:
             client.Configuration.set_default(self.client_config)
             self.watch_resource = watch.Watch()
             # Get the logger for the kubernetes client
-            kubernetes_logger = logging.getLogger('kubernetes')
+            kubernetes_logger = logging.getLogger("kubernetes")
 
-            # Set the logging level to a higher level than DEBUG, 
+            # Set the logging level to a higher level than DEBUG,
             # such as INFO, WARNING, or ERROR
             # This will effectively disable DEBUG level messages.
             kubernetes_logger.setLevel(logging.INFO)
@@ -1388,6 +1388,304 @@ class KrknKubernetes:
                 str(e),
             )
 
+    def get_vm(self, name: str, namespace: str) -> Optional[Dict]:
+        """
+        Get a Virtual Machine by name and namespace.
+
+        :param name: Name of the VM to retrieve
+        :param namespace: Namespace of the VM
+        :return: The VM object if found, None otherwise
+        """
+        try:
+            vm = self.custom_object_client.get_namespaced_custom_object(
+                group="kubevirt.io",
+                version="v1",
+                namespace=namespace,
+                plural="virtualmachines",
+                name=name,
+            )
+            return vm
+        except ApiException as e:
+            if e.status == 404:
+                logging.warning(
+                    f"VM {name} not found in namespace {namespace}"
+                )
+                return None
+            else:
+                logging.error(f"Error getting VM {name}: {e}")
+                raise
+        except Exception as e:
+            logging.error(f"Unexpected error getting VM {name}: {e}")
+            raise
+
+    def get_vmi(self, name: str, namespace: str) -> Optional[Dict]:
+        """
+        Get a Virtual Machine Instance by name and namespace.
+
+        :param name: Name of the VMI to retrieve
+        :param namespace: Namespace of the VMI
+        :return: The VMI object if found, None otherwise
+        """
+        try:
+            vmi = self.custom_object_client.get_namespaced_custom_object(
+                group="kubevirt.io",
+                version="v1",
+                namespace=namespace,
+                plural="virtualmachineinstances",
+                name=name,
+            )
+            return vmi
+        except ApiException as e:
+            if e.status == 404:
+                logging.warning(
+                    f"VMI {name} not found in namespace {namespace}"
+                )
+                return None
+            else:
+                logging.error(f"Error getting VMI {name}: {e}")
+                raise
+        except Exception as e:
+            logging.error(f"Unexpected error getting VMI {name}: {e}")
+            raise
+
+    def get_vmis(self, regex_name: str, namespace: str) -> Optional[Dict]:
+        """
+        Get a Virtual Machine Instance by name and namespace.
+
+        :param name: Name of the VMI to retrieve
+        :param namespace: Namespace of the VMI
+        :return: The VMI object if found, None otherwise
+        """
+        try:
+            vmis_list = []
+            namespaces = self.list_namespaces_by_regex(namespace)
+            for namespace in namespaces:
+                vmis = self.custom_object_client.list_namespaced_custom_object(
+                    group="kubevirt.io",
+                    version="v1",
+                    namespace=namespace,
+                    plural="virtualmachineinstances",
+                )
+
+                for vmi in vmis.get("items"):
+                    vmi_name = vmi.get("metadata", {}).get("name")
+                    match = re.match(regex_name, vmi_name)
+                    if match:
+                        vmis_list.append(vmi)
+        except ApiException as e:
+            if e.status == 404:
+                logging.warning(
+                    f"VMI {regex_name} not found in namespace {namespace}"
+                )
+                return []
+            else:
+                logging.error(f"Error getting VMI {regex_name}: {e}")
+                raise
+        except Exception as e:
+            logging.error(f"Unexpected error getting VMI {regex_name}: {e}")
+            raise
+        return vmis_list
+
+    def create_vmi(self, name: str, namespace: str, vm_name: str, vmi_body: dict) -> Optional[Dict]:
+        """
+        Create a Virtual Machine Instance by name and namespace.
+
+        :param name: Name of the VMI to create
+        :param namespace: Namespace of the VMI
+        :param vm_name: Name of the Virtual Machine to create the VMI from
+        :return: The VMI object if created, None otherwise
+        """
+        try:
+            vmi = self.custom_object_client.create_namespaced_custom_object(
+                group="kubevirt.io",
+                version="v1",
+                namespace=namespace,
+                plural="virtualmachineinstances",
+                name=name,
+                body=vmi_body,
+            )
+            return vmi
+        except ApiException as e:
+            if e.status == 404:
+                logging.warning(f"VMI {name} not found in namespace {namespace}")
+                return None
+            else:
+                logging.error(f"Error creating VMI {name}: {e}")
+                raise
+        except Exception as e:
+            logging.error(f"Unexpected error creating VMI {name}: {e}")
+            raise
+
+    def patch_vm(self, name: str, namespace: str, vm_body: dict) -> Optional[Dict]:
+        """
+        Patch a Virtual Machine by name and namespace.
+
+        :param name: Name of the VM to patch
+        :param namespace: Namespace of the VM
+        :param vm_body: Body of the VM to patch
+        :return: The VM object if patched, None otherwise
+        """
+        try:
+            vmi = self.custom_object_client.patch_namespaced_custom_object(
+                group="kubevirt.io",
+                version="v1",
+                namespace=namespace,
+                plural="virtualmachines",
+                name=name,
+                body=vm_body,
+            )
+            return vmi
+        except ApiException as e:
+            if e.status == 404:
+                logging.warning(f"VM {name} not found in namespace {namespace}")
+                return None
+            else:
+                logging.error(f"Error patching VM {name}: {e}")
+                raise
+        except Exception as e:
+            logging.error(f"Unexpected error patching VM {name}: {e}")
+            raise
+
+    def patch_vmi(self, name: str, namespace: str, vmi_body: dict) -> Optional[Dict]:
+        """
+        Patch a Virtual Machine Instance by name and namespace.
+
+        :param name: Name of the VMI to patch
+        :param namespace: Namespace of the VMI
+        :param vmi_body: Body of the VMI to patch
+        :return: The VMI object if patched, None otherwise
+        """
+        try:
+            vmi = self.custom_object_client.patch_namespaced_custom_object(
+                group="kubevirt.io",
+                version="v1",
+                namespace=namespace,
+                plural="virtualmachineinstances",
+                name=name,
+                body=vmi_body,
+            )
+            return vmi
+        except ApiException as e:
+            if e.status == 404:
+                logging.warning(f"VMI {name} not found in namespace {namespace}")
+                return None
+            else:
+                logging.error(f"Error patching VMI {name}: {e}")
+                raise
+        except Exception as e:
+            logging.error(f"Unexpected error patching VMI {name}: {e}")
+            raise
+
+    def get_vms(self, regex_name: str, namespace: str) -> Optional[Dict]:
+        """
+        Get a Virtual Machine by name and namespace.
+
+        :param name: Name of the VM to retrieve
+        :param namespace: Namespace of the VM
+        :return: The VM object if found, None otherwise
+        """
+        try:
+            vms_list = []
+            namespaces = self.list_namespaces_by_regex(namespace)
+            for namespace in namespaces:
+                vms = self.custom_object_client.list_namespaced_custom_object(
+                    group="kubevirt.io",
+                    version="v1",
+                    namespace=namespace,
+                    plural="virtualmachines",
+                )
+
+                for vm in vms.get("items"):
+                    vm_name = vm.get("metadata", {}).get("name")
+                    match = re.match(regex_name, vm_name)
+                    if match:
+                        vms_list.append(vm)
+            return vms_list
+        except ApiException as e:
+            if e.status == 404:
+                logging.warning(
+                    f"VM {regex_name} not found in namespace {namespace}"
+                )
+                return []
+            else:
+                logging.error(f"Error getting VM {regex_name}: {e}")
+                raise
+        except Exception as e:
+            logging.error(f"Unexpected error getting VM {regex_name}: {e}")
+            raise
+    
+    def get_snapshot(self, name: str, namespace: str) -> Optional[Dict]:
+        """
+        Get a Snapshot by name and namespace.
+
+        :param name: Name of the Snapshot to retrieve
+        :param namespace: Namespace of the Snapshot
+        :return: The Snapshot object if found, None otherwise
+        """
+        try:
+            vmi = self.custom_object_client.get_namespaced_custom_object(
+                group="kubevirt.io",
+                version="v1",
+                namespace=namespace,
+                plural="VirtualMachineSnapshot",
+                name=name,
+            )
+            return vmi
+        except ApiException as e:
+            if e.status == 404:
+                logging.warning(
+                    f"VMI {name} not found in namespace {namespace}"
+                )
+                return None
+            else:
+                logging.error(f"Error getting VMI {name}: {e}")
+                raise
+        except Exception as e:
+            logging.error(f"Unexpected error getting VMI {name}: {e}")
+            raise
+
+    def create_snapshot(
+        self, name: str, namespace: str, vm_name: str
+    ) -> Optional[Dict]:
+        """
+        Create a Snapshot by name and namespace.
+
+        :param name: Name of the Snapshot to create
+        :param namespace: Namespace of the Snapshot
+        :param vm_name: Name of the Virtual Machine to create the Snapshot from
+        :return: The Snapshot object if created, None otherwise
+        """
+        try:
+            file_loader = PackageLoader("krkn_lib.k8s", "templates")
+            env = Environment(loader=file_loader, autoescape=True)
+            snapshot_template = env.get_template("snapshot.j2")
+            ss_body = yaml.safe_load(
+                snapshot_template.render(
+                    name=name, namespace=namespace, vm_name=vm_name
+                )
+            )
+            vmi = self.custom_object_client.create_namespaced_custom_object(
+                group="kubevirt.io",
+                version="v1",
+                namespace=namespace,
+                plural="VirtualMachineSnapshot",
+                name=name,
+                body=ss_body,
+            )
+            return vmi
+        except ApiException as e:
+            if e.status == 404:
+                logging.warning(
+                    f"Snapshot {name} not found in namespace {namespace}"
+                )
+                return None
+            else:
+                logging.error(f"Error creating Snapshot {name}: {e}")
+                raise
+        except Exception as e:
+            logging.error(f"Unexpected error creating Snapshot {name}: {e}")
+            raise
+
     def get_job_status(
         self, name: str, namespace: str = "default"
     ) -> client.V1Job:
@@ -1410,6 +1708,82 @@ class KrknKubernetes:
             )
             raise
 
+    def delete_vm(self, name: str, namespace: str) -> Optional[Dict]:
+        """
+        Delete a Virtual Machine by name and namespace.
+
+        :param name: Name of the VM to delete
+        :param namespace: Namespace of the VM
+        :return: The VM object if found, None otherwise
+        """
+        try:
+            return self.custom_object_client.delete_namespaced_custom_object(
+                group="kubevirt.io",
+                version="v1",
+                namespace=namespace,
+                plural="virtualmachines",
+                name=name,
+            )
+        except ApiException as e:
+            if e.status == 404:
+                logging.warning(
+                    f"VM {name} not found in namespace {namespace}"
+                )
+                return None
+            else:
+                logging.error(f"Error deleting VM {name}: {e}")
+                raise
+        except Exception as e:
+            logging.error(f"Error deleting VM {name}: {e}")
+            raise
+
+    def delete_vmi(self, vm_name: str, namespace: str):
+        """
+        Delete a Virtual Machine Instance to simulate a VM outage.
+
+        :param vm_name: Name of the VMI to delete
+        :param namespace: Namespace of the VMI
+        :return: 0 for success, 1 for failure
+        """
+        logging.info(
+            f"Injecting chaos: Deleting VMI {vm_name} in namespace {namespace}"
+        )
+        try:
+            self.custom_object_client.delete_namespaced_custom_object(
+                group="kubevirt.io",
+                version="v1",
+                namespace=namespace,
+                plural="virtualmachineinstances",
+                name=vm_name,
+            )
+        except ApiException as e:
+            if e.status == 404:
+                logging.warning(f"VMI {vm_name} not found during deletion")
+                return 1
+            else:
+                logging.error(f"API error during VMI deletion: {e}")
+                return 1
+
+    def delete_snapshot(self, snapshot_name: str, namespace: str):
+        """Helper method to delete any snapshot created by the scenario."""
+        self.logger.info(f"Deleting snapshot '{self.snapshot_name}'...")
+        try:
+            self.custom_object_client.delete_namespaced_custom_object(
+                group="kubevirt.io",
+                version="v1",
+                namespace=namespace,
+                plural="VirtualMachineSnapshot",
+                name=snapshot_name,
+            )
+            self.logger.info(
+                f"Snapshot '{self.snapshot_name}' deleted successfully."
+            )
+        except Exception as e:
+            self.logger.warning(
+                "Failed to delete snapshot, "
+                f"might have been already deleted: {e}"
+            )
+
     def monitor_nodes(
         self,
     ) -> (bool, list[str]):

krkn_lib/k8s/pod_monitor/pod_monitor.py

@@ -1,10 +1,14 @@
+import logging
 import re
+import time
+import traceback
 from concurrent.futures import Future
 from concurrent.futures.thread import ThreadPoolExecutor
 from functools import partial
 
 from kubernetes import watch
 from kubernetes.client import V1Pod, CoreV1Api
+from urllib3.exceptions import ProtocolError
 
 from krkn_lib.models.pod_monitor.models import (
     PodsSnapshot,
@@ -47,67 +51,151 @@ def _monitor_pods(
     max_timeout: int,
     name_pattern: str = None,
     namespace_pattern: str = None,
+    max_retries: int = 3,
 ) -> PodsSnapshot:
-    w = watch.Watch(return_type=V1Pod)
+    """
+    Monitor pods with automatic retry on watch stream disconnection.
+
+    :param monitor_partial: Partial function for monitoring pods
+    :param snapshot: Snapshot to populate with pod events
+    :param max_timeout: Maximum time to monitor (seconds)
+    :param name_pattern: Regex pattern for pod names
+    :param namespace_pattern: Regex pattern for namespaces
+    :param max_retries: Maximum number of retries on connection error
+        (default: 3)
+    :return: PodsSnapshot with collected pod events
+    """
+
+    start_time = time.time()
+    retry_count = 0
     deleted_parent_pods = []
     restored_pods = []
     cluster_restored = False
-    for event in w.stream(monitor_partial, timeout_seconds=max_timeout):
-        match_name = True
-        match_namespace = True
-        event_type = event["type"]
-        pod = event["object"]
 
-        if namespace_pattern:
-            match = re.match(namespace_pattern, pod.metadata.namespace)
-            match_namespace = match is not None
-        if name_pattern:
-            match = re.match(name_pattern, pod.metadata.name)
-            match_name = match is not None
+    while retry_count <= max_retries:
+        try:
+            # Calculate remaining timeout if retrying
+            if retry_count > 0:
+                elapsed = time.time() - start_time
+                remain_timeout = max(1, int(max_timeout - elapsed))
+                logging.info("remain timeout " + str(remain_timeout))
+                if remain_timeout <= 0:
+                    logging.info("Maximum timeout reached, stopping monitoring")
+                    break
+                logging.info(
+                    "Reconnecting watch stream"
+                    f"(attempt {retry_count}/{max_retries}),"
+                    f"remaining timeout: {remain_timeout}s"
+                )
+            else:
+                remain_timeout = max_timeout
 
-        if match_name and match_namespace:
-            pod_event = PodEvent()
-            if event_type == "MODIFIED":
-                if pod.metadata.deletion_timestamp is not None:
-                    pod_event.status = PodStatus.DELETION_SCHEDULED
-                    deleted_parent_pods.append(pod.metadata.name)
-                elif _is_pod_ready(pod):
-                    pod_event.status = PodStatus.READY
-                    # if there are at least the same number of ready
-                    # pods as the snapshot.initial_pods set we assume that
-                    # the cluster is restored to the initial condition
-                    restored_pods.append(pod.metadata.name)
-                    if len(restored_pods) >= len(snapshot.initial_pods):
-                        cluster_restored = True
-                else:
-                    pod_event.status = PodStatus.NOT_READY
-
-            elif event_type == "DELETED":
-                pod_event.status = PodStatus.DELETED
-            elif event_type == "ADDED":
-                pod_event.status = PodStatus.ADDED
-
-            if pod_event.status == PodStatus.ADDED:
-                snapshot.added_pods.append(pod.metadata.name)
-                # in case a pod is respawn with the same name
-                # the dictionary must not be reinitialized
-                if pod.metadata.name not in snapshot.pods:
-                    snapshot.pods[pod.metadata.name] = MonitoredPod()
-                    snapshot.pods[pod.metadata.name].name = pod.metadata.name
-                    snapshot.pods[pod.metadata.name].namespace = (
-                        pod.metadata.namespace
-                    )
-            # skips events out of the snapshot
-            if pod.metadata.name in snapshot.pods:
-                snapshot.pods[pod.metadata.name].status_changes.append(
-                    pod_event
+            w = watch.Watch(return_type=V1Pod)
+
+            for e in w.stream(monitor_partial, timeout_seconds=remain_timeout):
+                match_name = True
+                match_namespace = True
+                event_type = e["type"]
+                pod = e["object"]
+
+                if namespace_pattern:
+                    match = re.match(namespace_pattern, pod.metadata.namespace)
+                    match_namespace = match is not None
+                if name_pattern:
+                    match = re.match(name_pattern, pod.metadata.name)
+                    match_name = match is not None
+
+                if match_name and match_namespace:
+                    pod_event = PodEvent()
+                    pod_name = pod.metadata.name
+                    if event_type == "MODIFIED":
+                        if pod.metadata.deletion_timestamp is not None:
+                            pod_event.status = PodStatus.DELETION_SCHEDULED
+                            if pod_name not in deleted_parent_pods:
+                                deleted_parent_pods.append(pod_name)
+                        elif _is_pod_ready(pod):
+                            pod_event.status = PodStatus.READY
+                            # if there are at least the same number of ready
+                            # pods as the snapshot.initial_pods set we assume
+                            # the cluster is restored to the initial condition
+                            if pod_name not in restored_pods:
+                                restored_pods.append(pod_name)
+                            inital_pod_len = len(snapshot.initial_pods)
+                            if len(restored_pods) >= inital_pod_len:
+                                cluster_restored = True
+                        else:
+                            pod_event.status = PodStatus.NOT_READY
+
+                    elif event_type == "DELETED":
+                        pod_event.status = PodStatus.DELETED
+                    elif event_type == "ADDED":
+                        pod_event.status = PodStatus.ADDED
+
+                    if pod_event.status == PodStatus.ADDED:
+                        
+                        if pod_name not in snapshot.added_pods:
+                            snapshot.added_pods.append(pod_name)
+                        # in case a pod is respawn with the same name
+                        # the dictionary must not be reinitialized
+                        if pod_name not in snapshot.pods:
+                            snapshot.pods[pod_name] = MonitoredPod()
+                            snapshot.pods[pod_name].name = pod_name
+                            snapshot.pods[pod_name].namespace = (
+                                pod.metadata.namespace
+                            )
+                    # skips events out of the snapshot
+                    if pod_name in snapshot.pods:
+                        snapshot.pods[pod_name].status_changes.append(
+                            pod_event
+                        )
+                    # this flag is set when all the pods
+                    # that has been deleted or not ready
+                    # have been restored, if True the
+                    # monitoring is stopped earlier
+                    if cluster_restored:
+                        logging.info("Cluster restored, stopping monitoring")
+                        w.stop()
+                        return snapshot
+
+            # If we exit the loop normally (timeout reached), we're done
+            logging.info("Watch stream completed normally")
+            break
+
+        except ProtocolError as e:
+
+            if retry_count > max_retries:
+                logging.warning(
+                    f"Watch stream connection broken after {max_retries}"
+                    f"retries. ProtocolError: {e}. Returning snapshot "
+                    "with data collected so far."
+                )
+                break
+
+            # Log retry attempt
+            logging.info(
+                f"Watch stream connection broken (ProtocolError): {e}. "
+                f"Retry {retry_count}/{max_retries} in progress..."
+            )
+            backoff_time = 1
+
+            # Check if we have time for backoff
+            elapsed = time.time() - start_time
+            if elapsed + backoff_time >= max_timeout:
+                logging.info(
+                    "Not enough time remaining for backoff, "
+                    "returning snapshot with data collected."
                 )
-            # this flag is set when all the pods
-            # that has been deleted or not ready
-            # have been restored, if True the
-            # monitoring is stopeed earlier
-            if cluster_restored:
-                w.stop()
+                break
+
+            logging.debug(f"Waiting {backoff_time}s before retry...")
+            time.sleep(backoff_time)
+
+        except Exception as e:
+            logging.error("Error in monitor pods: " + str(e))
+            logging.error("Stack trace:\n%s", traceback.format_exc())
+            raise Exception(e)
+        
+        retry_count += 1
 
     return snapshot
 

krkn_lib/k8s/templates/snapshot.j2

@@ -0,0 +1,10 @@
+apiVersion: snapshot.kubevirt.io/v1alpha1
+kind: VirtualMachineSnapshot
+metadata:
+  name: {{name}}
+  namespace: {{namespace}}
+spec:
+  source:
+    apiGroup: kubevirt.io
+    kind: VirtualMachine
+    name: {{vm_name}}

krkn_lib/models/elastic/models.py

@@ -144,6 +144,7 @@ class ElasticChaosRunTelemetry(Document):
     run_uuid = Text(fields={"keyword": Keyword()})
     health_checks = Nested(ElasticHealthChecks, multi=True)
     virt_checks = Nested(ElasticVirtChecks, multi=True)
+    post_virt_checks = Nested(ElasticVirtChecks, multi=True)
 
     class Index:
         name = "chaos_run_telemetry"
@@ -261,6 +262,28 @@ class ElasticChaosRunTelemetry(Document):
         else:
             self.virt_checks = None
 
+        if chaos_run_telemetry.post_virt_checks:
+            self.post_virt_checks = [
+                ElasticVirtChecks(
+                    vm_name=post_info.vm_name,
+                    ip_address=post_info.ip_address,
+                    new_ip_address=post_info.new_ip_address,
+                    namespace=post_info.namespace,
+                    node_name=post_info.node_name,
+                    status=post_info.status,
+                    start_timestamp=datetime.datetime.fromisoformat(
+                        str(post_info.start_timestamp)
+                    ),
+                    end_timestamp=datetime.datetime.fromisoformat(
+                        str(post_info.end_timestamp)
+                    ),
+                    duration=post_info.duration,
+                )
+                for post_info in chaos_run_telemetry.post_virt_checks
+            ]
+        else:
+            self.post_virt_checks = None
+
         self.timestamp = chaos_run_telemetry.timestamp
         self.total_node_count = chaos_run_telemetry.total_node_count
         self.cloud_infrastructure = chaos_run_telemetry.cloud_infrastructure