Merge pull request #209 from robusta-dev/openshift

Add Openshift support with loading sa token

6 files changed, 45 insertions, 1 deletions

diff --git a/robusta_krr/core/integrations/openshift/__init__.py b/robusta_krr/core/integrations/openshift/__init__.py
new file mode 100644
index 0000000..c54df54
--- /dev/null
+++ b/robusta_krr/core/integrations/openshift/__init__.py
@@ -0,0 +1,3 @@
+from .token import TOKEN_LOCATION, load_token
+
+__all__ = ["TOKEN_LOCATION", "load_token"]
diff --git a/robusta_krr/core/integrations/openshift/token.py b/robusta_krr/core/integrations/openshift/token.py
new file mode 100644
index 0000000..54a599f
--- /dev/null
+++ b/robusta_krr/core/integrations/openshift/token.py
@@ -0,0 +1,17 @@
+from typing import Optional
+
+from robusta_krr.core.models.config import settings
+
+# NOTE: This one should be mounted if openshift is enabled (done by Robusta Runner)
+TOKEN_LOCATION = '/var/run/secrets/kubernetes.io/serviceaccount/token'
+
+
+def load_token() -> Optional[str]:
+    if not settings.openshift:
+        return None
+
+    try:
+        with open(TOKEN_LOCATION, 'r') as file:
+            return file.read()
+    except FileNotFoundError:
+        return None
diff --git a/robusta_krr/core/integrations/prometheus/metrics_service/prometheus_metrics_service.py b/robusta_krr/core/integrations/prometheus/metrics_service/prometheus_metrics_service.py
index 091041e..761909b 100644
--- a/robusta_krr/core/integrations/prometheus/metrics_service/prometheus_metrics_service.py
+++ b/robusta_krr/core/integrations/prometheus/metrics_service/prometheus_metrics_service.py
@@ -9,6 +9,7 @@ from prometheus_api_client import PrometheusApiClientException
 from prometrix import PrometheusNotFound, get_custom_prometheus_connect
 
 from robusta_krr.core.abstract.strategies import PodsTimeData
+from robusta_krr.core.integrations import openshift
 from robusta_krr.core.models.config import settings
 from robusta_krr.core.models.objects import K8sObjectData, PodData
 from robusta_krr.utils.batched import batched
@@ -65,6 +66,16 @@ class PrometheusMetricsService(MetricsService):
         self.auth_header = settings.prometheus_auth_header
         self.ssl_enabled = settings.prometheus_ssl_enabled
 
+        if settings.openshift:
+            logging.info("Openshift flag is set, trying to load token from service account.")
+            openshift_token = openshift.load_token()
+
+            if openshift_token:
+                logging.info("Openshift token is loaded successfully.")
+                self.auth_header = self.auth_header or f"Bearer {openshift_token}"
+            else:
+                logging.warning("Openshift token is not found, trying to connect without it.")
+
         self.prometheus_discovery = self.service_discovery(api_client=self.api_client)
 
         self.url = settings.prometheus_url
diff --git a/robusta_krr/core/models/config.py b/robusta_krr/core/models/config.py
index 223b5df..54a9ed3 100644
--- a/robusta_krr/core/models/config.py
+++ b/robusta_krr/core/models/config.py
@@ -45,6 +45,7 @@ class Config(pd.BaseSettings):
     eks_service_name: Optional[str] = pd.Field(None)
     eks_managed_prom_region: Optional[str] = pd.Field(None)
     coralogix_token: Optional[str] = pd.Field(None)
+    openshift: bool = pd.Field(False)
 
     # Threading settings
     max_workers: int = pd.Field(6, ge=1)
diff --git a/robusta_krr/core/runner.py b/robusta_krr/core/runner.py
index 312efdd..787c3a4 100644
--- a/robusta_krr/core/runner.py
+++ b/robusta_krr/core/runner.py
@@ -240,6 +240,7 @@ class Runner:
                 self._strategy.settings.timeframe_duration = min_step
 
             result = await self._collect_result()
+            logger.info("Result collected, displaying...")
             self._process_result(result)
         except ClusterNotSpecifiedException as e:
             logger.error(e)
diff --git a/robusta_krr/main.py b/robusta_krr/main.py
index a8e1c2a..aef3848 100644
--- a/robusta_krr/main.py
+++ b/robusta_krr/main.py
@@ -170,6 +170,13 @@ def load_commands() -> None:
                     help="Adds the token needed to query Coralogix managed prometheus.",
                     rich_help_panel="Prometheus Coralogix Settings",
                 ),
+                openshift: bool = typer.Option(
+                    False,
+                    "--openshift",
+                    help="Used when running by Robusta inside an OpenShift cluster.",
+                    rich_help_panel="Prometheus Openshift Settings",
+                    hidden=True,
+                ),
                 cpu_min_value: int = typer.Option(
                     10,
                     "--cpu-min",
@@ -206,7 +213,10 @@ def load_commands() -> None:
                     False, "--logtostderr", help="Pass logs to stderr", rich_help_panel="Logging Settings"
                 ),
                 width: Optional[int] = typer.Option(
-                    None, "--width", help="Width of the output. Will use console width by default.", rich_help_panel="Logging Settings"
+                    None,
+                    "--width",
+                    help="Width of the output. Will use console width by default.",
+                    rich_help_panel="Logging Settings",
                 ),
                 file_output: Optional[str] = typer.Option(
                     None, "--fileoutput", help="Print the output to a file", rich_help_panel="Output Settings"
@@ -241,6 +251,7 @@ def load_commands() -> None:
                         eks_secret_key=eks_secret_key,
                         eks_service_name=eks_service_name,
                         coralogix_token=coralogix_token,
+                        openshift=openshift,
                         max_workers=max_workers,
                         format=format,
                         verbose=verbose,