Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add colors in CLI and fix bugs #1007

Merged
merged 1 commit into from
Jan 15, 2024
Merged

Add colors in CLI and fix bugs #1007

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 31 additions & 10 deletions ydb/public/lib/ydb_cli/common/format.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -439,6 +439,22 @@ void TQueryPlanPrinter::PrintPrettyTableImpl(const NJson::TJsonValue& plan, TStr
newRow.Column(2, std::move(nRows));
}

NColorizer::TColors colors = NColorizer::AutoColors(Cout);
TStringBuf color;
switch(offset.size() % 3) {
case 0:
color = colors.Red();
break;
case 1:
color = colors.Green();
break;
case 2:
color = colors.Blue();
break;
default:
color = colors.Default();
break;
}

if (node.contains("Operators")) {
for (const auto& op : node.at("Operators").GetArraySafe()) {
Expand All @@ -458,26 +474,25 @@ void TQueryPlanPrinter::PrintPrettyTableImpl(const NJson::TJsonValue& plan, TStr

TStringBuilder operation;
if (info.empty()) {
operation << offset << " -> " << op.GetMapSafe().at("Name").GetString();
operation << offset << color << " -> " << colors.Default() << op.GetMapSafe().at("Name").GetString();
} else {
operation << offset << " -> " << op.GetMapSafe().at("Name").GetString()
operation << offset << color << " -> " << colors.Default() << op.GetMapSafe().at("Name").GetString()
<< " (" << JoinStrings(info, ", ") << ")";
}

auto& row = table.AddRow();
row.Column(0, std::move(operation));
newRow.Column(0, std::move(operation));
if (AnalyzeMode) {
row.Column(3, std::move(eCost));
row.Column(4, std::move(eRows));
newRow.Column(3, std::move(eCost));
newRow.Column(4, std::move(eRows));
}
else {
row.Column(1, std::move(eCost));
row.Column(2, std::move(eRows));
newRow.Column(1, std::move(eCost));
newRow.Column(2, std::move(eRows));
}
}
} else {
TStringBuilder operation;
operation << offset << " -> " << node.at("Node Type").GetString();
operation << offset << color << " -> " << colors.Default() << node.at("Node Type").GetString();
newRow.Column(0, std::move(operation));
}

Expand Down Expand Up @@ -614,6 +629,8 @@ NJson::TJsonValue TQueryPlanPrinter::ReconstructQueryPlanRec(const NJson::TJsonV

TVector<NJson::TJsonValue> planInputs;

auto opName = op.GetMapSafe().at("Name").GetStringSafe();

for (auto opInput : op.GetMapSafe().at("Inputs").GetArraySafe()) {
if (opInput.GetMapSafe().contains("ExternalPlanNodeId")) {
auto inputPlanKey = opInput.GetMapSafe().at("ExternalPlanNodeId").GetIntegerSafe();
Expand All @@ -623,6 +640,10 @@ NJson::TJsonValue TQueryPlanPrinter::ReconstructQueryPlanRec(const NJson::TJsonV
auto inputPlanId = opInput.GetMapSafe().at("InternalOperatorId").GetIntegerSafe();
planInputs.push_back( ReconstructQueryPlanRec(plan, inputPlanId, planIndex, precomputes, nodeCounter));
}
// temp hack
if (opName == "Filter") {
break;
}
}

if (op.GetMapSafe().contains("Inputs")) {
Expand All @@ -642,7 +663,7 @@ NJson::TJsonValue TQueryPlanPrinter::ReconstructQueryPlanRec(const NJson::TJsonV
}
}

result["Node Type"] = op.GetMapSafe().at("Name").GetStringSafe();
result["Node Type"] = opName;
NJson::TJsonValue newOps;
newOps.AppendValue(op);
result["Operators"] = newOps;
Expand Down
51 changes: 49 additions & 2 deletions ydb/public/lib/ydb_cli/common/pretty_table.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,48 @@ TPrettyTable::TRow::TRow(size_t nColumns)
{
}

size_t TotalAnsiEscapeCodeLen(TStringBuf text) {
enum {
TEXT,
BEFORE_CODE,
IN_CODE,
} state = TEXT;

size_t totalLen = 0;
size_t curLen = 0;

for (auto it = text.begin(); it < text.end(); ++it) {
switch (state) {
case TEXT:
if (*it == '\033') {
state = BEFORE_CODE;
curLen = 1;
}
break;
case BEFORE_CODE:
if (*it == '[') {
state = IN_CODE;
curLen++;
} else {
state = TEXT;
}
break;
case IN_CODE:
if (*it == ';' || isdigit(*it)) {
curLen++;
} else {
if (*it == 'm') {
totalLen += curLen + 1;
}
state = TEXT;
}
break;
}
}

return totalLen;
}

size_t TPrettyTable::TRow::ExtraBytes(TStringBuf data) const {
// counter of previously uncounted bytes
size_t extraBytes = 0;
Expand All @@ -31,7 +73,9 @@ size_t TPrettyTable::TRow::ExtraBytes(TStringBuf data) const {
extraBytes += n - 1;
}
}

// update counter with len of color
extraBytes += TotalAnsiEscapeCodeLen(data);

return extraBytes;
}

Expand Down Expand Up @@ -70,8 +114,11 @@ bool TPrettyTable::TRow::PrintColumns(IOutputStream& o, const TVector<size_t>& w
for (const auto& line : column) {
data = line;
extraBytes = ExtraBytes(data);
if (data && l < lineNumber) {
data.Skip(extraBytes);
}
while (data && l < lineNumber) {
data.Skip(width + extraBytes);
data.Skip(width);
++l;
}
}
Expand Down
Loading