Add support for context types env.yaml

Add support for declaring a protobuf context message
whose top level fields are declared as variables in the CEL
environment.

PiperOrigin-RevId: 930613750

Author: jnthntatum

env/BUILD

@@ -130,6 +130,7 @@ cc_library(
         "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

env/config.h

@@ -32,6 +32,11 @@ class Config {
   void SetName(std::string name) { name_ = std::move(name); }
   std::string GetName() const { return name_; }
 
+  void SetContextType(std::string context_type) {
+    context_type_ = std::move(context_type);
+  }
+  std::string GetContextType() const { return context_type_; }
+
   struct ContainerConfig {
     std::string name;
     std::vector<std::string> abbreviations;
@@ -150,6 +155,7 @@ class Config {
 
  private:
   std::string name_;
+  std::string context_type_;
   ContainerConfig container_config_;
   std::vector<ExtensionConfig> extension_configs_;
   StandardLibraryConfig standard_library_config_;

env/env.cc

@@ -138,6 +138,11 @@ absl::StatusOr<std::unique_ptr<CompilerBuilder>> Env::NewCompilerBuilder() {
   for (const auto& abbr : config_.GetContainerConfig().abbreviations) {
     CEL_RETURN_IF_ERROR(container.AddAbbreviation(abbr));
   }
+
+  if (!config_.GetContextType().empty()) {
+    CEL_RETURN_IF_ERROR(
+        checker_builder.AddContextDeclaration(config_.GetContextType()));
+  }
   for (const auto& alias : config_.GetContainerConfig().aliases) {
     CEL_RETURN_IF_ERROR(container.AddAlias(alias.alias, alias.qualified_name));
   }

env/env_test.cc

@@ -344,6 +344,25 @@ TEST(ContainerConfigTest, ContainerConfigWithAliases) {
   EXPECT_THAT(result.GetIssues(), IsEmpty()) << result.FormatError();
 }
 
+TEST(ContextVariableConfigTest, Basic) {
+  Env env;
+  env.SetDescriptorPool(internal::GetSharedTestingDescriptorPool());
+  Config config;
+  config.SetContextType("cel.expr.conformance.proto3.TestAllTypes");
+  env.SetConfig(config);
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<Compiler> compiler, env.NewCompiler());
+
+  // Top-level fields of TestAllTypes like "single_int32" should resolve
+  // successfully.
+  ASSERT_OK_AND_ASSIGN(auto result, compiler->Compile("single_int32 > 10"));
+  EXPECT_THAT(result.GetIssues(), IsEmpty());
+
+  ASSERT_OK_AND_ASSIGN(auto result_invalid,
+                       compiler->Compile("non_existent_field > 10"));
+  EXPECT_THAT(result_invalid.GetIssues(), Not(IsEmpty()));
+}
+
 struct VariableConfigWithValueTestCase {
   Config::VariableConfig variable_config;
   std::string validate_type_expr;

env/env_yaml.cc

@@ -26,6 +26,7 @@
 #include "absl/base/no_destructor.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/log/absl_check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/escaping.h"
@@ -1245,6 +1246,34 @@ void EmitFunctionConfigs(const Config& env_config, YAML::Emitter& out,
   }
   out << YAML::EndSeq;
 }
+
+absl::Status ParseContextVariableConfig(Config& config, absl::string_view yaml,
+                                        const YAML::Node& root) {
+  const YAML::Node context_variable = root["context_variable"];
+  if (!context_variable.IsDefined()) {
+    return absl::OkStatus();
+  }
+  if (!context_variable.IsMap()) {
+    return YamlError(yaml, context_variable,
+                     "Node 'context_variable' is not a map");
+  }
+
+  const YAML::Node type_name = context_variable["type_name"];
+  const YAML::Node type = context_variable["type"];
+  const YAML::Node* type_node = nullptr;
+  if (type.IsDefined() && type.IsScalar()) {
+    type_node = &type;
+  } else if (type_name.IsDefined() && type_name.IsScalar()) {
+    type_node = &type_name;
+  } else {
+    return YamlError(yaml, context_variable,
+                     "Node 'context_variable' does not have a valid type");
+  }
+  ABSL_DCHECK(type_node != nullptr);
+  config.SetContextType(GetString(yaml, *type_node));
+  return absl::OkStatus();
+}
+
 }  // namespace
 
 absl::StatusOr<Config> EnvConfigFromYaml(const std::string& yaml) {
@@ -1263,6 +1292,7 @@ absl::StatusOr<Config> EnvConfigFromYaml(const std::string& yaml) {
   CEL_RETURN_IF_ERROR(ParseContainerConfig(config, yaml, root));
   CEL_RETURN_IF_ERROR(ParseExtensionConfigs(config, yaml, root));
   CEL_RETURN_IF_ERROR(ParseStandardLibraryConfig(config, yaml, root));
+  CEL_RETURN_IF_ERROR(ParseContextVariableConfig(config, yaml, root));
   CEL_RETURN_IF_ERROR(ParseVariableConfigs(config, yaml, root));
   CEL_RETURN_IF_ERROR(ParseFunctionConfigs(config, yaml, root));
   return config;

env/env_yaml_test.cc

@@ -216,6 +216,47 @@ TEST(EnvYamlTest, ParseVariableConfigWithTypeParams) {
   EXPECT_THAT(type_info.params[1].params, IsEmpty());
 }
 
+TEST(EnvYamlTest, ParseContextVariableConfig) {
+  ASSERT_OK_AND_ASSIGN(Config config, EnvConfigFromYaml(R"yaml(
+                    context_variable:
+                      type_name: "cel.expr.conformance.proto3.TestAllTypes"
+                  )yaml"));
+
+  EXPECT_EQ(config.GetContextType(),
+            "cel.expr.conformance.proto3.TestAllTypes");
+}
+
+TEST(EnvYamlTest, ParseContextVariableConfigAlternativeSyntax) {
+  ASSERT_OK_AND_ASSIGN(Config config, EnvConfigFromYaml(R"yaml(
+                    context_variable:
+                      type: "cel.expr.conformance.proto3.TestAllTypes"
+                  )yaml"));
+
+  EXPECT_EQ(config.GetContextType(),
+            "cel.expr.conformance.proto3.TestAllTypes");
+}
+
+TEST(EnvYamlTest, ParseContextVariableMalformedContextVariable) {
+  EXPECT_THAT(EnvConfigFromYaml(R"yaml(
+                    context_variable: 123
+
+                  )yaml"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Node 'context_variable' is not a map")));
+}
+
+TEST(EnvYamlTest, ParseContextVariableMalformedContextVariable2) {
+  EXPECT_THAT(
+      EnvConfigFromYaml(R"yaml(
+                    context_variable:
+                      type:
+                        foo: bar
+                  )yaml"),
+      StatusIs(
+          absl::StatusCode::kInvalidArgument,
+          HasSubstr("Node 'context_variable' does not have a valid type")));
+}
+
 TEST(EnvYamlTest, ParseVariableConfigWithTypeParamsLegacySyntax) {
   ASSERT_OK_AND_ASSIGN(Config config, EnvConfigFromYaml(R"yaml(
                     variables: