Skip to content

Commit 04ffade

Browse files
jnthntatumcopybara-github
authored andcommitted
Migrate gaer::CelProtoWrapper::CreateMessage to just call the modern
equivalent. Add overload for specifying the expected message factory and descriptor pool. PiperOrigin-RevId: 974637296
1 parent c580592 commit 04ffade

4 files changed

Lines changed: 155 additions & 22 deletions

File tree

eval/public/structs/BUILD

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,10 +30,17 @@ cc_library(
3030
deps = [
3131
":cel_proto_wrap_util",
3232
":proto_message_type_adapter",
33+
":trivial_legacy_type_info_internal",
34+
"//common:value",
3335
"//eval/public:cel_value",
3436
"//eval/public:message_wrapper",
3537
"//internal:proto_time_encoding",
36-
"@com_google_absl//absl/types:optional",
38+
"@com_google_absl//absl/base:no_destructor",
39+
"@com_google_absl//absl/base:nullability",
40+
"@com_google_absl//absl/log:absl_check",
41+
"@com_google_absl//absl/log:absl_log",
42+
"@com_google_absl//absl/status",
43+
"@com_google_absl//absl/status:statusor",
3744
"@com_google_protobuf//:duration_cc_proto",
3845
"@com_google_protobuf//:protobuf",
3946
"@com_google_protobuf//:timestamp_cc_proto",

eval/public/structs/cel_proto_wrapper.cc

Lines changed: 80 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,21 @@
1414

1515
#include "eval/public/structs/cel_proto_wrapper.h"
1616

17-
#include "absl/types/optional.h"
17+
#include <optional>
18+
19+
#include "absl/base/no_destructor.h"
20+
#include "absl/base/nullability.h"
21+
#include "absl/log/absl_check.h"
22+
#include "absl/log/absl_log.h"
23+
#include "absl/status/status.h"
24+
#include "absl/status/statusor.h"
25+
#include "common/legacy_value.h"
26+
#include "common/value.h"
1827
#include "eval/public/cel_value.h"
1928
#include "eval/public/message_wrapper.h"
2029
#include "eval/public/structs/cel_proto_wrap_util.h"
2130
#include "eval/public/structs/proto_message_type_adapter.h"
31+
#include "eval/public/structs/trivial_legacy_type_info_internal.h"
2232
#include "google/protobuf/arena.h"
2333
#include "google/protobuf/descriptor.h"
2434
#include "google/protobuf/message.h"
@@ -27,9 +37,36 @@ namespace google::api::expr::runtime {
2737

2838
namespace {
2939

40+
using ::cel::interop_internal::TrivialTypeInfo;
3041
using ::google::protobuf::Arena;
3142
using ::google::protobuf::Descriptor;
43+
using ::google::protobuf::DescriptorPool;
3244
using ::google::protobuf::Message;
45+
using ::google::protobuf::MessageFactory;
46+
47+
// Returns the arena for the given message, or the fallback arena if the
48+
// message does not have an arena.
49+
//
50+
// A global fallback arena is used to avoid allocation when the arena is not
51+
// specified. This is effectively a memory leak, but won't trigger leak
52+
// check analyzers.
53+
//
54+
// This emulates the old behavior of tolerating a nullptr arena without
55+
// triggering a crash.
56+
google::protobuf::Arena* GetArena(const Message* absl_nonnull message,
57+
google::protobuf::Arena* absl_nullable arena) {
58+
if (arena != nullptr) {
59+
return arena;
60+
}
61+
if (message->GetArena() != nullptr) {
62+
return message->GetArena();
63+
}
64+
static absl::NoDestructor<google::protobuf::Arena> fallback_arena;
65+
ABSL_LOG(WARNING)
66+
<< "CelValue: using fallback global arena for wrapping message: "
67+
<< message->GetTypeName();
68+
return fallback_arena.get();
69+
}
3370

3471
} // namespace
3572

@@ -38,14 +75,50 @@ CelValue CelProtoWrapper::InternalWrapMessage(const Message* message) {
3875
MessageWrapper(message, &GetGenericProtoTypeInfoInstance()));
3976
}
4077

41-
// CreateMessage creates CelValue from google::protobuf::Message.
42-
// As some of CEL basic types are subclassing google::protobuf::Message,
43-
// this method contains type checking and downcasts.
44-
CelValue CelProtoWrapper::CreateMessage(const Message* value, Arena* arena) {
45-
return internal::UnwrapMessageToValue(value, &InternalWrapMessage, arena);
78+
CelValue CelProtoWrapper::CreateMessage(
79+
const Message* absl_nonnull value,
80+
const google::protobuf::DescriptorPool* absl_nonnull pool,
81+
MessageFactory* absl_nonnull factory, Arena* absl_nonnull arena) {
82+
ABSL_DCHECK(value != nullptr);
83+
if (value->GetDescriptor() == nullptr || value->GetReflection() == nullptr) {
84+
// This only happens for custom google::protobuf::Message subclasses that CEL can't
85+
// support.
86+
return CelValue::CreateMessageWrapper(
87+
MessageWrapper(value, TrivialTypeInfo::GetInstance()));
88+
}
89+
90+
auto modern_value =
91+
cel::Value::WrapMessageUnsafe(value, pool, factory, arena);
92+
93+
absl::StatusOr<CelValue> cel_value = cel::LegacyValue(arena, modern_value);
94+
if (!cel_value.ok()) {
95+
// This only happens for custom google::protobuf::Message subclasses that CEL can't
96+
// support.
97+
auto* status =
98+
google::protobuf::Arena::Create<absl::Status>(arena, cel_value.status());
99+
return CelValue::CreateError(status);
100+
}
101+
return *cel_value;
102+
}
103+
104+
CelValue CelProtoWrapper::CreateMessage(const Message* absl_nullable value,
105+
Arena* absl_nullable arena) {
106+
if (value == nullptr) {
107+
return CelValue::CreateNull();
108+
}
109+
110+
if (value->GetDescriptor() == nullptr || value->GetReflection() == nullptr) {
111+
// This only happens for custom messages subclasses that CEL can't support.
112+
return CelValue::CreateMessageWrapper(
113+
MessageWrapper(value, TrivialTypeInfo::GetInstance()));
114+
}
115+
const auto* pool = value->GetDescriptor()->file()->pool();
116+
auto* factory = value->GetReflection()->GetMessageFactory();
117+
arena = GetArena(value, arena);
118+
return CreateMessage(value, pool, factory, arena);
46119
}
47120

48-
absl::optional<CelValue> CelProtoWrapper::MaybeWrapValue(
121+
std::optional<CelValue> CelProtoWrapper::MaybeWrapValue(
49122
const Descriptor* descriptor, google::protobuf::MessageFactory* factory,
50123
const CelValue& value, Arena* arena) {
51124
const Message* msg =

eval/public/structs/cel_proto_wrapper.h

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
#ifndef THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_STRUCTS_CEL_PROTO_WRAPPER_H_
22
#define THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_STRUCTS_CEL_PROTO_WRAPPER_H_
33

4+
#include <optional>
5+
46
#include "google/protobuf/duration.pb.h"
57
#include "google/protobuf/timestamp.pb.h"
6-
#include "absl/types/optional.h"
8+
#include "absl/base/nullability.h"
79
#include "eval/public/cel_value.h"
810
#include "internal/proto_time_encoding.h"
911
#include "google/protobuf/arena.h"
@@ -17,8 +19,19 @@ class CelProtoWrapper {
1719
// CreateMessage creates CelValue from google::protobuf::Message.
1820
// As some of CEL basic types are subclassing google::protobuf::Message,
1921
// this method contains type checking and downcasts.
20-
static CelValue CreateMessage(const google::protobuf::Message* value,
21-
google::protobuf::Arena* arena);
22+
static CelValue CreateMessage(const google::protobuf::Message* absl_nonnull value,
23+
const google::protobuf::DescriptorPool* absl_nonnull pool,
24+
google::protobuf::MessageFactory* absl_nonnull factory,
25+
google::protobuf::Arena* absl_nonnull arena);
26+
27+
// Prefer using the overload that takes an explicit descriptor pool and
28+
// message factory instead. This overload will use the ones associated with
29+
// the value.
30+
//
31+
// For backward compatibility, nullptr message is allowed and will result in
32+
// the CEL null_type value.
33+
static CelValue CreateMessage(const google::protobuf::Message* absl_nullable value,
34+
google::protobuf::Arena* absl_nullable arena);
2235

2336
// Internal utility for creating a CelValue wrapping a user defined type.
2437
// Assumes that the message has been properly unpacked.
@@ -43,7 +56,7 @@ class CelProtoWrapper {
4356
// message to native CelValue representation during a protobuf field read.
4457
// Just as CreateMessage should only be used when reading protobuf values,
4558
// MaybeWrapValue should only be used when assigning protobuf fields.
46-
static absl::optional<CelValue> MaybeWrapValue(
59+
static std::optional<CelValue> MaybeWrapValue(
4760
const google::protobuf::Descriptor* descriptor, google::protobuf::MessageFactory* factory,
4861
const CelValue& value, google::protobuf::Arena* arena);
4962
};

eval/public/structs/cel_proto_wrapper_test.cc

Lines changed: 50 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ class CelProtoWrapperTest : public ::testing::Test {
104104

105105
T dyn_value;
106106
CelValue cel_dyn_value =
107-
CelProtoWrapper::CreateMessage(ReflectedCopy(message).get(), arena());
107+
CelProtoWrapper::CreateMessage(ReflectedCopy(message), arena());
108108
EXPECT_THAT(cel_dyn_value.type(), Eq(cel_value.type()));
109109
EXPECT_TRUE(cel_dyn_value.GetValue(&dyn_value));
110110
EXPECT_THAT(value, Eq(dyn_value));
@@ -121,10 +121,9 @@ class CelProtoWrapperTest : public ::testing::Test {
121121
EXPECT_THAT(cel_value.MessageOrDie(), testutil::EqualsProto(*result));
122122
}
123123

124-
std::unique_ptr<google::protobuf::Message> ReflectedCopy(
125-
const google::protobuf::Message& message) {
126-
std::unique_ptr<google::protobuf::Message> dynamic_value(
127-
factory_.GetPrototype(message.GetDescriptor())->New());
124+
google::protobuf::Message* ReflectedCopy(const google::protobuf::Message& message) {
125+
google::protobuf::Message* dynamic_value =
126+
factory_.GetPrototype(message.GetDescriptor())->New(&arena_);
128127
dynamic_value->CopyFrom(message);
129128
return dynamic_value;
130129
}
@@ -213,7 +212,7 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicValueNull) {
213212
value_msg.set_null_value(protobuf::NULL_VALUE);
214213

215214
CelValue value =
216-
CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg).get(), arena());
215+
CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg), arena());
217216
EXPECT_TRUE(value.IsNull());
218217
}
219218

@@ -314,8 +313,8 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicStruct) {
314313
const std::string kFieldBool = "field_bool";
315314
(*struct_msg.mutable_fields())[kFieldInt].set_number_value(1.);
316315
(*struct_msg.mutable_fields())[kFieldBool].set_bool_value(true);
317-
CelValue value =
318-
CelProtoWrapper::CreateMessage(ReflectedCopy(struct_msg).get(), arena());
316+
auto reflected_copy = ReflectedCopy(struct_msg);
317+
CelValue value = CelProtoWrapper::CreateMessage(reflected_copy, arena());
319318
EXPECT_TRUE(value.IsMap());
320319
const CelMap* cel_map = value.MapOrDie();
321320
ASSERT_TRUE(cel_map != nullptr);
@@ -355,7 +354,7 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicValueStruct) {
355354
.set_number_value(2);
356355

357356
CelValue value =
358-
CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg).get(), arena());
357+
CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg), arena());
359358
EXPECT_TRUE(value.IsMap());
360359
EXPECT_TRUE(
361360
(*value.MapOrDie())[CelValue::CreateString(&kField1)].has_value());
@@ -398,7 +397,7 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicValueListValue) {
398397
value_msg.mutable_list_value()->add_values()->set_number_value(2.);
399398

400399
CelValue value =
401-
CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg).get(), arena());
400+
CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg), arena());
402401
EXPECT_TRUE(value.IsList());
403402
EXPECT_THAT((*value.ListOrDie())[0].DoubleOrDie(), testing::DoubleEq(1));
404403
EXPECT_THAT((*value.ListOrDie())[1].DoubleOrDie(), testing::DoubleEq(2));
@@ -426,6 +425,47 @@ TEST_F(CelProtoWrapperTest, UnwrapInvalidAny) {
426425
ASSERT_TRUE(CelProtoWrapper::CreateMessage(&any, arena()).IsError());
427426
}
428427

