From efc0a8ac45a0d9ea85a029254bedea223f5c1e7a Mon Sep 17 00:00:00 2001 From: arkady-e1ppa Date: Mon, 30 Oct 2023 11:28:10 +0300 Subject: YT-19731: Whitelist for TError::Truncate with tests --- yt/yt/core/misc/error.cpp | 22 +++++----- yt/yt/core/misc/error.h | 11 ++++- yt/yt/core/misc/unittests/error_ut.cpp | 78 ++++++++++++++++++++++++++++++++++ 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& 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& 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& attributeWhitelist = {}) const &; + + TError Truncate( + int maxInnerErrorCount = 2, + i64 stringLimit = 16_KB, + const THashSet& 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...", 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 myWhitelist = {"attr2"}; + + auto truncatedError = error.Truncate(2, 10, myWhitelist); + + EXPECT_EQ(error.GetCode(), truncatedError.GetCode()); + EXPECT_EQ(error.GetMessage(), truncatedError.GetMessage()); + + EXPECT_EQ("......", truncatedError.Attributes().Get("attr1")); + EXPECT_EQ("Some long long attr", truncatedError.Attributes().Get("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 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("......", truncatedError.Attributes().Get("attr1")); + EXPECT_EQ("Some long long attr", truncatedError.Attributes().Get("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 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("......", truncatedInnerError.Attributes().Get("attr1")); + EXPECT_EQ("Some long long attr", truncatedInnerError.Attributes().Get("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 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("......", truncatedInnerError.Attributes().Get("attr1")); + EXPECT_EQ("Some long long attr", truncatedInnerError.Attributes().Get("attr2")); +} + TEST(TErrorTest, YTExceptionToError) { try { -- cgit v1.3