YT: Deserialize protobuf message via pull parser

417000f2824a0c0961e28b31bb2b11463589d4ed

6 files changed, 54 insertions, 25 deletions

diff --git a/yt/yt/core/yson/pull_parser.cpp b/yt/yt/core/yson/pull_parser.cpp
index f79b81b2e6..ed2ffd78d9 100644
--- a/yt/yt/core/yson/pull_parser.cpp
+++ b/yt/yt/core/yson/pull_parser.cpp
@@ -498,7 +498,6 @@ void TYsonPullParser::TransferComplexValue(IYsonConsumer* consumer, const TYsonI
         /*stopAfterAttributes*/ false);
 }
 
-
 void TYsonPullParser::SkipAttributes(const TYsonItem& previousItem)
 {
     EnsureYsonToken("attributes", *this, previousItem, EYsonItemType::BeginAttributes);
diff --git a/yt/yt/core/yson/pull_parser_deserialize-inl.h b/yt/yt/core/yson/pull_parser_deserialize-inl.h
index 9d45eec2a7..b0f6b9d8ca 100644
--- a/yt/yt/core/yson/pull_parser_deserialize-inl.h
+++ b/yt/yt/core/yson/pull_parser_deserialize-inl.h
@@ -4,11 +4,6 @@
 #include "pull_parser_deserialize.h"
 #endif
 
-#include "public.h"
-
-#include "pull_parser.h"
-#include "pull_parser_deserialize.h"
-
 #include <yt/yt/core/misc/error.h>
 
 #include <yt/yt/core/yson/token_writer.h>
@@ -360,6 +355,19 @@ void Deserialize(TStrongTypedef<T, TTag>& value, TYsonPullParserCursor* cursor)
     Deserialize(value.Underlying(), cursor);
 }
 
