Add resource parameter

robusta-dev · Aug 4, 2023 · d35c4fe · d35c4fe
1 parent 92ad847
commit d35c4fe
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 6 deletions.
diff --git a/robusta_krr/core/integrations/kubernetes.py b/robusta_krr/core/integrations/kubernetes.py
@@ -63,6 +63,7 @@ async def list_scannable_objects(self) -> AsyncGenerator[K8sObjectData, None]:
 
         self.info(f"Listing scannable objects in {self.cluster}")
         self.debug(f"Namespaces: {self.config.namespaces}")
+        self.debug(f"Resources: {self.config.resources}")
 
         self.__hpa_list = await self._try_list_hpa()
 
@@ -121,10 +122,19 @@ def __build_obj(
             hpa=self.__hpa_list.get((namespace, kind, name)),
         )
 
+    def _should_list_resource(self, resource: str):
+        if self.config.resources == "*":
+            return True
+        return resource.lower() in self.config.resources
+
     async def _list_workflows(
         self, kind: str, all_namespaces_request: Callable, namespaced_request: Callable
     ) -> AsyncIterator[K8sObjectData]:
-        self.debug(f"Listing {kind} in {self.cluster}")
+        if not self._should_list_resource(kind):
+            self.debug(f"Skipping {kind}s in {self.cluster}")
+            return
+
+        self.debug(f"Listing {kind}s in {self.cluster}")
         loop = asyncio.get_running_loop()
 
         try:
@@ -169,7 +179,7 @@ async def _list_workflows(
 
     def _list_deployments(self) -> AsyncIterator[K8sObjectData]:
         return self._list_workflows(
-            kind="deployments",
+            kind="deployment",
             all_namespaces_request=self.apps.list_deployment_for_all_namespaces,
             namespaced_request=self.apps.list_namespaced_deployment,
         )
@@ -178,7 +188,7 @@ async def _list_rollouts(self) -> AsyncIterator[K8sObjectData]:
         # TODO: Mutlitple errors will throw here, we should catch them all
         try:
             async for rollout in self._list_workflows(
-                kind="ArgoCD Rollout",
+                kind="rollout",
                 all_namespaces_request=self.rollout.list_rollout_for_all_namespaces,
                 namespaced_request=self.rollout.list_namespaced_rollout,
             ):
@@ -191,21 +201,21 @@ async def _list_rollouts(self) -> AsyncIterator[K8sObjectData]:
 
     def _list_all_statefulsets(self) -> AsyncIterator[K8sObjectData]:
         return self._list_workflows(
-            kind="statefulsets",
+            kind="statefulset",
             all_namespaces_request=self.apps.list_stateful_set_for_all_namespaces,
             namespaced_request=self.apps.list_namespaced_stateful_set,
         )
 
     def _list_all_daemon_set(self) -> AsyncIterator[K8sObjectData]:
         return self._list_workflows(
-            kind="daemonsets",
+            kind="daemonset",
             all_namespaces_request=self.apps.list_daemon_set_for_all_namespaces,
             namespaced_request=self.apps.list_namespaced_daemon_set,
         )
 
     def _list_all_jobs(self) -> AsyncIterator[K8sObjectData]:
         return self._list_workflows(
-            kind="jobs",
+            kind="job",
             all_namespaces_request=self.batch.list_job_for_all_namespaces,
             namespaced_request=self.batch.list_namespaced_job,
         )

diff --git a/robusta_krr/core/models/config.py b/robusta_krr/core/models/config.py
@@ -16,6 +16,7 @@ class Config(pd.BaseSettings):
     clusters: Union[list[str], Literal["*"], None] = None
     kubeconfig: Optional[str] = None
     namespaces: Union[list[str], Literal["*"]] = pd.Field("*")
+    resources: Union[list[str], Literal["*"]] = pd.Field("*")
     selector: Optional[str] = None
 
     # Value settings
@@ -68,6 +69,13 @@ def validate_namespaces(cls, v: Union[list[str], Literal["*"]]) -> Union[list[st
         if v == []:
             return "*"
 
+        return [val.lower() for val in v]
+
+    @pd.validator("resources")
+    def validate_resources(cls, v: Union[list[str], Literal["*"]]) -> Union[list[str], Literal["*"]]:
+        if v == []:
+            return "*"
+
         return v
 
     def create_strategy(self) -> AnyStrategy:

diff --git a/robusta_krr/main.py b/robusta_krr/main.py
@@ -72,6 +72,13 @@ def {func_name}(
                     help="List of namespaces to run on. By default, will run on all namespaces.",
                     rich_help_panel="Kubernetes Settings"
                 ),
+                resources: List[str] = typer.Option(
+                    None,
+                    "--resource",
+                    "-r",
+                    help="List of resources to run on (Deployment, StatefullSet, DaemonSet, Job, Rollout). By default, will run on all resources. Case insensitive.",
+                    rich_help_panel="Kubernetes Settings"
+                ),
                 selector: Optional[str] = typer.Option(
                     None,
                     "--selector",
@@ -139,6 +146,7 @@ def {func_name}(
                     kubeconfig=kubeconfig,
                     clusters="*" if all_clusters else clusters,
                     namespaces="*" if "*" in namespaces else namespaces,
+                    resources="*" if "*" in resources else resources,
                     selector=selector,
                     prometheus_url=prometheus_url,
                     prometheus_auth_header=prometheus_auth_header,