chore: add regex capture to label application rule (#214)

hypertrace · Jun 6, 2024 · 31fc31e · 31fc31e
1 parent c0dd217
commit 31fc31e
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 1 deletion.
diff --git a/...roto/org/hypertrace/label/application/rule/config/service/v1/label_application_rule.proto b/...roto/org/hypertrace/label/application/rule/config/service/v1/label_application_rule.proto
@@ -110,6 +110,7 @@ message LabelApplicationRuleData {
         string key = 1;
         optional string json_path = 2;
         optional string alias = 3;
+        optional string regex_capture = 4;
       }
     }
   }

diff --git a/...g/hypertrace/label/application/rule/config/service/LabelApplicationRuleValidatorImpl.java b/...g/hypertrace/label/application/rule/config/service/LabelApplicationRuleValidatorImpl.java
@@ -243,9 +243,10 @@ public void validateLabelExpression(Action.DynamicLabel dynamicLabel) {
         .getTokenExtractionRulesList()
         .forEach(
             tokenExtractionRule -> {
-              validKeys.add(tokenExtractionRule.getKey());
               if (tokenExtractionRule.hasAlias()) {
                 validKeys.add(tokenExtractionRule.getAlias());
+              } else {
+                validKeys.add(tokenExtractionRule.getKey());
               }
             });
     Pattern pattern = Pattern.compile("\\{(\\\\}|[^}])*}");
@@ -264,6 +265,9 @@ public void validateLabelExpression(Action.DynamicLabel dynamicLabel) {
   private void validateTokenExtractionRule(
       Action.DynamicLabel.TokenExtractionRule tokenExtractionRule) {
     validateNonDefaultPresenceOrThrow(tokenExtractionRule, tokenExtractionRule.KEY_FIELD_NUMBER);
+    if (tokenExtractionRule.hasRegexCapture()) {
+      RegexValidator.validateCaptureGroupCount(tokenExtractionRule.getRegexCapture(), 1);
+    }
   }
 
   private void throwInvalidArgumentException(String description) {

diff --git a/validation-utils/src/main/java/org/hypertrace/config/validation/RegexValidator.java b/validation-utils/src/main/java/org/hypertrace/config/validation/RegexValidator.java
@@ -16,4 +16,20 @@ public static Status validate(String regexPattern) {
     }
     return Status.OK;
   }
+
+  public static Status validateCaptureGroupCount(String regexPattern, int expectedCount) {
+    // compiling an invalid regex throws PatternSyntaxException
+    try {
+      Pattern pattern = Pattern.compile(regexPattern);
+      if (pattern.groupCount() != expectedCount) {
+        return Status.INVALID_ARGUMENT.withDescription(
+            "Regex group count should be: " + expectedCount);
+      }
+    } catch (PatternSyntaxException e) {
+      return Status.INVALID_ARGUMENT
+          .withCause(e)
+          .withDescription("Invalid Regex pattern: " + regexPattern);
+    }
+    return Status.OK;
+  }
 }