From 452beeaf52ca291e77d71820ae2917163c460687 Mon Sep 17 00:00:00 2001 From: Jonathan Tatum Date: Wed, 26 Aug 2026 11:37:01 -0700 Subject: [PATCH] Migrate gaer::CelProtoWrapper::CreateMessage to just call the modern equivalent. - Add overload for specifying the expected message factory and descriptor pool. - Avoid cloning parsed messages already tagged as unsafe when converting to legacy. - Preserve legacy behavior for key errors on json maps. PiperOrigin-RevId: 971405326 --- common/legacy_value.cc | 59 ++++++---- common/legacy_value.h | 8 +- common/values/legacy_map_value.cc | 23 ++-- common/values/parsed_message_value.h | 6 ++ eval/eval/select_step.cc | 73 +++++++------ eval/internal/cel_value_equal_test.cc | 2 + eval/public/cel_options.cc | 1 + eval/public/cel_options.h | 10 ++ eval/public/structs/BUILD | 9 +- eval/public/structs/cel_proto_wrapper.cc | 77 +++++++++++-- eval/public/structs/cel_proto_wrapper.h | 16 ++- eval/public/structs/cel_proto_wrapper_test.cc | 19 ++-- .../proto_message_type_adapter_test.cc | 10 +- extensions/select_optimization.cc | 102 ++++++++++-------- runtime/runtime_options.h | 10 ++ 15 files changed, 296 insertions(+), 129 deletions(-) diff --git a/common/legacy_value.cc b/common/legacy_value.cc index b963e5071..34b0e000f 100644 --- a/common/legacy_value.cc +++ b/common/legacy_value.cc @@ -74,11 +74,9 @@ using ::cel::interop_internal::TrivialTypeInfo; using ::google::api::expr::runtime::CelList; using ::google::api::expr::runtime::CelMap; using ::google::api::expr::runtime::CelValue; -using ::google::api::expr::runtime::CreateCelValueFromField; using ::google::api::expr::runtime::GetGenericProtoTypeInfoInstance; using ::google::api::expr::runtime::LegacyTypeInfoApis; using ::google::api::expr::runtime::MessageWrapper; -using ::google::api::expr::runtime::internal::GetGenericProtoAccessApisInstance; using ::google::api::expr::runtime::internal::MaybeWrapValueToMessage; absl::Status InvalidMapKeyTypeError(ValueKind kind) { @@ -262,6 +260,11 @@ CelValue LegacyTrivialStructValue(google::protobuf::Arena* absl_nonnull arena, } if (auto parsed_message_value = value.AsParsedMessage(); parsed_message_value) { + if (interop_internal::IsUnsafeParsedMessageValue(*parsed_message_value)) { + return CelValue::CreateMessageWrapper( + AsMessageWrapper(cel::to_address(*parsed_message_value), + &GetGenericProtoTypeInfoInstance())); + } auto maybe_cloned = parsed_message_value->Clone(arena); return CelValue::CreateMessageWrapper(MessageWrapper( cel::to_address(maybe_cloned), &GetGenericProtoTypeInfoInstance())); @@ -923,17 +926,26 @@ absl::Status LegacyStructValue::GetFieldByName( const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool, google::protobuf::MessageFactory* absl_nonnull message_factory, google::protobuf::Arena* absl_nonnull arena, Value* absl_nonnull result) const { - auto message_wrapper = AsMessageWrapper(message_ptr_, legacy_type_info_); if (ABSL_PREDICT_FALSE(legacy_type_info_ == TrivialTypeInfo::GetInstance())) { *result = NoSuchFieldError(name); return absl::OkStatus(); } - CEL_ASSIGN_OR_RETURN(auto cel_value, - GetGenericProtoAccessApisInstance().GetField( - name, message_wrapper, unboxing_options, - MemoryManagerRef::Pooling(arena))); - CEL_RETURN_IF_ERROR(ModernValue(arena, cel_value, *result)); - return absl::OkStatus(); + + ParsedMessageValue parsed_message = UnsafeParsedMessageValue(message_ptr_); + const auto* descriptor = parsed_message.GetDescriptor(); + const auto* field = descriptor->FindFieldByName(name); + if (field == nullptr) { + field = descriptor->file()->pool()->FindExtensionByPrintableName(descriptor, + name); + if (field == nullptr) { + *result = NoSuchFieldError(name); + return absl::OkStatus(); + } + } + + return interop_internal::WrapLegacyMessageField( + message_ptr_, field, unboxing_options, descriptor_pool, message_factory, + arena, result); } absl::Status LegacyStructValue::GetFieldByNumber( @@ -985,7 +997,6 @@ absl::Status LegacyStructValue::Qualify( if (ABSL_PREDICT_FALSE(qualifiers.empty())) { return absl::InvalidArgumentError("invalid select qualifier path."); } - auto message_wrapper = AsMessageWrapper(message_ptr_, legacy_type_info_); if (ABSL_PREDICT_FALSE(legacy_type_info_ == TrivialTypeInfo::GetInstance())) { absl::string_view field_name = absl::visit( absl::Overload( @@ -1000,12 +1011,13 @@ absl::Status LegacyStructValue::Qualify( *count = -1; return absl::OkStatus(); } - CEL_ASSIGN_OR_RETURN(auto legacy_result, - GetGenericProtoAccessApisInstance().Qualify( - qualifiers, message_wrapper, presence_test, - MemoryManager::Pooling(arena))); - CEL_RETURN_IF_ERROR(ModernValue(arena, legacy_result.value, *result)); - *count = legacy_result.qualifier_count; + + ParsedMessageValue parsed_message = UnsafeParsedMessageValue(message_ptr_); + CEL_RETURN_IF_ERROR(parsed_message.Qualify(qualifiers, presence_test, + descriptor_pool, message_factory, + arena, result, count)); + + interop_internal::WrapLegacyFieldAccessResult(arena, result); return absl::OkStatus(); } @@ -1311,12 +1323,17 @@ const google::protobuf::Message* absl_nullable GetLegacyMessage(const Value& val absl::Status WrapLegacyMessageField( const google::protobuf::Message* absl_nonnull message, const google::protobuf::FieldDescriptor* absl_nonnull field_descriptor, - ProtoWrapperTypeOptions unboxing_option, google::protobuf::Arena* arena, + ProtoWrapperTypeOptions unboxing_option, + const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool, + google::protobuf::MessageFactory* absl_nonnull message_factory, google::protobuf::Arena* arena, Value* absl_nonnull out) { - CEL_ASSIGN_OR_RETURN(CelValue result, - CreateCelValueFromField(message, field_descriptor, - unboxing_option, arena)); - return ModernValue(arena, result, *out); + ParsedMessageValue parsed_message = UnsafeParsedMessageValue(message); + CEL_RETURN_IF_ERROR(parsed_message.GetField(field_descriptor, unboxing_option, + descriptor_pool, message_factory, + arena, out)); + WrapLegacyFieldAccessResult(arena, out); + + return absl::OkStatus(); } } // namespace interop_internal diff --git a/common/legacy_value.h b/common/legacy_value.h index 5b7140387..e71eb85a7 100644 --- a/common/legacy_value.h +++ b/common/legacy_value.h @@ -65,6 +65,10 @@ class MessageFactory; namespace cel::interop_internal { +inline bool IsUnsafeParsedMessageValue(const cel::ParsedMessageValue& value) { + return value.is_unsafe(); +} + // Returns the underlying `google::protobuf::Message` of a `cel::Value` if it is a legacy // message with the default type info, or `nullptr` otherwise. const google::protobuf::Message* absl_nullable GetLegacyMessage(const Value& value); @@ -82,7 +86,9 @@ void WrapLegacyFieldAccessResult(google::protobuf::Arena* absl_nonnull arena, absl::Status WrapLegacyMessageField( const google::protobuf::Message* absl_nonnull message, const google::protobuf::FieldDescriptor* absl_nonnull field_descriptor, - ProtoWrapperTypeOptions unboxing_option, google::protobuf::Arena* arena, + ProtoWrapperTypeOptions unboxing_option, + const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool, + google::protobuf::MessageFactory* absl_nonnull message_factory, google::protobuf::Arena* arena, Value* absl_nonnull out); absl::StatusOr FromLegacyValue( diff --git a/common/values/legacy_map_value.cc b/common/values/legacy_map_value.cc index e287df076..1c35ed327 100644 --- a/common/values/legacy_map_value.cc +++ b/common/values/legacy_map_value.cc @@ -401,21 +401,24 @@ class LegacyParsedJsonMapValue final if (arena == nullptr) { arena = arena_; } - if (auto status = - google::api::expr::runtime::CelValue::CheckMapKeyType(key); - !status.ok()) { - status.IgnoreError(); - return std::nullopt; - } Value modern_key; if (ABSL_PREDICT_FALSE(!ModernValue(arena, key, modern_key).ok())) { return std::nullopt; } Value modern_val; - auto status_or_found = value_.Find( - modern_key, google::protobuf::DescriptorPool::generated_pool(), - google::protobuf::MessageFactory::generated_factory(), arena, &modern_val); - if (!status_or_found.ok() || !*status_or_found) { + // Call custom map FindDirectly. MapValue normally handles coercing error + // results to value types, so emulate that here. + // + // We know that the descriptor pool and message factory aren't needed here, + // so fine to use generated. + auto found = + Find(modern_key, google::protobuf::DescriptorPool::generated_pool(), + google::protobuf::MessageFactory::generated_factory(), arena, &modern_val); + if (!found.ok()) { + return google::api::expr::runtime::CreateErrorValue(arena, + found.status()); + } + if (!(*found) && !modern_val.IsError()) { return std::nullopt; } return UnsafeLegacyValue(modern_val, /*stable=*/false, arena); diff --git a/common/values/parsed_message_value.h b/common/values/parsed_message_value.h index 3cad912e7..1d02a4369 100644 --- a/common/values/parsed_message_value.h +++ b/common/values/parsed_message_value.h @@ -48,6 +48,10 @@ namespace cel { +namespace interop_internal { +bool IsUnsafeParsedMessageValue(const ParsedMessageValue& value); +} + class MessageValue; class StructValue; class Value; @@ -189,6 +193,8 @@ class ParsedMessageValue final friend class common_internal::StructValueMixin; friend ParsedMessageValue UnsafeParsedMessageValue( const google::protobuf::Message* absl_nonnull value); + friend bool interop_internal::IsUnsafeParsedMessageValue( + const ParsedMessageValue& value); explicit ParsedMessageValue( const google::protobuf::Message* absl_nonnull value ABSL_ATTRIBUTE_LIFETIME_BOUND) diff --git a/eval/eval/select_step.cc b/eval/eval/select_step.cc index a57179017..e7974496b 100644 --- a/eval/eval/select_step.cc +++ b/eval/eval/select_step.cc @@ -86,17 +86,22 @@ absl::Status WrappedStructGet( ProtoWrapperTypeOptions unboxing_option, const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool, google::protobuf::MessageFactory* absl_nonnull message_factory, - google::protobuf::Arena* absl_nonnull arena, Value* absl_nonnull result) { - if (const google::protobuf::Message* message = - cel::interop_internal::GetLegacyMessage(target); - message != nullptr) { - CelValue::MessageWrapper message_wrapper( - message, &GetGenericProtoTypeInfoInstance()); - CEL_ASSIGN_OR_RETURN(CelValue cel_value, - internal::GetGenericProtoAccessApisInstance().GetField( - field, message_wrapper, unboxing_option, - cel::MemoryManagerRef::Pooling(arena))); - return cel::ModernValue(arena, cel_value, *result); + google::protobuf::Arena* absl_nonnull arena, + bool enable_use_new_field_select_implementation, + Value* absl_nonnull result) { + if (!enable_use_new_field_select_implementation) { + if (const google::protobuf::Message* message = + cel::interop_internal::GetLegacyMessage(target); + message != nullptr) { + CelValue::MessageWrapper message_wrapper( + message, &GetGenericProtoTypeInfoInstance()); + CEL_ASSIGN_OR_RETURN( + CelValue cel_value, + internal::GetGenericProtoAccessApisInstance().GetField( + field, message_wrapper, unboxing_option, + cel::MemoryManagerRef::Pooling(arena))); + return cel::ModernValue(arena, cel_value, *result); + } } return target.GetStruct().GetFieldByName( field, unboxing_option, descriptor_pool, message_factory, arena, result); @@ -132,7 +137,9 @@ absl::Status PerformGet(const Value& target, absl::string_view field, ProtoWrapperTypeOptions unboxing_option, const google::protobuf::DescriptorPool* descriptor_pool, google::protobuf::MessageFactory* message_factory, - google::protobuf::Arena* arena, Value& result) { + google::protobuf::Arena* arena, + bool enable_use_new_field_select_implementation, + Value& result) { switch (target.kind()) { case ValueKind::kMap: { auto status = target.GetMap().Get(field_value, descriptor_pool, @@ -143,9 +150,9 @@ absl::Status PerformGet(const Value& target, absl::string_view field, return absl::OkStatus(); } case ValueKind::kStruct: { - auto status = - WrappedStructGet(target, field, unboxing_option, descriptor_pool, - message_factory, arena, &result); + auto status = WrappedStructGet( + target, field, unboxing_option, descriptor_pool, message_factory, + arena, enable_use_new_field_select_implementation, &result); if (!status.ok()) { result = ErrorValue(std::move(status)); } @@ -161,7 +168,9 @@ absl::Status PerformOptionalGet(const Value& target, absl::string_view field, ProtoWrapperTypeOptions unboxing_option, const google::protobuf::DescriptorPool* descriptor_pool, google::protobuf::MessageFactory* message_factory, - google::protobuf::Arena* arena, Value& result) { + google::protobuf::Arena* arena, + bool enable_use_new_field_select_implementation, + Value& result) { switch (target.kind()) { case ValueKind::kMap: { CEL_ASSIGN_OR_RETURN( @@ -182,9 +191,9 @@ absl::Status PerformOptionalGet(const Value& target, absl::string_view field, result = OptionalValue::None(); return absl::OkStatus(); } - CEL_RETURN_IF_ERROR(WrappedStructGet(target, field, unboxing_option, - descriptor_pool, message_factory, - arena, &result)); + CEL_RETURN_IF_ERROR(WrappedStructGet( + target, field, unboxing_option, descriptor_pool, message_factory, + arena, enable_use_new_field_select_implementation, &result)); ABSL_DCHECK(!result.IsUnknown()); result = OptionalValue::Of(std::move(result), arena); @@ -247,7 +256,7 @@ absl::Status SelectStep::Evaluate(ExecutionFrame* frame) const { optional_arg = arg.GetOptional(); } - if (!(optional_arg || arg->Is() || arg->Is())) { + if (!(optional_arg || arg.IsMap() || arg.IsStruct())) { frame->value_stack().PopAndPush(cel::ErrorValue(InvalidSelectTargetError()), std::move(result_trail)); return absl::OkStatus(); @@ -290,7 +299,8 @@ absl::Status SelectStep::Evaluate(ExecutionFrame* frame) const { optional_arg->Value(&value); auto status = PerformOptionalGet( value, field_, field_value_, unboxing_option_, frame->descriptor_pool(), - frame->message_factory(), frame->arena(), result); + frame->message_factory(), frame->arena(), + frame->options().enable_use_new_field_select_implementation, result); if (!status.ok()) { result = ErrorValue(std::move(status)); } @@ -300,7 +310,8 @@ absl::Status SelectStep::Evaluate(ExecutionFrame* frame) const { CEL_RETURN_IF_ERROR(PerformGet( arg, field_, field_value_, unboxing_option_, frame->descriptor_pool(), - frame->message_factory(), frame->arena(), result)); + frame->message_factory(), frame->arena(), + frame->options().enable_use_new_field_select_implementation, result)); frame->value_stack().PopAndPush(std::move(result), std::move(result_trail)); return absl::OkStatus(); } @@ -380,19 +391,20 @@ class DirectSelectStep : public DirectExpressionStep { } Value value; optional_arg->Value(&value); - auto status = - PerformOptionalGet(value, field_, field_value_, unboxing_option_, - frame.descriptor_pool(), frame.message_factory(), - frame.arena(), result); + auto status = PerformOptionalGet( + value, field_, field_value_, unboxing_option_, + frame.descriptor_pool(), frame.message_factory(), frame.arena(), + frame.options().enable_use_new_field_select_implementation, result); if (!status.ok()) { result = ErrorValue(std::move(status)); } return absl::OkStatus(); } - return PerformGet(result, field_, field_value_, unboxing_option_, - frame.descriptor_pool(), frame.message_factory(), - frame.arena(), result); + return PerformGet( + result, field_, field_value_, unboxing_option_, frame.descriptor_pool(), + frame.message_factory(), frame.arena(), + frame.options().enable_use_new_field_select_implementation, result); } private: @@ -495,7 +507,8 @@ absl::Status ProtoSelectStep::EvaluateLegacyMessageGetField( return absl::OkStatus(); } return cel::interop_internal::WrapLegacyMessageField( - legacy_message, field_descriptor_, unboxing_option_, frame->arena(), + legacy_message, field_descriptor_, unboxing_option_, + frame->descriptor_pool(), frame->message_factory(), frame->arena(), &frame->value_stack().Peek()); } diff --git a/eval/internal/cel_value_equal_test.cc b/eval/internal/cel_value_equal_test.cc index 109a63795..efef9ec93 100644 --- a/eval/internal/cel_value_equal_test.cc +++ b/eval/internal/cel_value_equal_test.cc @@ -131,6 +131,8 @@ const std::vector& ValueExamples1() { result->push_back(CelValue::CreateMap(&CelMapExample1())); result->push_back(CelValue::CreateCelTypeView("type")); + ABSL_CHECK_EQ(arena.SpaceUsed(), 0) << "Arena should not be used."; + return result.release(); }(); return *examples; diff --git a/eval/public/cel_options.cc b/eval/public/cel_options.cc index 93b67ad35..100ef5e01 100644 --- a/eval/public/cel_options.cc +++ b/eval/public/cel_options.cc @@ -45,6 +45,7 @@ cel::RuntimeOptions ConvertToRuntimeOptions(const InterpreterOptions& options) { options.enable_fast_builtins, options.enable_precision_preserving_double_format, options.enable_typed_field_access, + options.enable_use_new_field_select_implementation, }; } diff --git a/eval/public/cel_options.h b/eval/public/cel_options.h index 001990431..b0d9e6db0 100644 --- a/eval/public/cel_options.h +++ b/eval/public/cel_options.h @@ -223,6 +223,16 @@ struct InterpreterOptions { // path for field access when the type is known at plan time, instead of using // the generic field access implementation. bool enable_typed_field_access = false; + + // Temporary flag to gate using a new field selection implementation for + // protos. + // + // For the cel::Runtime APIs, this is a no-op. + // + // For google::api::expr::runtime::CelExpression, this will enable updated + // implementations for field access on protobuf messages, aligned with the + // cel::Value implementation. + bool enable_use_new_field_select_implementation = false; }; // LINT.ThenChange(//depot/google3/runtime/runtime_options.h) diff --git a/eval/public/structs/BUILD b/eval/public/structs/BUILD index 468867294..41c0d7047 100644 --- a/eval/public/structs/BUILD +++ b/eval/public/structs/BUILD @@ -30,10 +30,17 @@ cc_library( deps = [ ":cel_proto_wrap_util", ":proto_message_type_adapter", + ":trivial_legacy_type_info_internal", + "//common:value", "//eval/public:cel_value", "//eval/public:message_wrapper", "//internal:proto_time_encoding", - "@com_google_absl//absl/types:optional", + "@com_google_absl//absl/base:no_destructor", + "@com_google_absl//absl/base:nullability", + "@com_google_absl//absl/log:absl_check", + "@com_google_absl//absl/log:absl_log", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", "@com_google_protobuf//:duration_cc_proto", "@com_google_protobuf//:protobuf", "@com_google_protobuf//:timestamp_cc_proto", diff --git a/eval/public/structs/cel_proto_wrapper.cc b/eval/public/structs/cel_proto_wrapper.cc index 6fad6aee3..2d2928c22 100644 --- a/eval/public/structs/cel_proto_wrapper.cc +++ b/eval/public/structs/cel_proto_wrapper.cc @@ -14,11 +14,21 @@ #include "eval/public/structs/cel_proto_wrapper.h" -#include "absl/types/optional.h" +#include + +#include "absl/base/no_destructor.h" +#include "absl/base/nullability.h" +#include "absl/log/absl_check.h" +#include "absl/log/absl_log.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "common/legacy_value.h" +#include "common/value.h" #include "eval/public/cel_value.h" #include "eval/public/message_wrapper.h" #include "eval/public/structs/cel_proto_wrap_util.h" #include "eval/public/structs/proto_message_type_adapter.h" +#include "eval/public/structs/trivial_legacy_type_info_internal.h" #include "google/protobuf/arena.h" #include "google/protobuf/descriptor.h" #include "google/protobuf/message.h" @@ -27,9 +37,28 @@ namespace google::api::expr::runtime { namespace { +using ::cel::interop_internal::TrivialTypeInfo; using ::google::protobuf::Arena; using ::google::protobuf::Descriptor; +using ::google::protobuf::DescriptorPool; using ::google::protobuf::Message; +using ::google::protobuf::MessageFactory; + +// Returns the arena for the given message, or the fallback arena if the +// message does not have an arena. +google::protobuf::Arena* GetArena(const Message* absl_nonnull message, + google::protobuf::Arena* absl_nullable arena) { + if (arena != nullptr) { + return arena; + } + if (message->GetArena() != nullptr) { + return message->GetArena(); + } + static absl::NoDestructor fallback_arena; + ABSL_LOG(WARNING) << "CEL: using fallback global arena for message: " + << message->GetTypeName(); + return fallback_arena.get(); +} } // namespace @@ -38,14 +67,50 @@ CelValue CelProtoWrapper::InternalWrapMessage(const Message* message) { MessageWrapper(message, &GetGenericProtoTypeInfoInstance())); } -// CreateMessage creates CelValue from google::protobuf::Message. -// As some of CEL basic types are subclassing google::protobuf::Message, -// this method contains type checking and downcasts. +CelValue CelProtoWrapper::CreateMessage(const Message* value, + const google::protobuf::DescriptorPool* pool, + MessageFactory* factory, Arena* arena) { + ABSL_DCHECK(value != nullptr); + if (value->GetDescriptor() == nullptr || value->GetReflection() == nullptr) { + // This only happens for custom messages subclasses that CEL can't support. + return CelValue::CreateMessageWrapper( + MessageWrapper(value, TrivialTypeInfo::GetInstance())); + } + + // TODO(jdtatum): special types (mainly the json wrappers) were implicitly + // copied before. This would allow some cases where the output could remain + // valid after the input is gone (though that was unsafe in the general case). + // Check TGP to see if there was anyone taking advantage of this. + auto modern_value = + cel::Value::WrapMessageUnsafe(value, pool, factory, arena); + + absl::StatusOr cel_value = cel::LegacyValue(arena, modern_value); + if (!cel_value.ok()) { + // Should not happen for a valid google::protobuf::Message. + auto* status = + google::protobuf::Arena::Create(arena, cel_value.status()); + return CelValue::CreateError(status); + } + return *cel_value; +} + CelValue CelProtoWrapper::CreateMessage(const Message* value, Arena* arena) { - return internal::UnwrapMessageToValue(value, &InternalWrapMessage, arena); + if (value == nullptr) { + return CelValue::CreateNull(); + } + + if (value->GetDescriptor() == nullptr || value->GetReflection() == nullptr) { + // This only happens for custom messages subclasses that CEL can't support. + return CelValue::CreateMessageWrapper( + MessageWrapper(value, TrivialTypeInfo::GetInstance())); + } + const auto* pool = value->GetDescriptor()->file()->pool(); + auto* factory = value->GetReflection()->GetMessageFactory(); + arena = GetArena(value, arena); + return CreateMessage(value, pool, factory, arena); } -absl::optional CelProtoWrapper::MaybeWrapValue( +std::optional CelProtoWrapper::MaybeWrapValue( const Descriptor* descriptor, google::protobuf::MessageFactory* factory, const CelValue& value, Arena* arena) { const Message* msg = diff --git a/eval/public/structs/cel_proto_wrapper.h b/eval/public/structs/cel_proto_wrapper.h index 73942c253..861d3b9c8 100644 --- a/eval/public/structs/cel_proto_wrapper.h +++ b/eval/public/structs/cel_proto_wrapper.h @@ -1,9 +1,11 @@ #ifndef THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_STRUCTS_CEL_PROTO_WRAPPER_H_ #define THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_STRUCTS_CEL_PROTO_WRAPPER_H_ +#include + #include "google/protobuf/duration.pb.h" #include "google/protobuf/timestamp.pb.h" -#include "absl/types/optional.h" +#include "absl/base/nullability.h" #include "eval/public/cel_value.h" #include "internal/proto_time_encoding.h" #include "google/protobuf/arena.h" @@ -17,6 +19,16 @@ class CelProtoWrapper { // CreateMessage creates CelValue from google::protobuf::Message. // As some of CEL basic types are subclassing google::protobuf::Message, // this method contains type checking and downcasts. + static CelValue CreateMessage(const google::protobuf::Message* absl_nonnull value, + const google::protobuf::DescriptorPool* absl_nonnull pool, + google::protobuf::MessageFactory* absl_nonnull factory, + google::protobuf::Arena* absl_nonnull arena); + + // Prefer using an explicit descriptor pool and message factory instead of + // the ones associated with the value. + // + // For backward compatibility, nullptr message is allowed and will result in a + // the CEL null_type value. static CelValue CreateMessage(const google::protobuf::Message* value, google::protobuf::Arena* arena); @@ -43,7 +55,7 @@ class CelProtoWrapper { // message to native CelValue representation during a protobuf field read. // Just as CreateMessage should only be used when reading protobuf values, // MaybeWrapValue should only be used when assigning protobuf fields. - static absl::optional MaybeWrapValue( + static std::optional MaybeWrapValue( const google::protobuf::Descriptor* descriptor, google::protobuf::MessageFactory* factory, const CelValue& value, google::protobuf::Arena* arena); }; diff --git a/eval/public/structs/cel_proto_wrapper_test.cc b/eval/public/structs/cel_proto_wrapper_test.cc index 408e33284..6d47406d8 100644 --- a/eval/public/structs/cel_proto_wrapper_test.cc +++ b/eval/public/structs/cel_proto_wrapper_test.cc @@ -104,7 +104,7 @@ class CelProtoWrapperTest : public ::testing::Test { T dyn_value; CelValue cel_dyn_value = - CelProtoWrapper::CreateMessage(ReflectedCopy(message).get(), arena()); + CelProtoWrapper::CreateMessage(ReflectedCopy(message), arena()); EXPECT_THAT(cel_dyn_value.type(), Eq(cel_value.type())); EXPECT_TRUE(cel_dyn_value.GetValue(&dyn_value)); EXPECT_THAT(value, Eq(dyn_value)); @@ -121,10 +121,9 @@ class CelProtoWrapperTest : public ::testing::Test { EXPECT_THAT(cel_value.MessageOrDie(), testutil::EqualsProto(*result)); } - std::unique_ptr ReflectedCopy( - const google::protobuf::Message& message) { - std::unique_ptr dynamic_value( - factory_.GetPrototype(message.GetDescriptor())->New()); + google::protobuf::Message* ReflectedCopy(const google::protobuf::Message& message) { + google::protobuf::Message* dynamic_value = + factory_.GetPrototype(message.GetDescriptor())->New(&arena_); dynamic_value->CopyFrom(message); return dynamic_value; } @@ -213,7 +212,7 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicValueNull) { value_msg.set_null_value(protobuf::NULL_VALUE); CelValue value = - CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg).get(), arena()); + CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg), arena()); EXPECT_TRUE(value.IsNull()); } @@ -314,8 +313,8 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicStruct) { const std::string kFieldBool = "field_bool"; (*struct_msg.mutable_fields())[kFieldInt].set_number_value(1.); (*struct_msg.mutable_fields())[kFieldBool].set_bool_value(true); - CelValue value = - CelProtoWrapper::CreateMessage(ReflectedCopy(struct_msg).get(), arena()); + auto reflected_copy = ReflectedCopy(struct_msg); + CelValue value = CelProtoWrapper::CreateMessage(reflected_copy, arena()); EXPECT_TRUE(value.IsMap()); const CelMap* cel_map = value.MapOrDie(); ASSERT_TRUE(cel_map != nullptr); @@ -355,7 +354,7 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicValueStruct) { .set_number_value(2); CelValue value = - CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg).get(), arena()); + CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg), arena()); EXPECT_TRUE(value.IsMap()); EXPECT_TRUE( (*value.MapOrDie())[CelValue::CreateString(&kField1)].has_value()); @@ -398,7 +397,7 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicValueListValue) { value_msg.mutable_list_value()->add_values()->set_number_value(2.); CelValue value = - CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg).get(), arena()); + CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg), arena()); EXPECT_TRUE(value.IsList()); EXPECT_THAT((*value.ListOrDie())[0].DoubleOrDie(), testing::DoubleEq(1)); EXPECT_THAT((*value.ListOrDie())[1].DoubleOrDie(), testing::DoubleEq(2)); diff --git a/eval/public/structs/proto_message_type_adapter_test.cc b/eval/public/structs/proto_message_type_adapter_test.cc index 529052025..b44c17062 100644 --- a/eval/public/structs/proto_message_type_adapter_test.cc +++ b/eval/public/structs/proto_message_type_adapter_test.cc @@ -1188,10 +1188,12 @@ TEST(ProtoMesssageTypeAdapter, InteropFieldAccess) { message.GetDescriptor()->FindFieldByName("string_value"); ASSERT_NE(field, nullptr); cel::Value field_value; - ASSERT_THAT(cel::interop_internal::WrapLegacyMessageField( - &message, field, ProtoWrapperTypeOptions::kUnsetNull, &arena, - &field_value), - IsOk()); + ASSERT_THAT( + cel::interop_internal::WrapLegacyMessageField( + &message, field, ProtoWrapperTypeOptions::kUnsetNull, + google::protobuf::DescriptorPool::generated_pool(), + google::protobuf::MessageFactory::generated_factory(), &arena, &field_value), + IsOk()); EXPECT_THAT(field_value, cel::test::StringValueIs("hello")); } diff --git a/extensions/select_optimization.cc b/extensions/select_optimization.cc index 83ea6abc6..4dcd7d594 100644 --- a/extensions/select_optimization.cc +++ b/extensions/select_optimization.cc @@ -276,26 +276,29 @@ absl::StatusOr MapKeyFromQualifier(const AttributeQualifier& qual, } } -// Helper for StructValue::GetFieldByName. Used for opting out of old reflection -// implementation. +// // Helper for StructValue::GetFieldByName. Used for opting out of old +// reflection implementation. absl::StatusOr WrappedStructGet( const Value& target, absl::string_view field, const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool, google::protobuf::MessageFactory* absl_nonnull message_factory, - google::protobuf::Arena* absl_nonnull arena) { - if (const google::protobuf::Message* message = - cel::interop_internal::GetLegacyMessage(target); - message != nullptr) { - CelValue::MessageWrapper message_wrapper( - message, &GetGenericProtoTypeInfoInstance()); - CEL_ASSIGN_OR_RETURN( - CelValue cel_value, - GetGenericProtoAccessApisInstance().GetField( - field, message_wrapper, ProtoWrapperTypeOptions::kUnsetProtoDefault, - MemoryManagerRef::Pooling(arena))); - Value result; - CEL_RETURN_IF_ERROR(cel::ModernValue(arena, cel_value, result)); - return result; + google::protobuf::Arena* absl_nonnull arena, + bool enable_use_new_field_select_implementation) { + if (!enable_use_new_field_select_implementation) { + if (const google::protobuf::Message* message = + cel::interop_internal::GetLegacyMessage(target); + message != nullptr) { + CelValue::MessageWrapper message_wrapper( + message, &GetGenericProtoTypeInfoInstance()); + CEL_ASSIGN_OR_RETURN(CelValue cel_value, + GetGenericProtoAccessApisInstance().GetField( + field, message_wrapper, + ProtoWrapperTypeOptions::kUnsetProtoDefault, + MemoryManagerRef::Pooling(arena))); + Value result; + CEL_RETURN_IF_ERROR(cel::ModernValue(arena, cel_value, result)); + return result; + } } return target.GetStruct().GetFieldByName(field, descriptor_pool, message_factory, arena); @@ -308,20 +311,23 @@ absl::StatusOr> WrappedStructQualify( absl::Span qualifiers, bool presence_test, const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool, google::protobuf::MessageFactory* absl_nonnull message_factory, - google::protobuf::Arena* absl_nonnull arena) { - if (const google::protobuf::Message* message = - cel::interop_internal::GetLegacyMessage(struct_value); - message != nullptr) { - CelValue::MessageWrapper message_wrapper( - message, &GetGenericProtoTypeInfoInstance()); - CEL_ASSIGN_OR_RETURN(auto legacy_result, - GetGenericProtoAccessApisInstance().Qualify( - qualifiers, message_wrapper, presence_test, - MemoryManagerRef::Pooling(arena))); - Value result; - CEL_RETURN_IF_ERROR(cel::ModernValue(arena, legacy_result.value, result)); - return std::pair{std::move(result), - legacy_result.qualifier_count}; + google::protobuf::Arena* absl_nonnull arena, + bool enable_use_new_field_select_implementation) { + if (!enable_use_new_field_select_implementation) { + if (const google::protobuf::Message* message = + cel::interop_internal::GetLegacyMessage(struct_value); + message != nullptr) { + CelValue::MessageWrapper message_wrapper( + message, &GetGenericProtoTypeInfoInstance()); + CEL_ASSIGN_OR_RETURN(auto legacy_result, + GetGenericProtoAccessApisInstance().Qualify( + qualifiers, message_wrapper, presence_test, + MemoryManagerRef::Pooling(arena))); + Value result; + CEL_RETURN_IF_ERROR(cel::ModernValue(arena, legacy_result.value, result)); + return std::pair{std::move(result), + legacy_result.qualifier_count}; + } } return struct_value.Qualify(qualifiers, presence_test, descriptor_pool, message_factory, arena); @@ -331,7 +337,8 @@ absl::StatusOr ApplyQualifier( const Value& operand, const SelectQualifier& qualifier, const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool, google::protobuf::MessageFactory* absl_nonnull message_factory, - google::protobuf::Arena* absl_nonnull arena) { + google::protobuf::Arena* absl_nonnull arena, + bool enable_use_new_field_select_implementation) { return absl::visit( absl::Overload( [&](const FieldSpecifier& field_specifier) -> absl::StatusOr { @@ -341,7 +348,8 @@ absl::StatusOr ApplyQualifier( "