ROB-1033 custom owners (#1798)

RoiGlinik · web-flow · commit f9f7b067f257 · 2025-05-07T09:20:17.000+03:00
* support custom crds in a more generic way. add ExecutionContext, cnpgcluter and StrimziPodSet
diff --git a/helm/robusta/templates/runner-service-account.yaml b/helm/robusta/templates/runner-service-account.yaml
@@ -286,6 +286,39 @@ rules:
     - patch
     - update
 {{- end }}
+{{- if has "StrimziPodSet" .Values.runner.customCRD }}
+  - apiGroups:
+    - core.strimzi.io
+    resources:
+    - strimzipodsets
+    verbs:
+    - get
+    - list
+    - patch
+    - update
+{{- end }}
+{{- if has "CNPGCluster" .Values.runner.customCRD }}
+  - apiGroups:
+    - postgresql.cnpg.io
+    resources:
+    - clusters
+    verbs:
+    - get
+    - list
+    - patch
+    - update
+{{- end }}
+{{- if has "ExecutionContext" .Values.runner.customCRD }}
+  - apiGroups:
+    - hub.knime.com
+    resources:
+    - executioncontexts
+    verbs:
+    - get
+    - list
+    - patch
+    - update
+{{- end }}
 
 ---
 apiVersion: v1
diff --git a/helm/robusta/templates/runner.yaml b/helm/robusta/templates/runner.yaml
@@ -106,6 +106,10 @@ spec:
           - name: ARGO_ROLLOUTS
             value: "True"
           {{- end }}
+          {{- if .Values.runner.customCRD }}
+          - name: CUSTOM_CRD
+            value: {{ .Values.runner.customCRD | toJson }}
+          {{- end -}}
           {{- if kindIs "string" .Values.runner.additional_env_vars }}
           {{- fail "The `additional_env_vars` string value is deprecated. Change the `additional_env_vars` value to an array" -}}
           {{- end }}
diff --git a/helm/robusta/values.yaml b/helm/robusta/values.yaml
@@ -672,6 +672,7 @@ runner:
       cpu: ~
   additional_env_vars: []
   additional_env_froms: []
+  customCRD: []
   priorityClassName: ""
   tolerations: []
   annotations: {}
diff --git a/playbooks/robusta_playbooks/event_enrichments.py b/playbooks/robusta_playbooks/event_enrichments.py
@@ -45,7 +45,21 @@
 from robusta.core.reporting.base import EnrichmentType
 from robusta.core.reporting.custom_rendering import render_value
 from robusta.utils.parsing import format_event_templated_string
+from robusta.core.model.env_vars import CUSTOM_CRD
 
+events_supported_types = [
+        "Pod",
+        "Deployment",
+        "DaemonSet",
+        "ReplicaSet",
+        "StatefulSet",
+        "Job",
+        "Node",
+        "DeploymentConfig",
+        "Rollout",
+    ]
+
+events_supported_types.extend(CUSTOM_CRD)
 
 class ExtendedEventEnricherParams(EventEnricherParams):
     """
@@ -155,17 +169,7 @@ def resource_events_enricher(event: KubernetesResourceEvent, params: ExtendedEve
     """
 
     resource = event.get_resource()
-    if resource.kind not in [
-        "Pod",
-        "Deployment",
-        "DaemonSet",
-        "ReplicaSet",
-        "StatefulSet",
-        "Job",
-        "Node",
-        "DeploymentConfig",
-        "Rollout",
-    ]:
+    if resource.kind not in events_supported_types:
         raise ActionException(
             ErrorCodes.RESOURCE_NOT_SUPPORTED, f"Resource events enricher is not supported for resource {resource.kind}"
         )
@@ -180,20 +184,13 @@ def resource_events_enricher(event: KubernetesResourceEvent, params: ExtendedEve
     )
 
     # append related pod data as well
-    if params.dependent_pod_mode and kind in [
-        "Deployment",
-        "DaemonSet",
-        "ReplicaSet",
-        "StatefulSet",
-        "Job",
-        "DeploymentConfig",
-        "Rollout",
-    ]:
+    if params.dependent_pod_mode and kind not in ["Pod", "Node"]:
         pods = []
         if kind == "Job":
             pods = get_job_all_pods(resource) or []
         else:
-            pods = list_pods_using_selector(resource.metadata.namespace, resource.spec.selector, "")
+            selector = resource.spec.get("selector", {}) if isinstance(resource.spec, dict) else resource.spec.selector
+            pods = list_pods_using_selector(resource.metadata.namespace, selector, "")
 
         selected_pods = pods[: min(len(pods), params.max_pods)]
         for pod in selected_pods:
diff --git a/playbooks/robusta_playbooks/k8s_resource_enrichments.py b/playbooks/robusta_playbooks/k8s_resource_enrichments.py
@@ -44,7 +44,7 @@
     pod_restarts,
 )
 from robusta.core.discovery import utils
