Skip to content

Commit d0aaa98

Browse files
author
mikari
committed
YsonStruct: ProtoAsString, fixed tuple support. Refactored flow yson serialization.
commit_hash:97979bcd3f3441cf99c984eb71174eeaf8c9d8a5
1 parent e344895 commit d0aaa98

File tree

7 files changed

+337
-26
lines changed

7 files changed

+337
-26
lines changed

‎yt/yt/core/ytree/proto_yson_struct.h‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
#pragma once
2+
3+
namespace NYT::NYTree {
4+
5+
////////////////////////////////////////////////////////////////////////////////
6+
7+
struct TProtoStorageTypeString
8+
{ };
9+
10+
////////////////////////////////////////////////////////////////////////////////
11+
12+
template <class T>
13+
concept CProtobufMessage = std::derived_from<std::decay_t<T>, google::protobuf::Message>;
14+
15+
template <class T>
16+
concept CProtobufMessageAsYson = CProtobufMessage<T> && !std::derived_from<std::decay_t<T>, TProtoStorageTypeString>;
17+
18+
template <class T>
19+
concept CProtobufMessageAsString = CProtobufMessage<T> && std::derived_from<std::decay_t<T>, TProtoStorageTypeString>;
20+
21+
////////////////////////////////////////////////////////////////////////////////
22+
23+
template <CProtobufMessage TMessage>
24+
class TProtoSerializedAsString
25+
: public TProtoStorageTypeString
26+
, public TMessage
27+
{
28+
public:
29+
using TMessage::TMessage;
30+
};
31+
32+
template <CProtobufMessage TMessage>
33+
class TProtoSerializedAsYson
34+
: public TMessage
35+
{
36+
public:
37+
using TMessage::TMessage;
38+
};
39+
40+
////////////////////////////////////////////////////////////////////////////////
41+
42+
} // namespace NYT::NYTree

‎yt/yt/core/ytree/serialize-inl.h‎

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -485,15 +485,22 @@ void SerializeProtobufMessage(
485485
const NYson::TProtobufMessageType* type,
486486
NYson::IYsonConsumer* consumer);
487487

488-
template <class T>
488+
template <CProtobufMessageAsYson T>
489489
void Serialize(
490490
const T& message,
491-
NYson::IYsonConsumer* consumer,
492-
typename std::enable_if<std::is_convertible<T*, ::google::protobuf::Message*>::value, void>::type*)
491+
NYson::IYsonConsumer* consumer)
493492
{
494493
SerializeProtobufMessage(message, NYson::ReflectProtobufMessageType<T>(), consumer);
495494
}
496495

496+
template <CProtobufMessageAsString T>
497+
void Serialize(
498+
const T& message,
499+
NYson::IYsonConsumer* consumer)
500+
{
501+
consumer->OnStringScalar(message.SerializeAsStringOrThrow());
502+
}
503+
497504
template <class T, class TTag>
498505
void Serialize(const TStrongTypedef<T, TTag>& value, NYson::IYsonConsumer* consumer)
499506
{
@@ -699,8 +706,7 @@ void DeserializeProtobufMessage(
699706
const INodePtr& node,
700707
const NYson::TProtobufWriterOptions& options = {});
701708

702-
template <class T>
703-
requires std::derived_from<T, google::protobuf::Message>
709+
template <CProtobufMessageAsYson T>
704710
void Deserialize(
705711
T& message,
706712
const INodePtr& node)
@@ -711,6 +717,16 @@ void Deserialize(
711717
DeserializeProtobufMessage(message, NYson::ReflectProtobufMessageType<T>(), node, options);
712718
}
713719

720+
template <CProtobufMessageAsString T>
721+
void Deserialize(
722+
T& message,
723+
const INodePtr& node)
724+
{
725+
TString string;
726+
Deserialize(string, node);
727+
message.ParseFromStringOrThrow(string);
728+
}
729+
714730
template <class T, class TTag>
715731
void Deserialize(TStrongTypedef<T, TTag>& value, INodePtr node)
716732
{

‎yt/yt/core/ytree/serialize.h‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

33
#include "public.h"
44

5+
#include "proto_yson_struct.h"
6+
57
#include <yt/yt/core/yson/producer.h>
68

79
#include <yt/yt/core/misc/guid.h>
@@ -169,11 +171,15 @@ template <class E, class T, E Min, E Max>
169171
void Serialize(const TEnumIndexedArray<E, T, Min, Max>& value, NYson::IYsonConsumer* consumer);
170172

171173
// Subtypes of google::protobuf::Message
172-
template <class T>
174+
template <CProtobufMessageAsYson T>
175+
void Serialize(
176+
const T& message,
177+
NYson::IYsonConsumer* consumer);
178+
179+
template <CProtobufMessageAsString T>
173180
void Serialize(
174181
const T& message,
175-
NYson::IYsonConsumer* consumer,
176-
typename std::enable_if<std::is_convertible<T*, google::protobuf::Message*>::value, void>::type* = nullptr);
182+
NYson::IYsonConsumer* consumer);
177183

178184
template <class T, class TTag>
179185
void Serialize(const TStrongTypedef<T, TTag>& value, NYson::IYsonConsumer* consumer);
@@ -286,8 +292,12 @@ template <class E, class T, E Min, E Max>
286292
void Deserialize(TEnumIndexedArray<E, T, Min, Max>& vector, INodePtr node);
287293

288294
// Subtypes of google::protobuf::Message
289-
template <class T>
290-
requires std::derived_from<T, google::protobuf::Message>
295+
template <CProtobufMessageAsYson T>
296+
void Deserialize(
297+
T& message,
298+
const INodePtr& node);
299+
300+
template <CProtobufMessageAsString T>
291301
void Deserialize(
292302
T& message,
293303
const INodePtr& node);

‎yt/yt/core/ytree/unittests/yson_schema_ut.cpp‎

Lines changed: 139 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ struct TTestYsonStruct
8484
registrar.Parameter("int_map", &TThis::IntMap)
8585
.Default();
8686
registrar.Parameter("my_string_list", &TThis::MyStringList)
87-
.Default();
87+
.Default();
8888
registrar.Parameter("nullable_int", &TThis::NullableInt)
8989
.Default();
9090
registrar.Parameter("my_uint", &TThis::MyUint)
@@ -124,6 +124,33 @@ struct TTestStructWithProtobuf
124124
}
125125
};
126126

