Aligning replace function in Regex ext with Google SQL

maskri17 · copybara-github · commit 20d8ca8498e4 · 2025-06-20T16:24:26.000-07:00
PiperOrigin-RevId: 773880708
diff --git a/extensions/src/main/java/dev/cel/extensions/CelRegexExtensions.java b/extensions/src/main/java/dev/cel/extensions/CelRegexExtensions.java
@@ -78,7 +78,7 @@ enum Function {
                   String pattern = (String) args[1];
                   String replaceStr = (String) args[2];
                   long count = (long) args[3];
-                  return CelRegexExtensions.replace(target, pattern, replaceStr, count);
+                  return CelRegexExtensions.replaceN(target, pattern, replaceStr, count);
                 }))),
     EXTRACT(
         CelFunctionDecl.newFunctionDeclaration(
@@ -153,18 +153,20 @@ private static Pattern compileRegexPattern(String regex) {
   }
 
   private static String replace(String target, String regex, String replaceStr) {
-    Pattern pattern = compileRegexPattern(regex);
-    Matcher matcher = pattern.matcher(target);
-    return matcher.replaceAll(replaceStr);
+    return replaceN(target, regex, replaceStr, -1);
   }
 
-  private static String replace(String target, String regex, String replaceStr, long replaceCount) {
-    Pattern pattern = compileRegexPattern(regex);
-
+  private static String replaceN(
+      String target, String regex, String replaceStr, long replaceCount) {
     if (replaceCount == 0) {
       return target;
     }
+    // For all negative replaceCount, do a replaceAll
+    if (replaceCount < 0) {
+      replaceCount = -1;
+    }
 
+    Pattern pattern = compileRegexPattern(regex);
     Matcher matcher = pattern.matcher(target);
     StringBuffer sb = new StringBuffer();
     int counter = 0;
@@ -173,14 +175,59 @@ private static String replace(String target, String regex, String replaceStr, lo
       if (replaceCount != -1 && counter >= replaceCount) {
         break;
       }
-      matcher.appendReplacement(sb, replaceStr);
+
+      String processedReplacement = replaceStrValidator(matcher, replaceStr);
+      matcher.appendReplacement(sb, Matcher.quoteReplacement(processedReplacement));
       counter++;
     }
     matcher.appendTail(sb);
 
     return sb.toString();
   }
 
+  private static String replaceStrValidator(Matcher matcher, String replacement) {
+    StringBuilder sb = new StringBuilder();
+    for (int i = 0; i < replacement.length(); i++) {
+      char c = replacement.charAt(i);
+
+      if (c != '\\') {
+        sb.append(c);
+        continue;
+      }
+
+      if (i + 1 >= replacement.length()) {
+        throw new IllegalArgumentException("Invalid replacement string: \\ not allowed at end");
+      }
+
+      char nextChar = replacement.charAt(++i);
+
+      if (Character.isDigit(nextChar)) {
+        int groupNum = Character.digit(nextChar, 10);
+        int groupCount = matcher.groupCount();
+
+        if (groupNum > groupCount) {
+          throw new IllegalArgumentException(
+              "Replacement string references group "
+                  + groupNum
+                  + " but regex has only "
+                  + groupCount
+                  + " group(s)");
+        }
+
+        String groupValue = matcher.group(groupNum);
+        if (groupValue != null) {
+          sb.append(groupValue);
+        }
+      } else if (nextChar == '\\') {
+        sb.append('\\');
+      } else {
+        throw new IllegalArgumentException(
+            "Invalid replacement string: \\ must be followed by a digit");
+      }
+    }
+    return sb.toString();
+  }
+
   private static Optional<String> extract(String target, String regex) {
     Pattern pattern = compileRegexPattern(regex);
     Matcher matcher = pattern.matcher(target);
@@ -215,11 +262,12 @@ private static ImmutableList<String> extractAll(String target, String regex) {
     while (matcher.find()) {
       if (hasOneGroup) {
         String group = matcher.group(1);
-        // Add the captured group's content only if it's not null (e.g. optional group didn't match)
+        // Add the captured group's content only if it's not null
         if (group != null) {
           builder.add(group);
         }
-      } else { // No capturing groups (matcher.groupCount() == 0)
+      } else {
+        // No capturing groups
         builder.add(matcher.group(0));
       }
     }
diff --git a/extensions/src/main/java/dev/cel/extensions/README.md b/extensions/src/main/java/dev/cel/extensions/README.md
@@ -717,12 +717,15 @@ chance for collision.
 
 ### Replace
 
-The `regex.replace` function replaces all occurrences of a regex pattern in a
-string with a replacement string. Optionally, you can limit the number of
-replacements by providing a count argument. Both numeric ($N) and named
-(${name}) capture group references are supported in the replacement string, with
-validation for correctness. An error will be thrown for invalid regex or replace
-string.
+The `regex.replace` function replaces all non-overlapping substring of a regex
+pattern in the target string with a replacement string. Optionally, you can
+limit the number of replacements by providing a count argument. When the count
+is a negative number, the function acts as replace all. Only numeric (\N)
+capture group references are supported in the replacement string, with
+validation for correctness. Backslashed-escaped digits (\1 to \9) within the
+replacement argument can be used to insert text matching the corresponding
+parenthesized group in the regexp pattern. An error will be thrown for invalid
+regex or replace string.
 
 ```
 regex.replace(target: string, pattern: string, replacement: string) -> string
@@ -732,14 +735,16 @@ regex.replace(target: string, pattern: string, replacement: string, count: int)
 Examples:
 
 ```
+regex.replace('hello world hello', 'hello', 'hi') == 'hi world hi'
 regex.replace('banana', 'a', 'x', 0) == 'banana'
 regex.replace('banana', 'a', 'x', 1) == 'bxnana'
 regex.replace('banana', 'a', 'x', 2) == 'bxnxna'
-regex.replace('foo bar', '(fo)o (ba)r', '$2 $1') == 'ba fo'
+regex.replace('banana', 'a', 'x', -12) == 'bxnxnx'
+regex.replace('foo bar', '(fo)o (ba)r', '\\2 \\1') == 'ba fo'
 
 regex.replace('test', '(.)', '$2') \\ Runtime Error invalid replace string
 regex.replace('foo bar', '(', '$2 $1') \\ Runtime Error invalid regex string
-regex.replace('id=123', 'id=(?P<value>\\\\d+)', 'value: ${values}') \\ Runtime Error invalid replace string
+regex.replace('id=123', 'id=(?P<value>\\\\d+)', 'value: \\values') \\ Runtime Error invalid replace string
 
 ```
 
diff --git a/extensions/src/test/java/dev/cel/extensions/CelRegexExtensionsTest.java b/extensions/src/test/java/dev/cel/extensions/CelRegexExtensionsTest.java
@@ -39,23 +39,31 @@ public final class CelRegexExtensionsTest {
       CelRuntimeFactory.standardCelRuntimeBuilder().addLibraries(CelExtensions.regex()).build();
 
   @Test
-  @TestParameters("{target: 'foo bar', regex: '(fo)o (ba)r', replaceStr: '$2 $1', res: 'ba fo'}")
+  @TestParameters("{target: 'abc', regex: '^', replaceStr: 'start_', res: 'start_abc'}")
+  @TestParameters("{target: 'abc', regex: '$', replaceStr: '_end', res: 'abc_end'}")
+  @TestParameters("{target: 'a-b', regex: '\\\\b', replaceStr: '|', res: '|a|-|b|'}")
+  @TestParameters(
+      "{target: 'foo bar', regex: '(fo)o (ba)r', replaceStr: '\\\\2 \\\\1', res: 'ba fo'}")
+  @TestParameters("{target: 'foo bar', regex: 'foo', replaceStr: '\\\\\\\\', res: '\\ bar'}")
   @TestParameters("{target: 'banana', regex: 'ana', replaceStr: 'x', res: 'bxna'}")
-  @TestParameters("{target: 'abc', regex: 'b(.)', replaceStr: 'x$1', res: 'axc'}")
+  @TestParameters("{target: 'abc', regex: 'b(.)', replaceStr: 'x\\\\1', res: 'axc'}")
   @TestParameters(
       "{target: 'hello world hello', regex: 'hello', replaceStr: 'hi', res: 'hi world hi'}")
+  @TestParameters("{target: 'ac', regex: 'a(b)?c', replaceStr: '[\\\\1]', res: '[]'}")
   @TestParameters("{target: 'apple pie', regex: 'p', replaceStr: 'X', res: 'aXXle Xie'}")
   @TestParameters(
       "{target: 'remove all spaces', regex: '\\\\s', replaceStr: '', res: 'removeallspaces'}")
   @TestParameters("{target: 'digit:99919291992', regex: '\\\\d+', replaceStr: '3', res: 'digit:3'}")
   @TestParameters(
-      "{target: 'foo bar baz', regex: '\\\\w+', replaceStr: '($0)', res: '(foo) (bar) (baz)'}")
+      "{target: 'foo bar baz', regex: '\\\\w+', replaceStr: '(\\\\0)', res: '(foo) (bar) (baz)'}")
   @TestParameters("{target: '', regex: 'a', replaceStr: 'b', res: ''}")
   @TestParameters(
       "{target: 'User: Alice, Age: 30', regex: 'User: (?P<name>\\\\w+), Age: (?P<age>\\\\d+)',"
-          + " replaceStr: '${name} is ${age} years old', res: 'Alice is 30 years old'}")
+          + " replaceStr: '${name} is ${age} years old', res: '${name} is ${age} years old'}")
   @TestParameters(
-      "{target: 'abc', regex: '(?P<letter>b)', replaceStr: '[${letter}]', res: 'a[b]c'}")
+      "{target: 'User: Alice, Age: 30', regex: 'User: (?P<name>\\\\w+), Age: (?P<age>\\\\d+)',"
+          + " replaceStr: '\\\\1 is \\\\2 years old', res: 'Alice is 30 years old'}")
+  @TestParameters("{target: 'hello ☃', regex: '☃', replaceStr: '❄', res: 'hello ❄'}")
   public void replaceAll_success(String target, String regex, String replaceStr, String res)
       throws Exception {
     String expr = String.format("regex.replace('%s', '%s', '%s')", target, regex, replaceStr);
@@ -70,8 +78,8 @@ public void replaceAll_success(String target, String regex, String replaceStr, S
   public void replace_nested_success() throws Exception {
     String expr =
         "regex.replace("
-            + "    regex.replace('%(foo) %(bar) %2','%\\\\((\\\\w+)\\\\)','\\\\${$1}'),"
-            + "    '%(\\\\d+)', '\\\\$$1')";
+            + "    regex.replace('%(foo) %(bar) %2','%\\\\((\\\\w+)\\\\)','${\\\\1}'),"
+            + "    '%(\\\\d+)', '$\\\\1')";
     CelRuntime.Program program = RUNTIME.createProgram(COMPILER.compile(expr).getAst());
 
     Object result = program.eval();
@@ -85,19 +93,18 @@ public void replace_nested_success() throws Exception {
   @TestParameters("{t: 'banana', re: 'a', rep: 'x', i: 2, res: 'bxnxna'}")
   @TestParameters("{t: 'banana', re: 'a', rep: 'x', i: 100, res: 'bxnxnx'}")
   @TestParameters("{t: 'banana', re: 'a', rep: 'x', i: -1, res: 'bxnxnx'}")
-  @TestParameters("{t: 'banana', re: 'a', rep: 'x', i: -100, res: 'banana'}")
+  @TestParameters("{t: 'banana', re: 'a', rep: 'x', i: -100, res: 'bxnxnx'}")
   @TestParameters(
-      "{t: 'cat-dog dog-cat cat-dog dog-cat', re: '(cat)-(dog)', rep: '$2-$1', i: 1,"
+      "{t: 'cat-dog dog-cat cat-dog dog-cat', re: '(cat)-(dog)', rep: '\\\\2-\\\\1', i: 1,"
           + " res: 'dog-cat dog-cat cat-dog dog-cat'}")
   @TestParameters(
-      "{t: 'cat-dog dog-cat cat-dog dog-cat', re: '(cat)-(dog)', rep: '$2-$1', i: 2, res: 'dog-cat"
-          + " dog-cat dog-cat dog-cat'}")
+      "{t: 'cat-dog dog-cat cat-dog dog-cat', re: '(cat)-(dog)', rep: '\\\\2-\\\\1', i: 2, res:"
+          + " 'dog-cat dog-cat dog-cat dog-cat'}")
   @TestParameters("{t: 'a.b.c', re: '\\\\.', rep: '-', i: 1, res: 'a-b.c'}")
   @TestParameters("{t: 'a.b.c', re: '\\\\.', rep: '-', i: -1, res: 'a-b-c'}")
   public void replaceCount_success(String t, String re, String rep, long i, String res)
       throws Exception {
     String expr = String.format("regex.replace('%s', '%s', '%s', %d)", t, re, rep, i);
-    System.out.println("expr: " + expr);
     CelRuntime.Program program = RUNTIME.createProgram(COMPILER.compile(expr).getAst());
 
     Object result = program.eval();
@@ -108,7 +115,7 @@ public void replaceCount_success(String t, String re, String rep, long i, String
   @Test
   @TestParameters("{target: 'foo bar', regex: '(', replaceStr: '$2 $1'}")
   @TestParameters("{target: 'foo bar', regex: '[a-z', replaceStr: '$2 $1'}")
-  public void replace_invalid_regex(String target, String regex, String replaceStr)
+  public void replace_invalidRegex_throwsException(String target, String regex, String replaceStr)
       throws Exception {
     String expr = String.format("regex.replace('%s', '%s', '%s')", target, regex, replaceStr);
     CelAbstractSyntaxTree ast = COMPILER.compile(expr).getAst();
@@ -121,32 +128,48 @@ public void replace_invalid_regex(String target, String regex, String replaceStr
   }
 
   @Test
-  @TestParameters("{target: 'test', regex: '(.)', replaceStr: '$2'}")
-  public void replace_invalid_captureGroup(String target, String regex, String replaceStr)
-      throws Exception {
-    String expr = String.format("regex.replace('%s', '%s', '%s')", target, regex, replaceStr);
+  public void replace_invalidCaptureGroupReplaceStr_throwsException() throws Exception {
+    String expr = "regex.replace('test', '(.)', '\\\\2')";
     CelAbstractSyntaxTree ast = COMPILER.compile(expr).getAst();
 
     CelEvaluationException e =
         assertThrows(CelEvaluationException.class, () -> RUNTIME.createProgram(ast).eval());
 
-    assertThat(e).hasCauseThat().isInstanceOf(IndexOutOfBoundsException.class);
-    assertThat(e).hasCauseThat().hasMessageThat().contains("n > number of groups");
+    assertThat(e).hasCauseThat().isInstanceOf(IllegalArgumentException.class);
+    assertThat(e)
+        .hasCauseThat()
+        .hasMessageThat()
+        .contains("Replacement string references group 2 but regex has only 1 group(s)");
   }
 
   @Test
-  @TestParameters(
-      "{target: 'id=123', regex: 'id=(?P<value>\\\\d+)', replaceStr: 'value: ${values}'}")
-  public void replace_invalid_replaceStr(String target, String regex, String replaceStr)
-      throws Exception {
-    String expr = String.format("regex.replace('%s', '%s', '%s')", target, regex, replaceStr);
+  public void replace_trailingBackslashReplaceStr_throwsException() throws Exception {
+    String expr = "regex.replace('id=123', 'id=(?P<value>\\\\d+)', '\\\\')";
+    CelAbstractSyntaxTree ast = COMPILER.compile(expr).getAst();
+
+    CelEvaluationException e =
+        assertThrows(CelEvaluationException.class, () -> RUNTIME.createProgram(ast).eval());
+
+    assertThat(e).hasCauseThat().isInstanceOf(IllegalArgumentException.class);
+    assertThat(e)
+        .hasCauseThat()
+        .hasMessageThat()
+        .contains("Invalid replacement string: \\ not allowed at end");
+  }
+
+  @Test
+  public void replace_invalidGroupReferenceReplaceStr_throwsException() throws Exception {
+    String expr = "regex.replace('id=123', 'id=(?P<value>\\\\d+)', '\\\\a')";
     CelAbstractSyntaxTree ast = COMPILER.compile(expr).getAst();
 
     CelEvaluationException e =
         assertThrows(CelEvaluationException.class, () -> RUNTIME.createProgram(ast).eval());
 
     assertThat(e).hasCauseThat().isInstanceOf(IllegalArgumentException.class);
-    assertThat(e).hasCauseThat().hasMessageThat().contains("group 'values' not found");
+    assertThat(e)
+        .hasCauseThat()
+        .hasMessageThat()
+        .contains("Invalid replacement string: \\ must be followed by a digit");
   }
 
   @Test
@@ -208,6 +231,9 @@ private enum ExtractAllTestCase {
     NO_MATCH("regex.extractAll('id:123, id:456', 'assa')", ImmutableList.of()),
     NO_CAPTURE_GROUP(
         "regex.extractAll('id:123, id:456', 'id:\\\\d+')", ImmutableList.of("id:123", "id:456")),
+    CAPTURE_GROUP(
+        "regex.extractAll('key=\"\", key=\"val\"', 'key=\"([^\"]*)\"')",
+        ImmutableList.of("", "val")),
     SINGLE_NAMED_GROUP(
         "regex.extractAll('testuser@testdomain', '(?P<username>.*)@')",
         ImmutableList.of("testuser")),
