From 31fc31e0ef2277f514191c14d347b05541ee0789 Mon Sep 17 00:00:00 2001
From: Kshitiz Saxena <84707889+saxenakshitiz@users.noreply.github.com>
Date: Thu, 6 Jun 2024 19:26:41 +0530
Subject: [PATCH] chore: add regex capture to label application rule (#214)

---
 .../service/v1/label_application_rule.proto      |  1 +
 .../LabelApplicationRuleValidatorImpl.java       |  6 +++++-
 .../config/validation/RegexValidator.java        | 16 ++++++++++++++++
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/label-application-rule-config-service-api/src/main/proto/org/hypertrace/label/application/rule/config/service/v1/label_application_rule.proto b/label-application-rule-config-service-api/src/main/proto/org/hypertrace/label/application/rule/config/service/v1/label_application_rule.proto
index 4489dd1c..81af4850 100644
--- a/label-application-rule-config-service-api/src/main/proto/org/hypertrace/label/application/rule/config/service/v1/label_application_rule.proto
+++ b/label-application-rule-config-service-api/src/main/proto/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/label-application-rule-config-service-impl/src/main/java/org/hypertrace/label/application/rule/config/service/LabelApplicationRuleValidatorImpl.java b/label-application-rule-config-service-impl/src/main/java/org/hypertrace/label/application/rule/config/service/LabelApplicationRuleValidatorImpl.java
index f9934b84..f03a369a 100644
--- a/label-application-rule-config-service-impl/src/main/java/org/hypertrace/label/application/rule/config/service/LabelApplicationRuleValidatorImpl.java
+++ b/label-application-rule-config-service-impl/src/main/java/org/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
index ed159ab8..034dc648 100644
--- 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;
+  }
 }