flag for write protobuf error to exception message instead of cerr

2 files changed, 74 insertions, 19 deletions

diff --git a/library/cpp/protobuf/util/pb_io.cpp b/library/cpp/protobuf/util/pb_io.cpp
index 6270ee0624..3ec95f0a9b 100644
--- a/library/cpp/protobuf/util/pb_io.cpp
+++ b/library/cpp/protobuf/util/pb_io.cpp
@@ -3,6 +3,7 @@
 #include <library/cpp/binsaver/bin_saver.h>
 #include <library/cpp/string_utils/base64/base64.h>
 
+#include <google/protobuf/io/tokenizer.h>
 #include <google/protobuf/message.h>
 #include <google/protobuf/messagext.h>
 #include <google/protobuf/text_format.h>
@@ -79,7 +80,43 @@ namespace NProtoBuf {
     bool MergeFromString(NProtoBuf::Message& m, const TStringBuf serializedProtoMessage) {
         return MergePartialFromString(m, serializedProtoMessage) && m.IsInitialized();
     }
-}
+}  // end of namespace NProtoBuf
+
+
+namespace {
+    class TErrorCollector: public NProtoBuf::io::ErrorCollector {
+    public:
+        TErrorCollector(const NProtoBuf::Message& m, IOutputStream* errorOut, IOutputStream* warningOut)
+          : TypeName_(m.GetTypeName())
+        {
+            ErrorOut_ = errorOut ? errorOut : &Cerr;
+            WarningOut_ = warningOut ? warningOut : &Cerr;
+        }
+        void AddError(int line, int column, const TProtoStringType& message) override {
+            PrintErrorMessage(ErrorOut_, "Error", line, column, message);
+        }
+        void AddWarning(int line, int column, const TProtoStringType& message) override {
+            PrintErrorMessage(WarningOut_, "Warning", line, column, message);
+        }
+
+    private:
+        void PrintErrorMessage(IOutputStream* out, TStringBuf errorLevel, int line, int column, const TProtoStringType& message) {
+            (*out) << errorLevel << " parsing text-format ";
+            if (line >= 0) {
+                (*out) << TypeName_ << ": " << (line + 1) << ":" << (column + 1) << ": " << message;
+            } else {
+                (*out) << TypeName_ << ": " << message;
+            }
+            out->Flush();
+        }
+
+    private:
+        const TProtoStringType TypeName_;
+        IOutputStream* ErrorOut_;
+        IOutputStream* WarningOut_;
+    };
+}  // end of anonymous namespace
+
 
 int operator&(NProtoBuf::Message& m, IBinSaver& f) {
     TStringStream ss;
@@ -138,30 +175,46 @@ static void ConfigureParser(const EParseFromTextFormatOptions options,
 }
 
 void ParseFromTextFormat(IInputStream& in, NProtoBuf::Message& m,
-                         const EParseFromTextFormatOptions options) {
+                         const EParseFromTextFormatOptions options, IOutputStream* warningStream) {
     NProtoBuf::io::TCopyingInputStreamAdaptor adaptor(&in);
     NProtoBuf::TextFormat::Parser p;
     ConfigureParser(options, p);
 
+    bool writeErrorToException = options & EParseFromTextFormatOption::WriteErrorMessageToException;
+    TStringStream errorLog;
+    THolder<TErrorCollector> errorCollector;
+
+    if (writeErrorToException) {
+        errorCollector = MakeHolder<TErrorCollector>(m, &errorLog, warningStream);
+        p.RecordErrorsTo(errorCollector.Get());
+    } else if (warningStream) {
+        errorCollector = MakeHolder<TErrorCollector>(m, &Cerr, warningStream);
+        p.RecordErrorsTo(errorCollector.Get());
+    }
+
     if (!p.Parse(&adaptor, &m)) {
         // remove everything that may have been read
         m.Clear();
-        ythrow yexception() << "ParseFromTextFormat failed on Parse for " << m.GetTypeName();
+        if (Y_UNLIKELY(writeErrorToException)) {
+            ythrow yexception() << errorLog.Str();
+        } else {
+            ythrow yexception() << "ParseFromTextFormat failed on Parse for " << m.GetTypeName();
+        }
     }
 }
 
 void ParseFromTextFormat(const TString& fileName, NProtoBuf::Message& m,
-                         const EParseFromTextFormatOptions options) {
+                         const EParseFromTextFormatOptions options, IOutputStream* warningStream) {
     /* TUnbufferedFileInput is unbuffered, but TCopyingInputStreamAdaptor adds
     * a buffer on top of it. */
     TUnbufferedFileInput stream(fileName);
-    ParseFromTextFormat(stream, m, options);
+    ParseFromTextFormat(stream, m, options, warningStream);
 }
 
 bool TryParseFromTextFormat(const TString& fileName, NProtoBuf::Message& m,
-                            const EParseFromTextFormatOptions options) {
+                            const EParseFromTextFormatOptions options, IOutputStream* warningStream) {
     try {
-        ParseFromTextFormat(fileName, m, options);
+        ParseFromTextFormat(fileName, m, options, warningStream);
     } catch (std::exception&) {
         return false;
     }
@@ -170,9 +223,9 @@ bool TryParseFromTextFormat(const TString& fileName, NProtoBuf::Message& m,
 }
 
 bool TryParseFromTextFormat(IInputStream& in, NProtoBuf::Message& m,
-                            const EParseFromTextFormatOptions options) {
+                            const EParseFromTextFormatOptions options, IOutputStream* warningStream) {
     try {
-        ParseFromTextFormat(in, m, options);
+        ParseFromTextFormat(in, m, options, warningStream);
     } catch (std::exception&) {
         return false;
     }
diff --git a/library/cpp/protobuf/util/pb_io.h b/library/cpp/protobuf/util/pb_io.h
index 493c84cb5f..04c6ebb305 100644
--- a/library/cpp/protobuf/util/pb_io.h
+++ b/library/cpp/protobuf/util/pb_io.h
@@ -58,44 +58,46 @@ void SerializeToTextFormatWithEnumId(const NProtoBuf::Message& m, IOutputStream&
 
 enum class EParseFromTextFormatOption : ui64 {
     // Unknown fields will be ignored by the parser
-    AllowUnknownField = 1
+    AllowUnknownField = 1,
+    // Error message will be writen to exception message instead of cerr
+    WriteErrorMessageToException = 2
 };
 
 Y_DECLARE_FLAGS(EParseFromTextFormatOptions, EParseFromTextFormatOption);
 
 // Parse a text-format protocol message from the given file into message object.
 void ParseFromTextFormat(const TString& fileName, NProtoBuf::Message& m,
-                         const EParseFromTextFormatOptions options = {});
+                         const EParseFromTextFormatOptions options = {}, IOutputStream* warningStream = nullptr);
 // NOTE: will read `in` till the end.
 void ParseFromTextFormat(IInputStream& in, NProtoBuf::Message& m,
-                         const EParseFromTextFormatOptions options = {});
+                         const EParseFromTextFormatOptions options = {}, IOutputStream* warningStream = nullptr);
 
 /* @return              `true` if parsing was successfull and `false` otherwise.
  *
  * @see `ParseFromTextFormat`
  */
 bool TryParseFromTextFormat(const TString& fileName, NProtoBuf::Message& m,
-                            const EParseFromTextFormatOptions options = {});
+                            const EParseFromTextFormatOptions options = {}, IOutputStream* warningStream = nullptr);
 // NOTE: will read `in` till the end.
 bool TryParseFromTextFormat(IInputStream& in, NProtoBuf::Message& m,
-                            const EParseFromTextFormatOptions options = {});
+                            const EParseFromTextFormatOptions options = {}, IOutputStream* warningStream = nullptr);
 
 // @see `ParseFromTextFormat`
 template <typename T>
 static T ParseFromTextFormat(const TString& fileName,
-                             const EParseFromTextFormatOptions options = {}) {
+                             const EParseFromTextFormatOptions options = {}, IOutputStream* warningStream = nullptr) {
     T message;
-    ParseFromTextFormat(fileName, message, options);
+    ParseFromTextFormat(fileName, message, options, warningStream);
     return message;
 }
 
 // @see `ParseFromTextFormat`
 // NOTE: will read `in` till the end.
 template <typename T>
-static T ParseFromTextFormat(IInputStream& in,
-                             const EParseFromTextFormatOptions options = {}) {
+static T ParseFromTextFormat(IInputStream& in, const EParseFromTextFormatOptions options = {},
+                             IOutputStream* warningStream = nullptr) {
     T message;
-    ParseFromTextFormat(in, message, options);
+    ParseFromTextFormat(in, message, options, warningStream);
     return message;
 }