-from robusta.core.model.env_vars import RESOURCE_YAML_BLOCK_LIST
+from robusta.core.model.env_vars import RESOURCE_YAML_BLOCK_LIST, CUSTOM_CRD
 from robusta.core.model.pods import ResourceAttributes
 
 
@@ -96,7 +96,7 @@ class RelatedPod(BaseModel):
 
 
 supported_resources = ["Deployment", "DaemonSet", "ReplicaSet", "Pod", "StatefulSet", "Job", "Node", "DeploymentConfig", "Rollout"]
-
+supported_resources.extend(CUSTOM_CRD)
 
 def to_pod_row(pod: V1Pod, cluster_name: str) -> List:
     resource_requests = k8s_pod_requests(pod)
@@ -153,7 +153,8 @@ def get_related_pods_with_extra_info(resource, limit: Optional[int]=None, _conti
             _continue=_continue,
         )
     else:
-        selector = build_selector_query(resource.spec.selector)
+        selector = resource.spec.get("selector", {}) if isinstance(resource.spec, dict) else resource.spec.selector
+        selector = build_selector_query(selector)
         result = client.CoreV1Api().list_namespaced_pod(
             namespace=resource.metadata.namespace,
             label_selector=selector,
diff --git a/playbooks/robusta_playbooks/workload_actions.py b/playbooks/robusta_playbooks/workload_actions.py
@@ -16,7 +16,9 @@
     DaemonSet,
 )
 from robusta.integrations.kubernetes.custom_models import DeploymentConfig, Rollout
-from kubernetes.client import AppsV1Api
+from robusta.integrations.kubernetes.custom_crds import CRDBase
+
+from kubernetes.client import AppsV1Api, CustomObjectsApi
 
 
 @action
@@ -44,12 +46,7 @@ def rollout_restart(event: KubernetesResourceEvent):
     """
     resource = event.get_resource()
     if not resource:
-        raise ActionException(ErrorCodes.RESOURCE_NOT_FOUND, f"Couldn't found resource")
-
-    if resource.kind not in ["Deployment", "DaemonSet", "StatefulSet", "DeploymentConfig", "Rollout"]:
-        raise ActionException(
-            ErrorCodes.RESOURCE_NOT_SUPPORTED, f"Rollout restart is not supported for resource {resource.kind}"
-        )
+        raise ActionException(ErrorCodes.RESOURCE_NOT_FOUND, "Couldn't found resource")
 
     name = resource.metadata.name
     namespace = resource.metadata.namespace
@@ -72,6 +69,17 @@ def rollout_restart(event: KubernetesResourceEvent):
             resource.update()
             return
 
+        if isinstance(resource, (CRDBase)):
+            CustomObjectsApi().patch_namespaced_custom_object(
+                group=resource.group,
+                version=resource.version,
+                namespace=namespace,
+                plural=resource.plural,
+                name=name,
+                body=resource.rollout_restart_patch()
+            )
+            return
+
         if isinstance(resource, Deployment):
             func_name = "patch_namespaced_deployment"
         elif isinstance(resource, StatefulSet):
diff --git a/scripts/generate_kubernetes_code.py b/scripts/generate_kubernetes_code.py
@@ -81,6 +81,7 @@ def autogenerate_events(f: TextIO):
         from ....core.reporting.consts import FindingSubjectType, FindingSource
         from ....core.reporting.finding_subjects import KubeObjFindingSubject
         from  robusta.integrations.kubernetes.custom_models import {CUSTOM_MODELS_IMPORTS}
+        from  robusta.integrations.kubernetes.custom_crds import CRDS_map
         """
         )
     )
@@ -137,6 +138,16 @@ def autogenerate_events(f: TextIO):
 
     f.write(f"{'}'}\n\n\n")
 
