From ae207749de1ec486cd94b90ad3a4bbd9557222a9 Mon Sep 17 00:00:00 2001
From: NohaIhab <noha.ihab@canonical.com>
Date: Wed, 27 Nov 2024 10:44:25 +0200
Subject: [PATCH] feat: add config for ARGO_KEY_FORMAT envvar

---
 charms/kfp-ui/config.yaml                         | 5 +++++
 charms/kfp-ui/src/charm.py                        | 1 +
 charms/kfp-ui/src/components/pebble_components.py | 6 ++----
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/charms/kfp-ui/config.yaml b/charms/kfp-ui/config.yaml
index e439f866..e5658a35 100644
--- a/charms/kfp-ui/config.yaml
+++ b/charms/kfp-ui/config.yaml
@@ -11,6 +11,11 @@ options:
     type: boolean
     default: true
     description: Enable Argo log archive
+  argo-key-format:
+    type: string
+    # Must have the same value as the default of `keyformat` config in `argo-controller` charm.
+    default: "artifacts/{{workflow.name}}/{{workflow.creationTimestamp.Y}}/{{workflow.creationTimestamp.m}}/{{workflow.creationTimestamp.d}}/{{pod.name}}"
+    description: The keyFormat for pod logs artifacts stored
   disable-gke-metadata:
     type: boolean
     default: true
diff --git a/charms/kfp-ui/src/charm.py b/charms/kfp-ui/src/charm.py
index 577f551b..876b63f4 100755
--- a/charms/kfp-ui/src/charm.py
+++ b/charms/kfp-ui/src/charm.py
@@ -220,6 +220,7 @@ def __init__(self, *args):
                 inputs_getter=lambda: MlPipelineUiInputs(
                     ALLOW_CUSTOM_VISUALIZATIONS=self.model.config["allow-custom-visualizations"],
                     ARGO_ARCHIVE_LOGS=self.model.config["argo-archive-logs"],
+                    ARGO_KEY_FORMAT=self.model.config["argo-key-format"],
                     DISABLE_GKE_METADATA=self.model.config["disable-gke-metadata"],
                     FRONTEND_SERVER_NAMESPACE=self.model.name,
                     HIDE_SIDENAV=self.model.config["hide-sidenav"],
diff --git a/charms/kfp-ui/src/components/pebble_components.py b/charms/kfp-ui/src/components/pebble_components.py
index fab64d28..8d3c1a1c 100644
--- a/charms/kfp-ui/src/components/pebble_components.py
+++ b/charms/kfp-ui/src/components/pebble_components.py
@@ -13,6 +13,7 @@ class MlPipelineUiInputs:
 
     ALLOW_CUSTOM_VISUALIZATIONS: bool
     ARGO_ARCHIVE_LOGS: bool
+    ARGO_KEY_FORMAT: str
     DISABLE_GKE_METADATA: bool
     FRONTEND_SERVER_NAMESPACE: str
     HIDE_SIDENAV: bool
@@ -54,10 +55,7 @@ def get_layer(self) -> Layer:
                             "ARGO_ARCHIVE_BUCKETNAME": "mlpipeline",
                             "ARGO_ARCHIVE_LOGS": inputs.ARGO_ARCHIVE_LOGS,
                             "ARGO_ARCHIVE_PREFIX": "logs",
-                            # Must have the same value as the `keyFormat` specified in the
-                            # `argo-workflow-controller-configmap` ConfigMap owned by
-                            # the `argo-controller` charm.
-                            "ARGO_KEYFORMAT": "artifacts/{{workflow.name}}/{{workflow.creationTimestamp.Y}}/{{workflow.creationTimestamp.m}}/{{workflow.creationTimestamp.d}}/{{pod.name}}",  # noqa E501
+                            "ARGO_KEYFORMAT": inputs.ARGO_KEY_FORMAT,
                             # TODO: This should come from relation to kfp-profile-controller.
                             #  It is the name/port of the user-specific artifact accessor
                             "ARTIFACTS_SERVICE_PROXY_NAME": "ml-pipeline-ui-artifact",