428+
TEST_F(CelProtoWrapperTest, CreateMessageExplicitPoolAndFactory) {
429+
TestMessage test_message;
430+
test_message.set_string_value("test");
431+
432+
CelValue value = CelProtoWrapper::CreateMessage(
433+
&test_message, google::protobuf::DescriptorPool::generated_pool(),
434+
google::protobuf::MessageFactory::generated_factory(), arena());
435+
ASSERT_TRUE(value.IsMessage());
436+
EXPECT_THAT(value.MessageOrDie(), testutil::EqualsProto(test_message));
437+
}
438+
439+
TEST_F(CelProtoWrapperTest, CreateMessageExplicitPoolAndFactoryUnpackAny) {
440+
TestMessage test_message;
441+
test_message.set_string_value("test");
442+
443+
Any any;
444+
any.PackFrom(test_message);
445+
446+
google::protobuf::DynamicMessageFactory factory(
447+
google::protobuf::DescriptorPool::generated_pool());
448+
CelValue value = CelProtoWrapper::CreateMessage(
449+
&any, google::protobuf::DescriptorPool::generated_pool(), &factory, arena());
450+
ASSERT_TRUE(value.IsMessage());
451+
EXPECT_THAT(value.MessageOrDie(), testutil::EqualsProto(test_message));
452+
}
453+
454+
TEST_F(CelProtoWrapperTest,
455+
CreateMessageExplicitPoolAndFactoryUnpackAnyNotFound) {
456+
TestMessage test_message;
457+
test_message.set_string_value("test");
458+
459+
Any any;
460+
any.PackFrom(test_message);
461+
462+
google::protobuf::DescriptorPool empty_pool;
463+
google::protobuf::DynamicMessageFactory factory(&empty_pool);
464+
CelValue value =
465+
CelProtoWrapper::CreateMessage(&any, &empty_pool, &factory, arena());
466+
EXPECT_TRUE(value.IsError());
467+
}
468+
429469
// Test support of google.protobuf.<Type>Value wrappers in CelValue.
430470
TEST_F(CelProtoWrapperTest, UnwrapBoolWrapper) {
431471
bool value = true;

0 commit comments

Comments
 (0)