Add a shorthand for declaring policy variables

PiperOrigin-RevId: 868430821

Author: l46kok

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

@@ -29,6 +29,30 @@ public interface CelPolicyParserBuilder<T> {
   @CanIgnoreReturnValue
   CelPolicyParserBuilder<T> addTagVisitor(TagVisitor<T> tagVisitor);
 
+  /**
+   * Configures the parser to allow for key-value pairs to declare a variable name and expression.
+   *
+   * <p>For example:
+   *
+   * <pre>{@code
+   * variables:
+   * - foo: bar
+   * - baz: qux
+   * }</pre>
+   *
+   * <p>This is in contrast to the default behavior, which requires the following syntax:
+   *
+   * <pre>{@code
+   * variables:
+   * - name: foo
+   *   expression: bar
+   * - name: baz
+   *   expression: qux
+   * }</pre>
+   */
+  @CanIgnoreReturnValue
+  CelPolicyParserBuilder<T> enableSimpleVariables(boolean enable);
+
   /** Builds a new instance of {@link CelPolicyParser}. */
   @CheckReturnValue
   CelPolicyParser build();

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

@@ -49,6 +49,7 @@ final class CelPolicyYamlParser implements CelPolicyParser {
       Variable.newBuilder().setExpression(ERROR_VALUE).setName(ERROR_VALUE).build();
 
   private final TagVisitor<Node> tagVisitor;
+  private final boolean enableSimpleVariables;
 
   @Override
   public CelPolicy parse(String policySource) throws CelPolicyValidationException {
@@ -58,13 +59,15 @@ public CelPolicy parse(String policySource) throws CelPolicyValidationException
   @Override
   public CelPolicy parse(String policySource, String description)
       throws CelPolicyValidationException {
-    ParserImpl parser = new ParserImpl(tagVisitor, policySource, description);
+    ParserImpl parser =
+        new ParserImpl(tagVisitor, enableSimpleVariables, policySource, description);
     return parser.parseYaml();
   }
 
   private static class ParserImpl implements PolicyParserContext<Node> {
 
     private final TagVisitor<Node> tagVisitor;
+    private final boolean enableSimpleVariables;
     private final CelPolicySource policySource;
     private final ParserContext<Node> ctx;
 
@@ -336,9 +339,44 @@ public CelPolicy.Variable parseVariable(
       if (!assertYamlType(ctx, id, node, YamlNodeType.MAP)) {
         return ERROR_VARIABLE;
       }
+
       MappingNode variableMap = (MappingNode) node;
       Variable.Builder builder = Variable.newBuilder();
 
+      if (enableSimpleVariables) {
+        return parseVariableInline(ctx, id, variableMap, builder);
+      }
+      return parseVariableObject(ctx, policyBuilder, id, variableMap, builder);
+    }
+
+    private Variable parseVariableInline(
+        PolicyParserContext<Node> ctx, long id, MappingNode variableMap, Variable.Builder builder) {
+      int iterations = 0;
+      for (NodeTuple nodeTuple : variableMap.getValue()) {
+        Node keyNode = nodeTuple.getKeyNode();
+        long keyId = ctx.collectMetadata(keyNode);
+        builder.setName(ctx.newValueString(keyNode));
+        builder.setExpression(ctx.newValueString(nodeTuple.getValueNode()));
+        iterations++;
+
+        if (iterations > 1) {
+          ctx.reportError(keyId, "Only one variable may be defined inline");
+        }
+      }
+
+      if (!assertRequiredFields(ctx, id, builder.getMissingRequiredFieldNames())) {
+        return ERROR_VARIABLE;
+      }
+
+      return builder.build();
+    }
+
+    private Variable parseVariableObject(
+        PolicyParserContext<Node> ctx,
+        CelPolicy.Builder policyBuilder,
+        long id,
+        MappingNode variableMap,
+        Variable.Builder builder) {
       for (NodeTuple nodeTuple : variableMap.getValue()) {
         Node keyNode = nodeTuple.getKeyNode();
         long keyId = ctx.collectMetadata(keyNode);
@@ -370,8 +408,13 @@ public CelPolicy.Variable parseVariable(
       return builder.build();
     }
 
-    private ParserImpl(TagVisitor<Node> tagVisitor, String source, String description) {
+    private ParserImpl(
+        TagVisitor<Node> tagVisitor,
+        boolean enableSimpleVariables,
+        String source,
+        String description) {
       this.tagVisitor = tagVisitor;
+      this.enableSimpleVariables = enableSimpleVariables;
       this.policySource =
           CelPolicySource.newBuilder(CelCodePointArray.fromString(source))
               .setDescription(description)
@@ -413,9 +456,11 @@ public ValueString newValueString(Node node) {
   static final class Builder implements CelPolicyParserBuilder<Node> {
 
     private TagVisitor<Node> tagVisitor;
+    private boolean enableSimpleVariables;
 
     private Builder() {
       this.tagVisitor = new TagVisitor<Node>() {};
+      this.enableSimpleVariables = false;
     }
 
     @Override
@@ -424,17 +469,24 @@ public CelPolicyParserBuilder<Node> addTagVisitor(TagVisitor<Node> tagVisitor) {
       return this;
     }
 
+    @Override
+    public CelPolicyParserBuilder<Node> enableSimpleVariables(boolean enable) {
+      this.enableSimpleVariables = enable;
+      return this;
+    }
+
     @Override
     public CelPolicyParser build() {
-      return new CelPolicyYamlParser(tagVisitor);
+      return new CelPolicyYamlParser(tagVisitor, enableSimpleVariables);
     }
   }
 
-  static Builder newBuilder() {
-    return new Builder();
+  static CelPolicyParserBuilder<Node> newBuilder() {
+    return new Builder().enableSimpleVariables(false);
   }
 
-  private CelPolicyYamlParser(TagVisitor<Node> tagVisitor) {
+  private CelPolicyYamlParser(TagVisitor<Node> tagVisitor, boolean enableSimpleVariables) {
     this.tagVisitor = checkNotNull(tagVisitor);
+    this.enableSimpleVariables = enableSimpleVariables;
   }
 }

policy/src/test/java/dev/cel/policy/CelPolicyCompilerImplTest.java

@@ -302,6 +302,31 @@ public void evaluateYamlPolicy_lateBoundFunction() throws Exception {
     assertThat(evalResult).isEqualTo("foo" + exampleValue);
   }
 
+  @Test
+  @SuppressWarnings("unchecked") // Test only
+  public void evaluateYamlPolicy_withSimpleVariable() throws Exception {
+    Cel cel = newCel();
+    String policySource =
+        "name: shorthand_variables_policy\n"
+            + "rule:\n"
+            + "  variables:\n"
+            + "    - first: 'true'\n"
+            + "    - second: 'false'\n"
+            + "  match:\n"
+            + "    - condition: 'variables.first && variables.second'\n"
+            + "      output: 'false'\n";
+    CelPolicyParser parser =
+        CelPolicyParserFactory.newYamlParserBuilder().enableSimpleVariables(true).build();
+    CelPolicy policy = parser.parse(policySource);
+
+    CelAbstractSyntaxTree compiledPolicyAst =
+        CelPolicyCompilerFactory.newPolicyCompiler(cel).build().compile(policy);
+
+    Optional<Object> evalResult = (Optional<Object>) cel.createProgram(compiledPolicyAst).eval();
+
+    assertThat(evalResult).hasValue(false);
+  }
+
   private static final class EvaluablePolicyTestData {
     private final TestYamlPolicy yamlPolicy;
     private final PolicyTestCase testCase;

policy/src/test/java/dev/cel/policy/CelPolicyYamlParserTest.java

@@ -148,6 +148,30 @@ public void parseYamlPolicy_withImports() throws Exception {
         .inOrder();
   }
 
+  @Test
+  public void parseYamlPolicy_withSimpleVariable_multipleInlinedVariables() {
+    String policySource =
+        "name: shorthand_variables_policy\n"
+            + "rule:\n"
+            + "  variables:\n"
+            + "    - first: 'true'\n"
+            + "      second: 'false'\n"
+            + "  match:\n"
+            + "    - condition: 'variables.my_var'\n"
+            + "      output: 'true'\n";
+    CelPolicyParser parser =
+        CelPolicyParserFactory.newYamlParserBuilder().enableSimpleVariables(true).build();
+
+    CelPolicyValidationException e =
+        assertThrows(CelPolicyValidationException.class, () -> parser.parse(policySource));
+    assertThat(e)
+        .hasMessageThat()
+        .contains(
+            "ERROR: <input>:5:7: Only one variable may be defined inline\n"
+                + " |       second: 'false'\n"
+                + " | ......^");
+  }
+
   @Test
   public void parseYamlPolicy_errors(@TestParameter PolicyParseErrorTestCase testCase) {
     CelPolicyValidationException e =