Intermediate changes

4 files changed, 35 insertions, 1 deletions

diff --git a/yt/yt/core/yson/protobuf_interop.cpp b/yt/yt/core/yson/protobuf_interop.cpp
index 450ac2bbfc..1a1db0a5b7 100644
--- a/yt/yt/core/yson/protobuf_interop.cpp
+++ b/yt/yt/core/yson/protobuf_interop.cpp
@@ -785,7 +785,8 @@ TProtobufElement TProtobufField::GetElement(bool insideRepeated) const
         });
     } else {
         return std::make_unique<TProtobufScalarElement>(TProtobufScalarElement{
-            static_cast<TProtobufScalarElement::TType>(GetType())
+            static_cast<TProtobufScalarElement::TType>(GetType()),
+            GetEnumYsonStorageType()
         });
     }
 }
diff --git a/yt/yt/core/yson/protobuf_interop.h b/yt/yt/core/yson/protobuf_interop.h
index 8ee1ed2afc..3e6faa6a53 100644
--- a/yt/yt/core/yson/protobuf_interop.h
+++ b/yt/yt/core/yson/protobuf_interop.h
@@ -79,6 +79,9 @@ struct TProtobufScalarElement
 {
     YT_DEFINE_STRONG_TYPEDEF(TType, int);
     TType Type;
+
+    // Meaningful only when TYPE == TYPE_ENUM.
+    EEnumYsonStorageType EnumStorageType;
 };
 
 struct TProtobufAttributeDictionaryElement
diff --git a/yt/yt/core/yson/unittests/proto/protobuf_scalar_type_ut.proto b/yt/yt/core/yson/unittests/proto/protobuf_scalar_type_ut.proto
index de5f80bc53..927fc3b356 100644
--- a/yt/yt/core/yson/unittests/proto/protobuf_scalar_type_ut.proto
+++ b/yt/yt/core/yson/unittests/proto/protobuf_scalar_type_ut.proto
@@ -46,4 +46,14 @@ message TExampleMessage
     map<string, string> string_to_string_map = 35[(NYT.NYson.NProto.yson_map) = true];
     map<int32, int32> int32_to_int32_map = 36[(NYT.NYson.NProto.yson_map) = true];
     map<int32, string> int32_to_string_map = 37[(NYT.NYson.NProto.yson_map) = true];
+
+    enum EEnum {
+        VALUE0 = 0;
+        VALUE1 = 1;
+    };
+
+    optional EEnum enum_int = 38 [(NYT.NYson.NProto.enum_yson_storage_type) = EYST_INT];
+    repeated EEnum enum_int_repeated = 39 [(NYT.NYson.NProto.enum_yson_storage_type) = EYST_INT];
+    optional EEnum enum_string = 40 [(NYT.NYson.NProto.enum_yson_storage_type) = EYST_STRING];
+    repeated EEnum enum_string_repeated = 41 [(NYT.NYson.NProto.enum_yson_storage_type) = EYST_STRING];
 }
diff --git a/yt/yt/core/yson/unittests/protobuf_scalar_type_ut.cpp b/yt/yt/core/yson/unittests/protobuf_scalar_type_ut.cpp
index a0f11ccda6..857788bbd0 100644
--- a/yt/yt/core/yson/unittests/protobuf_scalar_type_ut.cpp
+++ b/yt/yt/core/yson/unittests/protobuf_scalar_type_ut.cpp
@@ -187,6 +187,26 @@ TEST(TProtobufScalarTypeTest, MapNestedTypes)
         FieldDescriptor::TYPE_STRING);
 }
 
+#define EXPECT_ENUM_STORAGE_TYPE(messageType, path, expectedEnumType) \
+    do { \
+        auto type = ReflectProtobufMessageType<NProto::messageType>(); \
+        auto element = ResolveProtobufElementByYPath(type, path).Element; \
+        EXPECT_TRUE(std::holds_alternative<std::unique_ptr<TProtobufScalarElement>>(element)); \
+        auto elementPtr = std::get<std::unique_ptr<TProtobufScalarElement>>(element).get(); \
+        int fieldType = static_cast<int>(elementPtr->Type); \
+        EXPECT_EQ(fieldType, FieldDescriptor::TYPE_ENUM); \
+        EXPECT_EQ(elementPtr->EnumStorageType, expectedEnumType); \
+    } while (false);
+
+TEST(TProtobufScalarTypeTest, EnumTypes)
+{
+    EXPECT_ENUM_STORAGE_TYPE(TExampleMessage, "/enum_int", EEnumYsonStorageType::Int);
+    EXPECT_ENUM_STORAGE_TYPE(TExampleMessage, "/enum_int_repeated/0", EEnumYsonStorageType::Int);
+
+    EXPECT_ENUM_STORAGE_TYPE(TExampleMessage, "/enum_string", EEnumYsonStorageType::String);
+    EXPECT_ENUM_STORAGE_TYPE(TExampleMessage, "/enum_string_repeated/0", EEnumYsonStorageType::String);
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace