From b8943d0c2c2346826fa4d97a7e254e860293f44a Mon Sep 17 00:00:00 2001
From: Vaibhav <vaibhav.verma@traceable.ai>
Date: Wed, 15 Jan 2025 17:52:54 +0530
Subject: [PATCH 1/4] Allow user to update label rules

---
 .../service/LabelsConfigServiceImpl.java      | 37 ++++++++++---------
 .../service/LabelsConfigServiceImplTest.java  | 11 ++----
 2 files changed, 23 insertions(+), 25 deletions(-)

diff --git a/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java b/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
index 0a51b4b8..0459c036 100644
--- a/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
+++ b/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
@@ -21,6 +21,7 @@
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
@@ -194,15 +195,13 @@ public void getLabel(GetLabelRequest request, StreamObserver<GetLabelResponse> r
     RequestContext requestContext = RequestContext.CURRENT.get();
     String labelId = request.getId();
     try {
-      Label label;
-      if (systemLabelsIdLabelMap.containsKey(labelId)) {
-        label = systemLabelsIdLabelMap.get(labelId);
-      } else {
-        label =
-            labelStore
-                .getData(requestContext, labelId)
-                .orElseThrow(Status.NOT_FOUND::asRuntimeException);
-      }
+      Label label =
+          labelStore
+              .getData(requestContext, labelId)
+              .orElseGet(
+                  () ->
+                      Optional.ofNullable(systemLabelsIdLabelMap.get(labelId))
+                          .orElseThrow(Status.NOT_FOUND::asRuntimeException));
       responseObserver.onNext(GetLabelResponse.newBuilder().setLabel(label).build());
       responseObserver.onCompleted();
     } catch (Exception e) {
@@ -214,12 +213,17 @@ public void getLabel(GetLabelRequest request, StreamObserver<GetLabelResponse> r
   public void getLabels(
       GetLabelsRequest request, StreamObserver<GetLabelsResponse> responseObserver) {
     RequestContext requestContext = RequestContext.CURRENT.get();
-    List<Label> allLabels = new ArrayList<>(systemLabels);
     List<Label> tenantLabels =
         labelStore.getAllObjects(requestContext).stream()
             .map(ContextualConfigObject::getData)
             .collect(Collectors.toUnmodifiableList());
-    allLabels.addAll(tenantLabels);
+    Map<String, Label> tenantLabelsMap =
+        tenantLabels.stream()
+            .collect(Collectors.toUnmodifiableMap(Label::getId, Function.identity()));
+    List<Label> allLabels = new ArrayList<>(tenantLabels);
+    systemLabels.stream()
+        .filter(label -> !tenantLabelsMap.containsKey(label.getId()))
+        .forEach(allLabels::add);
     responseObserver.onNext(GetLabelsResponse.newBuilder().addAllLabels(allLabels).build());
     responseObserver.onCompleted();
   }
@@ -230,11 +234,7 @@ public void updateLabel(
     try {
       LabelData updateLabelData = request.getData();
       RequestContext requestContext = RequestContext.CURRENT.get();
-      if (systemLabelsIdLabelMap.containsKey(request.getId())) {
-        // Updating a system label will error
-        responseObserver.onError(new StatusRuntimeException(Status.INVALID_ARGUMENT));
-        return;
-      }
+
       if (isDuplicateKey(requestContext, request.getId(), updateLabelData.getKey())) {
         responseObserver.onError(new StatusRuntimeException(Status.ALREADY_EXISTS));
         return;
@@ -242,7 +242,10 @@ public void updateLabel(
       Label oldLabel =
           labelStore
               .getData(requestContext, request.getId())
-              .orElseThrow(Status.NOT_FOUND::asRuntimeException);
+              .orElseGet(
+                  () ->
+                      Optional.ofNullable(systemLabelsIdLabelMap.get(request.getId()))
+                          .orElseThrow(Status.NOT_FOUND::asRuntimeException));
       Label updateLabel = oldLabel.toBuilder().setData(updateLabelData).build();
       Label updateLabelInRes = labelStore.upsertObject(requestContext, updateLabel).getData();
       responseObserver.onNext(UpdateLabelResponse.newBuilder().setLabel(updateLabelInRes).build());
diff --git a/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java b/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
index cca2ae14..aee69c1a 100644
--- a/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
+++ b/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
@@ -1,5 +1,6 @@
 package org.hypertrace.label.config.service;
 
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.Mockito.mock;
@@ -333,15 +334,9 @@ void test_system_updateLabel() {
       UpdateLabelRequest updateLabelRequest =
           UpdateLabelRequest.newBuilder()
               .setId(systemLabel.getId())
-              .setData(LabelData.newBuilder().setKey("1").build())
+              .setData(LabelData.newBuilder().setKey(systemLabel.getData().getKey() + "_1").build())
               .build();
-      Throwable exception =
-          assertThrows(
-              StatusRuntimeException.class,
-              () -> {
-                labelConfigStub.updateLabel(updateLabelRequest);
-              });
-      assertEquals(Status.INVALID_ARGUMENT, Status.fromThrowable(exception));
+      assertDoesNotThrow(() -> labelConfigStub.updateLabel(updateLabelRequest));
     }
     for (Label systemLabel : systemLabels) {
       UpdateLabelRequest updateLabelRequest =

From 071d313a94242e07e845b800dab5e947ff1ad7e0 Mon Sep 17 00:00:00 2001
From: Vaibhav <vaibhav.verma@traceable.ai>
Date: Fri, 17 Jan 2025 16:13:29 +0530
Subject: [PATCH 2/4] add unit test

---
 .../label/config/service/LabelsConfigServiceImplTest.java       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java b/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
index aee69c1a..dfbc5b23 100644
--- a/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
+++ b/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
@@ -334,7 +334,7 @@ void test_system_updateLabel() {
       UpdateLabelRequest updateLabelRequest =
           UpdateLabelRequest.newBuilder()
               .setId(systemLabel.getId())
-              .setData(LabelData.newBuilder().setKey(systemLabel.getData().getKey() + "_1").build())
+              .setData(LabelData.newBuilder().setKey(systemLabel.getData().getKey()).build())
               .build();
       assertDoesNotThrow(() -> labelConfigStub.updateLabel(updateLabelRequest));
     }

From c3ba749ddc120dc8d1e54198d15d82bffc88f10a Mon Sep 17 00:00:00 2001
From: Vaibhav <vaibhav.verma@traceable.ai>
Date: Mon, 20 Jan 2025 14:19:11 +0530
Subject: [PATCH 3/4] remove usage of system label key map and add unit test

---
 .../config/service/LabelsConfigServiceImpl.java | 12 ++++--------
 .../service/LabelsConfigServiceImplTest.java    | 17 +++++++++++++++--
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java b/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
index 0459c036..19220a94 100644
--- a/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
+++ b/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
@@ -53,7 +53,6 @@ public class LabelsConfigServiceImpl extends LabelsConfigServiceGrpc.LabelsConfi
   private final LabelStore labelStore;
   private final List<Label> systemLabels;
   private final Map<String, Label> systemLabelsIdLabelMap;
-  private final Map<String, Label> systemLabelsKeyLabelMap;
 
   public LabelsConfigServiceImpl(
       Channel configChannel, Config config, ConfigChangeEventGenerator configChangeEventGenerator) {
@@ -75,13 +74,9 @@ public LabelsConfigServiceImpl(
       systemLabels = buildSystemLabelList(systemLabelsObjectList);
       systemLabelsIdLabelMap =
           systemLabels.stream().collect(Collectors.toUnmodifiableMap(Label::getId, identity()));
-      systemLabelsKeyLabelMap =
-          systemLabels.stream()
-              .collect(Collectors.toUnmodifiableMap(label -> label.getData().getKey(), identity()));
     } else {
       systemLabels = Collections.emptyList();
       systemLabelsIdLabelMap = Collections.emptyMap();
-      systemLabelsKeyLabelMap = Collections.emptyMap();
     }
   }
 
@@ -287,10 +282,11 @@ private boolean isDuplicateKey(RequestContext requestContext, String key) {
   }
 
   private Map<String, Label> getLabelsMap(RequestContext requestContext) {
-    Map<String, Label> existingLabelsMap = new HashMap<>(systemLabelsKeyLabelMap);
+    Map<String, Label> existingLabelsMap = new HashMap<>(systemLabelsIdLabelMap);
     labelStore.getAllObjects(requestContext).stream()
         .map(ContextualConfigObject::getData)
-        .forEach(label -> existingLabelsMap.put(label.getData().getKey(), label));
-    return Collections.unmodifiableMap(existingLabelsMap);
+        .forEach(label -> existingLabelsMap.put(label.getId(), label));
+    return existingLabelsMap.values().stream()
+        .collect(Collectors.toUnmodifiableMap(label -> label.getData().getKey(), identity()));
   }
 }
diff --git a/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java b/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
index dfbc5b23..5c9566c6 100644
--- a/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
+++ b/labels-config-service-impl/src/test/java/org/hypertrace/label/config/service/LabelsConfigServiceImplTest.java
@@ -330,19 +330,22 @@ void test_updateLabel() {
   @Test
   void test_system_updateLabel() {
     List<Label> createdLabelsList = createLabels();
+    // updating system label key, should not throw error.
     for (Label systemLabel : systemLabels) {
       UpdateLabelRequest updateLabelRequest =
           UpdateLabelRequest.newBuilder()
               .setId(systemLabel.getId())
-              .setData(LabelData.newBuilder().setKey(systemLabel.getData().getKey()).build())
+              .setData(LabelData.newBuilder().setKey(systemLabel.getData().getKey() + "_1").build())
               .build();
       assertDoesNotThrow(() -> labelConfigStub.updateLabel(updateLabelRequest));
     }
+
+    // updating existing label with updated system label key, should throw error.
     for (Label systemLabel : systemLabels) {
       UpdateLabelRequest updateLabelRequest =
           UpdateLabelRequest.newBuilder()
               .setId(createdLabelsList.get(0).getId())
-              .setData(LabelData.newBuilder().setKey(systemLabel.getData().getKey()).build())
+              .setData(LabelData.newBuilder().setKey(systemLabel.getData().getKey() + "_1").build())
               .build();
       Throwable exception =
           assertThrows(
@@ -352,6 +355,16 @@ void test_system_updateLabel() {
               });
       assertEquals(Status.ALREADY_EXISTS, Status.fromThrowable(exception));
     }
+
+    // updating existing label with system label keys present in config, should not throw error.
+    for (Label systemLabel : systemLabels) {
+      UpdateLabelRequest updateLabelRequest =
+          UpdateLabelRequest.newBuilder()
+              .setId(createdLabelsList.get(0).getId())
+              .setData(LabelData.newBuilder().setKey(systemLabel.getData().getKey()).build())
+              .build();
+      assertDoesNotThrow(() -> labelConfigStub.updateLabel(updateLabelRequest));
+    }
   }
 
   @Test

From 31d669c87aae2ce7a5f8fd284c9bc460cccd4479 Mon Sep 17 00:00:00 2001
From: Vaibhav <vaibhav.verma@traceable.ai>
Date: Mon, 20 Jan 2025 14:20:48 +0530
Subject: [PATCH 4/4] rename variable

---
 .../label/config/service/LabelsConfigServiceImpl.java       | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java b/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
index 19220a94..1c54da97 100644
--- a/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
+++ b/labels-config-service-impl/src/main/java/org/hypertrace/label/config/service/LabelsConfigServiceImpl.java
@@ -282,11 +282,11 @@ private boolean isDuplicateKey(RequestContext requestContext, String key) {
   }
 
   private Map<String, Label> getLabelsMap(RequestContext requestContext) {
-    Map<String, Label> existingLabelsMap = new HashMap<>(systemLabelsIdLabelMap);
+    Map<String, Label> mergedLabels = new HashMap<>(systemLabelsIdLabelMap);
     labelStore.getAllObjects(requestContext).stream()
         .map(ContextualConfigObject::getData)
-        .forEach(label -> existingLabelsMap.put(label.getId(), label));
-    return existingLabelsMap.values().stream()
+        .forEach(label -> mergedLabels.put(label.getId(), label));
+    return mergedLabels.values().stream()
         .collect(Collectors.toUnmodifiableMap(label -> label.getData().getKey(), identity()));
   }
 }