Add columns to csv (#359)

.gitignore

@@ -133,4 +133,4 @@ dmypy.json
 .DS_Store
 robusta_lib
 .idea
-.vscode
+.vscode

robusta_krr/core/models/config.py

@@ -58,10 +58,11 @@ class Config(pd.BaseSettings):
     strategy: str
     log_to_stderr: bool
     width: Optional[int] = pd.Field(None, ge=1)
+    show_severity: bool = True
 
     # Output Settings
     file_output: Optional[str] = pd.Field(None)
-    file_output_dynamic = bool = pd.Field(False)
+    file_output_dynamic: bool = pd.Field(False)
     slack_output: Optional[str] = pd.Field(None)
 
     other_args: dict[str, Any]
@@ -105,7 +106,7 @@ def validate_namespaces(cls, v: Union[list[str], Literal["*"]]) -> Union[list[st
             for val in v:
                 if val.startswith("*"):
                     raise ValueError("Namespace's values cannot start with an asterisk (*)")
-            
+
         return [val.lower() for val in v]
 
     @pd.validator("resources", pre=True)

robusta_krr/core/models/result.py

@@ -17,8 +17,8 @@ class Recommendation(pd.BaseModel):
 
 
 class ResourceRecommendation(pd.BaseModel):
-    requests: dict[ResourceType, RecommendationValue]
-    limits: dict[ResourceType, RecommendationValue]
+    requests: dict[ResourceType, Union[RecommendationValue, Recommendation]]
+    limits: dict[ResourceType, Union[RecommendationValue, Recommendation]]
     info: dict[ResourceType, Optional[str]]
 
 
@@ -40,6 +40,7 @@ def calculate(cls, object: K8sObjectData, recommendation: ResourceAllocations) -
 
                 current_severity = Severity.calculate(current, recommended, resource_type)
 
+                #TODO: consider... changing field after model created doesn't validate it.
                 getattr(recommendation_processed, selector)[resource_type] = Recommendation(
                     value=recommended, severity=current_severity
                 )

robusta_krr/formatters/csv.py

@@ -1,14 +1,31 @@
-import itertools
 import csv
-import logging
 import io
+import itertools
+import logging
+from typing import Any
 
 from robusta_krr.core.abstract import formatters
-from robusta_krr.core.models.allocations import RecommendationValue, format_recommendation_value, format_diff, NONE_LITERAL, NAN_LITERAL
+from robusta_krr.core.models.allocations import NONE_LITERAL, format_diff, format_recommendation_value
+from robusta_krr.core.models.config import settings
 from robusta_krr.core.models.result import ResourceScan, ResourceType, Result
 
 logger = logging.getLogger("krr")
 
+
+NAMESPACE_HEADER = "Namespace"
+NAME_HEADER = "Name"
+PODS_HEADER = "Pods"
+OLD_PODS_HEADER = "Old Pods"
+TYPE_HEADER = "Type"
+CONTAINER_HEADER = "Container"
+CLUSTER_HEADER = "Cluster"
+SEVERITY_HEADER = "Severity"
+
+RESOURCE_DIFF_HEADER = "{resource_name} Diff"
+RESOURCE_REQUESTS_HEADER = "{resource_name} Requests"
+RESOURCE_LIMITS_HEADER = "{resource_name} Limits"
+
+
 def _format_request_str(item: ResourceScan, resource: ResourceType, selector: str) -> str:
     allocated = getattr(item.object.allocations, selector)[resource]
     recommended = getattr(item.recommended, selector)[resource]
@@ -20,12 +37,8 @@ def _format_request_str(item: ResourceScan, resource: ResourceType, selector: st
     if diff != "":
         diff = f"({diff}) "
 
-    return (
-        diff
-        + format_recommendation_value(allocated)
-        + " -> "
-        + format_recommendation_value(recommended.value)
-    )
+    return diff + format_recommendation_value(allocated) + " -> " + format_recommendation_value(recommended.value)
+
 
 def _format_total_diff(item: ResourceScan, resource: ResourceType, pods_current: int) -> str:
     selector = "requests"
@@ -34,43 +47,57 @@ def _format_total_diff(item: ResourceScan, resource: ResourceType, pods_current:
 
     return format_diff(allocated, recommended, selector, pods_current)
 
+
 @formatters.register("csv")
 def csv_exporter(result: Result) -> str:
     # We need to order the resource columns so that they are in the format of Namespace,Name,Pods,Old Pods,Type,Container,CPU Diff,CPU Requests,CPU Limits,Memory Diff,Memory Requests,Memory Limits
-    resource_columns = []
+    csv_columns = ["Namespace", "Name", "Pods", "Old Pods", "Type", "Container"]
+
+    if settings.show_cluster_name:
+        csv_columns.insert(0, "Cluster")
+
+    if settings.show_severity:
+        csv_columns.append("Severity")
+
     for resource in ResourceType:
-        resource_columns.append(f"{resource.name} Diff")
-        resource_columns.append(f"{resource.name} Requests")
-        resource_columns.append(f"{resource.name} Limits")
+        csv_columns.append(RESOURCE_DIFF_HEADER.format(resource_name=resource.name))
+        csv_columns.append(RESOURCE_REQUESTS_HEADER.format(resource_name=resource.name))
+        csv_columns.append(RESOURCE_LIMITS_HEADER.format(resource_name=resource.name))
 
     output = io.StringIO()
-    csv_writer = csv.writer(output)
-    csv_writer.writerow([
-        "Namespace", "Name", "Pods", "Old Pods", "Type", "Container",
-        *resource_columns
-    ])
+    csv_writer = csv.DictWriter(output, csv_columns, extrasaction="ignore")
+    csv_writer.writeheader()
 
     for _, group in itertools.groupby(
         enumerate(result.scans), key=lambda x: (x[1].object.cluster, x[1].object.namespace, x[1].object.name)
     ):
         group_items = list(group)
 
-        for j, (i, item) in enumerate(group_items):
+        for j, (_, item) in enumerate(group_items):
             full_info_row = j == 0
 
-            row = [
-                item.object.namespace if full_info_row else "",
-                item.object.name if full_info_row else "",
-                f"{item.object.current_pods_count}" if full_info_row else "",
-                f"{item.object.deleted_pods_count}" if full_info_row else "",
-                item.object.kind if full_info_row else "",
-                item.object.container,
-            ]
+            row: dict[str, Any] = {
+                NAMESPACE_HEADER: item.object.namespace if full_info_row else "",
+                NAME_HEADER: item.object.name if full_info_row else "",
+                PODS_HEADER: f"{item.object.current_pods_count}" if full_info_row else "",
+                OLD_PODS_HEADER: f"{item.object.deleted_pods_count}" if full_info_row else "",
+                TYPE_HEADER: item.object.kind if full_info_row else "",
+                CONTAINER_HEADER: item.object.container,
+                SEVERITY_HEADER: item.severity,
+                CLUSTER_HEADER: item.object.cluster,
+            }
 
             for resource in ResourceType:
-                row.append(_format_total_diff(item, resource, item.object.current_pods_count))
-                row += [_format_request_str(item, resource, selector) for selector in ["requests", "limits"]]
+                row[RESOURCE_DIFF_HEADER.format(resource_name=resource.name)] = _format_total_diff(
+                    item, resource, item.object.current_pods_count
+                )
+                row[RESOURCE_REQUESTS_HEADER.format(resource_name=resource.name)] = _format_request_str(
+                    item, resource, "requests"
+                )
+                row[RESOURCE_LIMITS_HEADER.format(resource_name=resource.name)] = _format_request_str(
+                    item, resource, "limits"
+                )
 
             csv_writer.writerow(row)
-             
+
     return output.getvalue()

robusta_krr/main.py

@@ -7,6 +7,7 @@
 from typing import List, Optional
 from uuid import UUID
 
+import click
 import typer
 import urllib3
 from pydantic import ValidationError  # noqa: F401
@@ -19,7 +20,12 @@
 from robusta_krr.core.runner import Runner
 from robusta_krr.utils.version import get_version
 
-app = typer.Typer(pretty_exceptions_show_locals=False, pretty_exceptions_short=True, no_args_is_help=True, help="IMPORTANT: Run `krr simple --help` to see all cli flags!")
+app = typer.Typer(
+    pretty_exceptions_show_locals=False,
+    pretty_exceptions_short=True,
+    no_args_is_help=True,
+    help="IMPORTANT: Run `krr simple --help` to see all cli flags!",
+)
 
 # NOTE: Disable insecure request warnings, as it might be expected to use self-signed certificates inside the cluster
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
@@ -216,7 +222,16 @@ def run_strategy(
                     rich_help_panel="Logging Settings",
                 ),
                 show_cluster_name: bool = typer.Option(
-                    False, "--show-cluster-name", help="In table output, always show the cluster name even for a single cluster", rich_help_panel="Output Settings"
+                    False,
+                    "--show-cluster-name",
+                    help="In table output, always show the cluster name even for a single cluster",
+                    rich_help_panel="Output Settings",
+                ),
+                show_severity: bool = typer.Option(
+                    True,
+                    " /--exclude-severity",
+                    help="Whether to include the severity in the output or not",
+                    rich_help_panel="Output Settings",
                 ),
                 verbose: bool = typer.Option(
                     False, "--verbose", "-v", help="Enable verbose mode", rich_help_panel="Logging Settings"
@@ -234,10 +249,16 @@ def run_strategy(
                     rich_help_panel="Logging Settings",
                 ),
                 file_output: Optional[str] = typer.Option(
-                    None, "--fileoutput", help="Filename to write output to (if not specified, file output is disabled)", rich_help_panel="Output Settings"
+                    None,
+                    "--fileoutput",
+                    help="Filename to write output to (if not specified, file output is disabled)",
+                    rich_help_panel="Output Settings",
                 ),
                 file_output_dynamic: bool = typer.Option(
-                    False, "--fileoutput-dynamic", help="Ignore --fileoutput and write files to the current directory in the format krr-{datetime}.{format} (e.g. krr-20240518223924.csv)", rich_help_panel="Output Settings"
+                    False,
+                    "--fileoutput-dynamic",
+                    help="Ignore --fileoutput and write files to the current directory in the format krr-{datetime}.{format} (e.g. krr-20240518223924.csv)",
+                    rich_help_panel="Output Settings",
                 ),
                 slack_output: Optional[str] = typer.Option(
                     None,
@@ -248,6 +269,8 @@ def run_strategy(
                 **strategy_args,
             ) -> None:
                 f"""Run KRR using the `{_strategy_name}` strategy"""
+                if not show_severity and format != "csv":
+                    raise click.BadOptionUsage("--exclude-severity", "--exclude-severity works only with format=csv")
 
                 try:
                     config = Config(
@@ -284,6 +307,7 @@ def run_strategy(
                         file_output=file_output,
                         file_output_dynamic=file_output_dynamic,
                         slack_output=slack_output,
+                        show_severity=show_severity,
                         strategy=_strategy_name,
                         other_args=strategy_args,
                     )