YT-19731: Whitelist for TError::Truncate with tests

3 files changed, 99 insertions, 12 deletions

diff --git a/yt/yt/core/misc/error.cpp b/yt/yt/core/misc/error.cpp
index dad759be23f..d17d2ea459c 100644
--- a/yt/yt/core/misc/error.cpp
+++ b/yt/yt/core/misc/error.cpp
@@ -606,14 +606,14 @@ TError TError::Sanitize() const
 
 const TString InnerErrorsTruncatedKey("inner_errors_truncated");
 
-TError TError::Truncate(int maxInnerErrorCount, i64 stringLimit) const &
+TError TError::Truncate(int maxInnerErrorCount, i64 stringLimit, const THashSet<TStringBuf>& attributeWhitelist) const &
 {
     if (!Impl_) {
         return TError();
     }
 
-    auto truncateInnerError = [=] (const TError& innerError) {
-        return innerError.Truncate(maxInnerErrorCount, stringLimit);
+    auto truncateInnerError = [=, &attributeWhitelist] (const TError& innerError) {
+        return innerError.Truncate(maxInnerErrorCount, stringLimit, attributeWhitelist);
     };
 
     auto truncateString = [stringLimit] (TString string) {
@@ -623,10 +623,12 @@ TError TError::Truncate(int maxInnerErrorCount, i64 stringLimit) const &
         return string;
     };
 
-    auto truncateAttributes = [stringLimit] (const IAttributeDictionary& attributes) {
+    auto truncateAttributes = [stringLimit, &attributeWhitelist] (const IAttributeDictionary& attributes) {
         auto truncatedAttributes = CreateEphemeralAttributes();
         for (const auto& key : attributes.ListKeys()) {
-            if (const auto& value = attributes.FindYson(key); std::ssize(value.AsStringBuf()) > stringLimit) {
+            const auto& value = attributes.FindYson(key);
+
+            if (std::ssize(value.AsStringBuf()) > stringLimit && !attributeWhitelist.contains(key)) {
                 truncatedAttributes->SetYson(
                     key,
                     BuildYsonStringFluently()
@@ -663,14 +665,14 @@ TError TError::Truncate(int maxInnerErrorCount, i64 stringLimit) const &
     return TError(std::move(result));
 }
 
-TError TError::Truncate(int maxInnerErrorCount, i64 stringLimit) &&
+TError TError::Truncate(int maxInnerErrorCount, i64 stringLimit, const THashSet<TStringBuf>& attributeWhitelist) &&
 {
     if (!Impl_) {
         return TError();
     }
 
-    auto truncateInnerError = [=] (TError& innerError) {
-        innerError = std::move(innerError).Truncate(maxInnerErrorCount, stringLimit);
+    auto truncateInnerError = [=, &attributeWhitelist] (TError& innerError) {
+        innerError = std::move(innerError).Truncate(maxInnerErrorCount, stringLimit, attributeWhitelist);
     };
 
     auto truncateString = [stringLimit] (TString* string) {
@@ -679,9 +681,9 @@ TError TError::Truncate(int maxInnerErrorCount, i64 stringLimit) &&
         }
     };
 
-    auto truncateAttributes = [stringLimit] (IAttributeDictionary* attributes) {
+    auto truncateAttributes = [stringLimit, &attributeWhitelist] (IAttributeDictionary* attributes) {
         for (const auto& key : attributes->ListKeys()) {
-            if (std::ssize(attributes->FindYson(key).AsStringBuf()) > stringLimit) {
+            if (std::ssize(attributes->FindYson(key).AsStringBuf()) > stringLimit && !attributeWhitelist.contains(key)) {
                 attributes->SetYson(
                     key,
                     BuildYsonStringFluently()
diff --git a/yt/yt/core/misc/error.h b/yt/yt/core/misc/error.h
index b4edbf50f45..689007b7f4c 100644
--- a/yt/yt/core/misc/error.h
+++ b/yt/yt/core/misc/error.h
@@ -144,8 +144,15 @@ public:
     // TODO(gritukan): This method is used only outside of YT. Remove it.
     TError Sanitize() const;
 
-    TError Truncate(int maxInnerErrorCount = 2, i64 stringLimit = 16_KB) const &;
-    TError Truncate(int maxInnerErrorCount = 2, i64 stringLimit = 16_KB) &&;
+    TError Truncate(
+        int maxInnerErrorCount = 2,
+        i64 stringLimit = 16_KB,
+        const THashSet<TStringBuf>& attributeWhitelist = {}) const &;
+
+    TError Truncate(
+        int maxInnerErrorCount = 2,
+        i64 stringLimit = 16_KB,
+        const THashSet<TStringBuf>& attributeWhitelist = {}) &&;
 
     bool IsOK() const;
 
diff --git a/yt/yt/core/misc/unittests/error_ut.cpp b/yt/yt/core/misc/unittests/error_ut.cpp
index 22997c51189..8f3e07476ec 100644
--- a/yt/yt/core/misc/unittests/error_ut.cpp
+++ b/yt/yt/core/misc/unittests/error_ut.cpp
@@ -147,6 +147,84 @@ TEST(TErrorTest, TruncateLargeRValue)
     EXPECT_EQ("Fourth inn...<message truncated>", truncatedError.InnerErrors()[2].GetMessage());
 }
 
+TEST(TErrorTest, TruncateWhitelist)
+{
+    auto error = TError("Some error");
+    error.MutableAttributes()->Set("attr1", "Some long long attr");
+    error.MutableAttributes()->Set("attr2", "Some long long attr");
+
+    THashSet<TStringBuf> myWhitelist = {"attr2"};
+
+    auto truncatedError = error.Truncate(2, 10, myWhitelist);
+
+    EXPECT_EQ(error.GetCode(), truncatedError.GetCode());
+    EXPECT_EQ(error.GetMessage(), truncatedError.GetMessage());
+
+    EXPECT_EQ("...<attribute truncated>...", truncatedError.Attributes().Get<TString>("attr1"));
+    EXPECT_EQ("Some long long attr", truncatedError.Attributes().Get<TString>("attr2"));
+}
+
+TEST(TErrorTest, TruncateWhitelistRValue)
+{
+    auto error = TError("Some error");
+    error.MutableAttributes()->Set("attr1", "Some long long attr");
+    error.MutableAttributes()->Set("attr2", "Some long long attr");
+
+    THashSet<TStringBuf> myWhitelist = {"attr2"};
+
+    auto errorCopy = error;
+    auto truncatedError = std::move(errorCopy).Truncate(2, 10, myWhitelist);
+    EXPECT_TRUE(errorCopy.IsOK());
+
+    EXPECT_EQ(error.GetCode(), truncatedError.GetCode());
+    EXPECT_EQ(error.GetMessage(), truncatedError.GetMessage());
+
+    EXPECT_EQ("...<attribute truncated>...", truncatedError.Attributes().Get<TString>("attr1"));
+    EXPECT_EQ("Some long long attr", truncatedError.Attributes().Get<TString>("attr2"));
+}
+
+TEST(TErrorTest, TruncateWhitelistInnerErrors)
+{
+    auto innerError = TError("Inner error");
+    innerError.MutableAttributes()->Set("attr1", "Some long long attr");
+    innerError.MutableAttributes()->Set("attr2", "Some long long attr");
+
+    auto error = TError("Error") << innerError;
+
+    THashSet<TStringBuf> myWhitelist = {"attr2"};
+
+    auto truncatedError = error.Truncate(2, 20, myWhitelist);
+    EXPECT_EQ(truncatedError.InnerErrors().size(), 1u);
+
+    auto truncatedInnerError = truncatedError.InnerErrors()[0];
+    EXPECT_EQ(truncatedInnerError.GetCode(), innerError.GetCode());
+    EXPECT_EQ(truncatedInnerError.GetMessage(), innerError.GetMessage());
+    EXPECT_EQ("...<attribute truncated>...", truncatedInnerError.Attributes().Get<TString>("attr1"));
+    EXPECT_EQ("Some long long attr", truncatedInnerError.Attributes().Get<TString>("attr2"));
+}
+
+TEST(TErrorTest, TruncateWhitelistInnerErrorsRValue)
+{
+    auto innerError = TError("Inner error");
+    innerError.MutableAttributes()->Set("attr1", "Some long long attr");
+    innerError.MutableAttributes()->Set("attr2", "Some long long attr");
+
+    auto error = TError("Error") << innerError;
+
+    THashSet<TStringBuf> myWhitelist = {"attr2"};
+
+    auto errorCopy = error;
+    auto truncatedError = std::move(errorCopy).Truncate(2, 20, myWhitelist);
+    EXPECT_TRUE(errorCopy.IsOK());
+    EXPECT_EQ(truncatedError.InnerErrors().size(), 1u);
+
+    auto truncatedInnerError = truncatedError.InnerErrors()[0];
+    EXPECT_EQ(truncatedInnerError.GetCode(), innerError.GetCode());
+    EXPECT_EQ(truncatedInnerError.GetMessage(), innerError.GetMessage());
+    EXPECT_EQ("...<attribute truncated>...", truncatedInnerError.Attributes().Get<TString>("attr1"));
+    EXPECT_EQ("Some long long attr", truncatedInnerError.Attributes().Get<TString>("attr2"));
+}
+
 TEST(TErrorTest, YTExceptionToError)
 {
     try {