127+
struct TTestStructWithProtobufAsYson
128+
: public TYsonStruct
129+
{
130+
std::optional<TProtoSerializedAsYson<NProto::TTestMessage>> MyMessage;
131+
132+
REGISTER_YSON_STRUCT(TTestStructWithProtobufAsYson);
133+
134+
static void Register(TRegistrar registrar)
135+
{
136+
registrar.Parameter("my_message", &TThis::MyMessage)
137+
.Optional();
138+
}
139+
};
140+
141+
struct TTestStructWithProtobufAsString
142+
: public TYsonStruct
143+
{
144+
std::optional<TProtoSerializedAsString<NProto::TTestMessage>> MyMessage;
145+
146+
REGISTER_YSON_STRUCT(TTestStructWithProtobufAsString);
147+
148+
static void Register(TRegistrar registrar)
149+
{
150+
registrar.Parameter("my_message", &TThis::MyMessage)
151+
.Optional();
152+
}
153+
};
127154

128155
////////////////////////////////////////////////////////////////////////////////
129156

@@ -170,6 +197,34 @@ struct TTestStructWithUndefinedType
170197
}
171198
};
172199

200+
struct TTestStructWithTuples
201+
: public TYsonStruct
202+
{
203+
std::tuple<TString, ui64, double> Tuple;
204+
std::pair<TString, TString> Pair;
205+
206+
REGISTER_YSON_STRUCT(TTestStructWithTuples)
207+
208+
static void Register(TRegistrar registrar)
209+
{
210+
registrar.Parameter("tuple", &TThis::Tuple);
211+
registrar.Parameter("pair", &TThis::Pair);
212+
}
213+
};
214+
215+
struct TTestStructWithArray
216+
: public TYsonStruct
217+
{
218+
std::array<TString, 3> StringArray;
219+
220+
REGISTER_YSON_STRUCT(TTestStructWithArray)
221+
222+
static void Register(TRegistrar registrar)
223+
{
224+
registrar.Parameter("string_array", &TThis::StringArray);
225+
}
226+
};
227+
173228
////////////////////////////////////////////////////////////////////////////////
174229

175230
void CheckSchema(const TYsonStructPtr& ysonStruct, TStringBuf expected)
@@ -216,7 +271,7 @@ TEST(TYsonStructSchemaTest, TestYsonStruct)
216271
];})");
217272
}
218273

219-
TEST(TYsonStructSchemaTest, TestSchemaForProtobufMessage)
274+
TEST(TYsonStructSchemaTest, TestSchemaForProtobuf)
220275
{
221276
CheckSchema(
222277
New<TTestStructWithProtobuf>(),
@@ -239,6 +294,46 @@ TEST(TYsonStructSchemaTest, TestSchemaForProtobufMessage)
239294
];})");
240295
}
241296

297+
TEST(TYsonStructSchemaTest, TestSchemaForProtobufAsYson)
298+
{
299+
CheckSchema(
300+
New<TTestStructWithProtobufAsYson>(),
301+
R"({
302+
type_name="struct";
303+
members=[
304+
{
305+
name="my_message";
306+
type={
307+
type_name="optional";
308+
item={
309+
type_name="struct";
310+
members=[
311+
{name="int32_field";type="int32";};
312+
{name="string_field";"type"="utf8";};
313+
];
314+
};
315+
};
316+
};
317+
];})");
318+
}
319+
320+
TEST(TYsonStructSchemaTest, TestSchemaForProtobufAsString)
321+
{
322+
CheckSchema(
323+
New<TTestStructWithProtobufAsString>(),
324+
R"({
325+
type_name="struct";
326+
members=[
327+
{
328+
name="my_message";
329+
type={
330+
"type_name"="optional";
331+
"item"="string";
332+
};
333+
};
334+
];})");
335+
}
336+
242337
TEST(TYsonStructSchemaTest, TestYsonStructWithCustomType)
243338
{
244339
CheckSchema(
@@ -273,6 +368,48 @@ TEST(TYsonStructSchemaTest, TestYsonStructWithUndefinedType)
273368
]})");
274369
}
275370