+    f.write(
+        textwrap.dedent(
+            """\
+        for cls in CRDS_map.values():
+            LOADERS_MAPPINGS[cls.name.lower()] = (True, cls.readNamespaced)
+
+        """
+        )
+    )
+
     f.write(
         textwrap.dedent(
             """\
diff --git a/src/robusta/core/discovery/discovery.py b/src/robusta/core/discovery/discovery.py
@@ -5,7 +5,7 @@
 from collections import defaultdict
 from concurrent.futures.process import BrokenProcessPool, ProcessPoolExecutor
 from typing import Dict, List, Optional, Union
-
+import dpath.util
 import prometheus_client
 from hikaru.model.rel_1_26 import (
     Container,
@@ -50,6 +50,7 @@
     DISCOVERY_PROCESS_TIMEOUT_SEC,
     IS_OPENSHIFT,
     OPENSHIFT_GROUPS,
+    CUSTOM_CRD
 )
 from robusta.core.model.helm_release import HelmRelease
 from robusta.core.model.jobs import JobInfo
@@ -58,6 +59,7 @@
 from robusta.core.model.openshift_group import OpenshiftGroup
 from robusta.core.model.services import ContainerInfo, ServiceConfig, ServiceInfo, VolumeInfo
 from robusta.integrations.kubernetes.custom_models import DeploymentConfig, DictToK8sObj, Rollout
+from robusta.integrations.kubernetes.custom_crds import CRDS_map
 from robusta.patch.patch import create_monkey_patches
 from robusta.utils.cluster_provider_discovery import cluster_provider
 from robusta.utils.stack_tracer import StackTracer
@@ -191,11 +193,53 @@ def discovery_process() -> DiscoveryResults:
         node_requests = defaultdict(list)  # map between node name, to request of pods running on it
         active_services: List[ServiceInfo] = []
         openshift_groups: List[OpenshiftGroup] = []
-
+        continue_ref: Optional[str] = None
         # discover micro services
+
         try:
+            for cls_name in CUSTOM_CRD:
+                if (cls := CRDS_map.get(cls_name)) is None:
+                    continue
+
+                for _ in range(DISCOVERY_MAX_BATCHES):
+                    try:
+                        crd_res = client.CustomObjectsApi().list_cluster_custom_object(
+                            group=cls.group,
+                            version=cls.version,
+                            plural=cls.plural,
+                            limit=DISCOVERY_BATCH_SIZE,
+                            _continue=continue_ref,
+                        )
+                    except Exception:
+                        logging.exception(msg=f"Failed to list {cls.name} from api.")
+                        break
+
+                    for crd in crd_res.get("items", []):
+                        try:
+                            meta = DictToK8sObj(crd.get("metadata"), V1ObjectMeta)
+                            active_services.extend(
+                                [
+                                    Discovery.__create_service_info(
+                                        meta=meta,
+                                        kind=cls.name,
+                                        containers=[],
+                                        volumes=[],
+                                        total_pods=dpath.util.get(crd, cls.total_pods_path, default=0),
+                                        ready_pods=dpath.util.get(crd, cls.ready_pods_path, default=0),
+                                        is_helm_release=is_release_managed_by_helm(
+                                            annotations=meta.annotations, labels=meta.labels
+                                        ),
+                                    )
+                                ]
+                            )
+                        except Exception:
+                            logging.exception(msg=f"Failed to parse {cls.name} {crd}")
+                            continue
+
+                    continue_ref = crd_res.get("metadata", {}).get("continue")
+                    if not continue_ref:
+                        break
 
-            continue_ref: Optional[str] = None
             if IS_OPENSHIFT:
                 for _ in range(DISCOVERY_MAX_BATCHES):
                     try:
diff --git a/src/robusta/core/discovery/resource_names.py b/src/robusta/core/discovery/resource_names.py
@@ -5,7 +5,7 @@
 
 from robusta.integrations.kubernetes.custom_models import DeploymentConfig, Rollout
 from robusta.utils.error_codes import ActionException, ErrorCodes
-
+from robusta.integrations.kubernetes.custom_crds import CRDS_map
 
 class ResourceLister(BaseModel):
     list_all: Callable
@@ -66,6 +66,13 @@ class ResourceLister(BaseModel):
 }
 
 
+for cls in CRDS_map.values():
+    LISTERS[cls.name.lower()] = ResourceLister(
+        list_all=cls.list_for_all_namespaces,
+        list_namespaced=cls.list_namespaced,
+    )
+
+
 class ResourceNameLister:
     @staticmethod
     def list_resource_names(kind: str, namespace: str = None) -> List[str]:
diff --git a/src/robusta/core/model/env_vars.py b/src/robusta/core/model/env_vars.py
@@ -135,3 +135,6 @@ def load_bool(env_var, default: bool):
 
 # simple calculated values (not direct environment vars)
 SENTRY_ENABLED = SEND_ADDITIONAL_TELEMETRY and SENTRY_DSN
+
+# enable custom CRDs supported by robusta "["StrimziPodSet", "Cluster"]"
+CUSTOM_CRD = json.loads(os.environ.get("CUSTOM_CRD", "[]"))
diff --git a/src/robusta/integrations/kubernetes/autogenerated/events.py b/src/robusta/integrations/kubernetes/autogenerated/events.py
diff --git a/src/robusta/integrations/kubernetes/autogenerated/triggers.py b/src/robusta/integrations/kubernetes/autogenerated/triggers.py
diff --git a/src/robusta/integrations/kubernetes/autogenerated/v1/models.py b/src/robusta/integrations/kubernetes/autogenerated/v1/models.py
diff --git a/src/robusta/integrations/kubernetes/custom_crds.py b/src/robusta/integrations/kubernetes/custom_crds.py
