Add capability to visit custom tags

PiperOrigin-RevId: 649184086

Author: l46kok

policy/BUILD.bazel

@@ -8,6 +8,11 @@ java_library(
     exports = ["//policy/src/main/java/dev/cel/policy"],
 )
 
+java_library(
+    name = "value_string",
+    exports = ["//policy/src/main/java/dev/cel/policy:value_string"],
+)
+
 java_library(
     name = "source",
     exports = ["//policy/src/main/java/dev/cel/policy:source"],
@@ -33,6 +38,11 @@ java_library(
     exports = ["//policy/src/main/java/dev/cel/policy:parser"],
 )
 
+java_library(
+    name = "parser_context",
+    exports = ["//policy/src/main/java/dev/cel/policy:parser_context"],
+)
+
 java_library(
     name = "parser_factory",
     exports = ["//policy/src/main/java/dev/cel/policy:parser_factory"],

policy/src/main/java/dev/cel/policy/BUILD.bazel

@@ -109,6 +109,7 @@ java_library(
         ":source",
         ":validation_exception",
         ":value_string",
+        "//common:source",
         "//common/internal",
         "@maven//:com_google_guava_guava",
         "@maven//:org_yaml_snakeyaml",
@@ -227,7 +228,6 @@ java_library(
     srcs = [
         "ValueString.java",
     ],
-    visibility = ["//visibility:private"],
     deps = [
         "//:auto_value",
     ],
@@ -262,8 +262,9 @@ java_library(
     srcs = [
         "ParserContext.java",
     ],
-    visibility = ["//visibility:private"],
     deps = [
+        ":policy",
+        ":value_string",
         "//common:source",
     ],
 )

policy/src/main/java/dev/cel/policy/CelPolicy.java

@@ -19,9 +19,11 @@
 import com.google.auto.value.AutoOneOf;
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.util.Arrays;
+import java.util.Map;
 import java.util.Optional;
 
 /**
@@ -37,11 +39,14 @@ public abstract class CelPolicy {
 
   public abstract CelPolicySource policySource();
 
+  public abstract ImmutableMap<String, Object> metadata();
+
   /** Creates a new builder to construct a {@link CelPolicy} instance. */
   public static Builder newBuilder() {
     return new AutoValue_CelPolicy.Builder()
         .setName(ValueString.of(0, ""))
-        .setRule(Rule.newBuilder().build());
+        .setRule(Rule.newBuilder().build())
+        .setMetadata(ImmutableMap.of());
   }
 
   /** Builder for {@link CelPolicy}. */
@@ -55,6 +60,23 @@ public abstract static class Builder {
 
     public abstract Builder setPolicySource(CelPolicySource policySource);
 
+    // This should stay package-private to encourage add/set methods to be used instead.
+    abstract ImmutableMap.Builder<String, Object> metadataBuilder();
+
+    public abstract Builder setMetadata(ImmutableMap<String, Object> value);
+
+    @CanIgnoreReturnValue
+    public Builder putMetadata(String key, Object value) {
+      metadataBuilder().put(key, value);
+      return this;
+    }
+
+    @CanIgnoreReturnValue
+    public Builder putMetadata(Map<String, Object> map) {
+      metadataBuilder().putAll(map);
+      return this;
+    }
+
     public abstract CelPolicy build();
   }
 

policy/src/main/java/dev/cel/policy/CelPolicyParser.java

@@ -14,6 +14,8 @@
 
 package dev.cel.policy;
 
+import dev.cel.policy.ParserContext.PolicyParserContext;
+
 /** CelPolicyParser is the interface for parsing policies into a canonical Policy representation. */
 public interface CelPolicyParser {
 
@@ -41,47 +43,54 @@ interface TagVisitor<T> {
      * allow for continued parsing within a custom tag.
      */
     default void visitPolicyTag(
-        ParserContext<T> ctx, long id, String fieldName, T node, CelPolicy.Builder policyBuilder) {
-      ctx.reportError(id, String.format("Unsupported policy tag: %s", fieldName));
+        PolicyParserContext<T> ctx,
+        long id,
+        String tagName,
+        T node,
+        CelPolicy.Builder policyBuilder) {
+      ctx.reportError(id, String.format("Unsupported policy tag: %s", tagName));
     }
 
     /**
      * visitRuleTag accepts a parser context, field id, tag name, yaml node, as well as the parent
      * policy and current rule to allow for continued parsing within custom tags.
      */
     default void visitRuleTag(
-        ParserContext<T> ctx,
+        PolicyParserContext<T> ctx,
         long id,
-        String fieldName,
+        String tagName,
         T node,
+        CelPolicy.Builder policyBuilder,
         CelPolicy.Rule.Builder ruleBuilder) {
-      ctx.reportError(id, String.format("Unsupported rule tag: %s", fieldName));
+      ctx.reportError(id, String.format("Unsupported rule tag: %s", tagName));
     }
 
     /**
      * visitMatchTag accepts a parser context, field id, tag name, yaml node, as well as the parent
      * policy and current match to allow for continued parsing within custom tags.
      */
     default void visitMatchTag(
-        ParserContext<T> ctx,
+        PolicyParserContext<T> ctx,
         long id,
-        String fieldName,
+        String tagName,
         T node,
+        CelPolicy.Builder policyBuilder,
         CelPolicy.Match.Builder matchBuilder) {
-      ctx.reportError(id, String.format("Unsupported match tag: %s", fieldName));
+      ctx.reportError(id, String.format("Unsupported match tag: %s", tagName));
     }
 
     /**
      * visitVariableTag accepts a parser context, field id, tag name, yaml node, as well as the
      * parent policy and current variable to allow for continued parsing within custom tags.
      */
     default void visitVariableTag(
-        ParserContext<T> ctx,
+        PolicyParserContext<T> ctx,
         long id,
-        String fieldName,
+        String tagName,
         T node,
+        CelPolicy.Builder policyBuilder,
         CelPolicy.Variable.Builder variableBuilder) {
-      ctx.reportError(id, String.format("Unsupported variable tag: %s", fieldName));
+      ctx.reportError(id, String.format("Unsupported variable tag: %s", tagName));
     }
   }
 }