From b6156d908a9a6cb95bcb7b8181386f3c7fb52099 Mon Sep 17 00:00:00 2001
From: Dark-Avery <ingv0rr@ydb.tech>
Date: Mon, 15 Jan 2024 13:41:28 +0000
Subject: [PATCH] add colors in cli and fix minor bugs

---
 ydb/public/lib/ydb_cli/common/format.cpp      | 41 +++++++++++----
 .../lib/ydb_cli/common/pretty_table.cpp       | 51 ++++++++++++++++++-
 2 files changed, 80 insertions(+), 12 deletions(-)

diff --git a/ydb/public/lib/ydb_cli/common/format.cpp b/ydb/public/lib/ydb_cli/common/format.cpp
index 08b76da9781d..dc8f37c59d8b 100644
--- a/ydb/public/lib/ydb_cli/common/format.cpp
+++ b/ydb/public/lib/ydb_cli/common/format.cpp
@@ -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()) {
@@ -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));
     }
 
@@ -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();
@@ -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")) {
@@ -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;
diff --git a/ydb/public/lib/ydb_cli/common/pretty_table.cpp b/ydb/public/lib/ydb_cli/common/pretty_table.cpp
index 0dc07572ee6e..b32c8ae9d00f 100644
--- a/ydb/public/lib/ydb_cli/common/pretty_table.cpp
+++ b/ydb/public/lib/ydb_cli/common/pretty_table.cpp
@@ -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;
@@ -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;
 }
 
@@ -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;
                 }
             }