+template <class T>
+    requires std::derived_from<T, google::protobuf::Message>
+void Deserialize(
+    T& message,
+    NYson::TYsonPullParserCursor* cursor)
+{
+    NYson::TProtobufWriterOptions options;
+    options.UnknownYsonFieldModeResolver = NYson::TProtobufWriterOptions::CreateConstantUnknownYsonFieldModeResolver(
+        NYson::EUnknownYsonFieldsMode::Keep);
+
+    DeserializeProtobufMessage(message, NYson::ReflectProtobufMessageType<T>(), cursor, options);
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 template <typename TTo>
diff --git a/yt/yt/core/yson/pull_parser_deserialize.cpp b/yt/yt/core/yson/pull_parser_deserialize.cpp
index 0ccc9cdc42..3659fd29fb 100644
--- a/yt/yt/core/yson/pull_parser_deserialize.cpp
+++ b/yt/yt/core/yson/pull_parser_deserialize.cpp
@@ -1,5 +1,11 @@
 #include "pull_parser_deserialize.h"
 
+#include "protobuf_interop.h"
+
+#include <yt/yt/core/misc/protobuf_helpers.h>
+
+#include <library/cpp/yt/yson/consumer.h>
+
 #include <library/cpp/yt/misc/cast.h>
 
 namespace NYT::NYson {
@@ -237,6 +243,24 @@ void Deserialize(TGuid& value, TYsonPullParserCursor* cursor)
     cursor->Next();
 }
 
+void DeserializeProtobufMessage(
+    google::protobuf::Message& message,
+    const TProtobufMessageType* type,
+    NYson::TYsonPullParserCursor* cursor,
+    const NYson::TProtobufWriterOptions& options)
+{
+    TProtobufString wireBytes;
+    google::protobuf::io::StringOutputStream outputStream(&wireBytes);
+    auto protobufWriter = CreateProtobufWriter(&outputStream, type, options);
+
+    cursor->TransferComplexValue(protobufWriter.get());
+
+    if (!message.ParseFromArray(wireBytes.data(), wireBytes.size())) {
+        THROW_ERROR_EXCEPTION("Error parsing %v from wire bytes",
+            message.GetTypeName());
+    }
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace NYT::NYson
diff --git a/yt/yt/core/yson/pull_parser_deserialize.h b/yt/yt/core/yson/pull_parser_deserialize.h
index 93af8deb78..9d10f5f47c 100644
--- a/yt/yt/core/yson/pull_parser_deserialize.h
+++ b/yt/yt/core/yson/pull_parser_deserialize.h
@@ -3,6 +3,7 @@
 #include "public.h"
 
 #include "pull_parser.h"
+#include "protobuf_interop.h"
 
 #include <library/cpp/yt/containers/enum_indexed_array.h>
 
@@ -160,6 +161,18 @@ void Deserialize(std::unique_ptr<T>& value, TYsonPullParserCursor* cursor);
 template <class T, class TTag>
 void Deserialize(TStrongTypedef<T, TTag>& value, TYsonPullParserCursor* cursor);
 
+void DeserializeProtobufMessage(
+    google::protobuf::Message& message,
+    const NYson::TProtobufMessageType* type,
+    NYson::TYsonPullParserCursor* cursor,
+    const NYson::TProtobufWriterOptions& options = {});
+
+template <class T>
+    requires std::derived_from<T, google::protobuf::Message>
+void Deserialize(
+    T& message,
+    NYson::TYsonPullParserCursor* cursor);
+
 ////////////////////////////////////////////////////////////////////////////////
 
 template <typename TTo>
diff --git a/yt/yt/core/ytree/serialize-inl.h b/yt/yt/core/ytree/serialize-inl.h
index 59d18ee68d..6d846e42dd 100644
--- a/yt/yt/core/ytree/serialize-inl.h
+++ b/yt/yt/core/ytree/serialize-inl.h
@@ -625,10 +625,10 @@ void DeserializeProtobufMessage(
     const NYson::TProtobufWriterOptions& options = {});
 
 template <class T>
+    requires std::derived_from<T, google::protobuf::Message>
 void Deserialize(
     T& message,
-    const INodePtr& node,
-    typename std::enable_if<std::is_convertible<T*, google::protobuf::Message*>::value, void>::type*)
+    const INodePtr& node)
 {
     NYson::TProtobufWriterOptions options;
     options.UnknownYsonFieldModeResolver = NYson::TProtobufWriterOptions::CreateConstantUnknownYsonFieldModeResolver(
@@ -636,15 +636,6 @@ void Deserialize(
     DeserializeProtobufMessage(message, NYson::ReflectProtobufMessageType<T>(), node, options);
 }
 
-template <class T>
-void Deserialize(
-    T& message,
-    NYson::TYsonPullParserCursor* cursor,
-    typename std::enable_if<std::is_convertible<T*, google::protobuf::Message*>::value, void>::type*)
-{
-    Deserialize(message, NYson::ExtractTo<NYTree::INodePtr>(cursor));
-}
-
 template <class T, class TTag>
 void Deserialize(TStrongTypedef<T, TTag>& value, INodePtr node)
 {
diff --git a/yt/yt/core/ytree/serialize.h b/yt/yt/core/ytree/serialize.h
index 5791c2e16f..2a9c0875ee 100644
--- a/yt/yt/core/ytree/serialize.h
+++ b/yt/yt/core/ytree/serialize.h
@@ -252,16 +252,10 @@ void Deserialize(TEnumIndexedArray<E, T, Min, Max>& vector, INodePtr node);
 
 // Subtypes of google::protobuf::Message
 template <class T>
+    requires std::derived_from<T, google::protobuf::Message>
 void Deserialize(
     T& message,
-    const INodePtr& node,
-    typename std::enable_if<std::is_convertible<T*, google::protobuf::Message*>::value, void>::type* = nullptr);
-
-template <class T>
-void Deserialize(
-    T& message,
-    NYson::TYsonPullParserCursor* cursor,
-    typename std::enable_if<std::is_convertible<T*, google::protobuf::Message*>::value, void>::type* = nullptr);
+    const INodePtr& node);
 
 template <class T, class TTag>
 void Deserialize(TStrongTypedef<T, TTag>& value, INodePtr node);