371+
TEST(TYsonStructSchemaTest, TestYsonStructWithTuples)
372+
{
373+
CheckSchema(
374+
New<TTestStructWithTuples>(),
375+
R"({type_name="struct";
376+
members=[
377+
{
378+
name="pair";
379+
required=%true;
380+
type={
381+
type_name="tuple";
382+
elements=[{"type"="string"};{"type"="string"};];
383+
};
384+
};
385+
{
386+
name="tuple";
387+
required=%true;
388+
type={
389+
type_name="tuple";
390+
elements=[{"type"="string"};{"type"="uint64"};{"type"="double"};];
391+
};
392+
};
393+
]})");
394+
}
395+
396+
TEST(TYsonStructSchemaTest, TestYsonStructWithArray)
397+
{
398+
CheckSchema(
399+
New<TTestStructWithArray>(),
400+
R"({type_name="struct";
401+
members=[
402+
{
403+
name="string_array";
404+
required=%true;
405+
type={
406+
type_name="tuple";
407+
elements=[{"type"="string"};{"type"="string"};{"type"="string"};];
408+
};
409+
};
410+
]})");
411+
}
412+
276413
////////////////////////////////////////////////////////////////////////////////
277414

278415
} // namespace

‎yt/yt/core/ytree/unittests/yson_struct_ut.cpp‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
#include <yt/yt/core/ytree/ypath_client.h>
1212
#include <yt/yt/core/ytree/yson_struct.h>
1313

14+
#include <yt/yt/core/ytree/unittests/proto/test.pb.h>
15+
1416
#include <util/stream/buffer.h>
1517

1618
#include <util/ysaveload.h>
@@ -3740,5 +3742,63 @@ TEST(TYsonStructTest, DefaultUnrecognizedStrategy3)
37403742
EXPECT_ANY_THROW(Deserialize(yson, source->AsMap()));
37413743
}
37423744

3745+
////////////////////////////////////////////////////////////////////////////////
3746+
3747+
struct TTestYsonStructWithProto
3748+
: public virtual TYsonStruct
3749+
{
3750+
NProto::TTestMessage DefaultProto;
3751+
TProtoSerializedAsYson<NProto::TTestMessage> YsonProto;
3752+
TProtoSerializedAsString<NProto::TTestMessage> StringProto;
3753+
3754+
REGISTER_YSON_STRUCT(TTestYsonStructWithProto);
3755+
3756+
static void Register(TRegistrar registrar)
3757+
{
3758+
registrar.Parameter("default_proto", &TThis::DefaultProto)
3759+
.Default();
3760+
registrar.Parameter("yson_proto", &TThis::YsonProto)
3761+
.Default();
3762+
registrar.Parameter("string_proto", &TThis::StringProto)
3763+
.Default();
3764+
}
3765+
};
3766+
3767+
using TTestYsonStructWithProtoPtr = TIntrusivePtr<TTestYsonStructWithProto>;
3768+
3769+
TEST(TYsonStructTest, ProtoSerialize)
3770+
{
3771+
NProto::TTestMessage proto;
3772+
proto.set_int32_field(532);
3773+
proto.set_string_field("abcdef");
3774+
auto serialized = proto.SerializeAsString();
3775+
3776+
auto ysonStruct = New<TTestYsonStructWithProto>();
3777+
ysonStruct->DefaultProto.CopyFrom(proto);
3778+
ysonStruct->YsonProto.CopyFrom(proto);
3779+
ysonStruct->StringProto.CopyFrom(proto);
3780+
3781+
auto node = ConvertToNode(ysonStruct);
3782+
const auto expectedNode = BuildYsonNodeFluently()
3783+
.BeginMap()
3784+
.Item("default_proto")
3785+
.BeginMap()
3786+
.Item("int32_field").Value(532)
3787+
.Item("string_field").Value("abcdef")
3788+
.EndMap()
3789+
.Item("yson_proto")
3790+
.BeginMap()
3791+
.Item("int32_field").Value(532)
3792+
.Item("string_field").Value("abcdef")
3793+
.EndMap()
3794+
.Item("string_proto").Value(serialized)
3795+
.EndMap();
3796+
EXPECT_TRUE(AreNodesEqual(node, expectedNode));
3797+
auto otherStruct = ConvertTo<TTestYsonStructWithProtoPtr>(node);
3798+
EXPECT_EQ(*otherStruct, *ysonStruct);
3799+
}
3800+
3801+
////////////////////////////////////////////////////////////////////////////////
3802+
37433803
} // namespace
37443804
} // namespace NYT::NYTree

0 commit comments

Comments
 (0)