fix bug with table (#1998)

3 files changed, 249 insertions, 101 deletions

diff --git a/ydb/public/lib/ydb_cli/common/format.cpp b/ydb/public/lib/ydb_cli/common/format.cpp
index 64f1217e82f..481dd305dcd 100644
--- a/ydb/public/lib/ydb_cli/common/format.cpp
+++ b/ydb/public/lib/ydb_cli/common/format.cpp
@@ -402,6 +402,20 @@ void TQueryPlanPrinter::PrintPrettyTable(const NJson::TJsonValue& plan) {
     }
 }
 
+TString replaceAll(TString str, const TString& from, const TString& to) {
+    if (!from) {
+        return str;
+    }
+        
+    size_t startPos = 0;
+    while ((startPos = str.find(from, startPos)) != TString::npos) {
+        str.replace(startPos, from.length(), to);
+        startPos += to.length();
+    }
+
+    return str;
+}
+
 void TQueryPlanPrinter::PrintPrettyTableImpl(const NJson::TJsonValue& plan, TString& offset, TPrettyTable& table) {
     const auto& node = plan.GetMapSafe();
 
@@ -431,13 +445,13 @@ void TQueryPlanPrinter::PrintPrettyTableImpl(const NJson::TJsonValue& plan, TStr
     TStringBuf color;
     switch(offset.size() % 3) {
         case 0: 
-            color = colors.Red();
+            color = colors.LightRed();
             break;
         case 1:
-            color = colors.Green();
+            color = colors.LightGreen();
             break;
         case 2:
-            color = colors.Blue();
+            color = colors.LightBlue();
             break;
         default:
             color = colors.Default();
@@ -459,15 +473,15 @@ void TQueryPlanPrinter::PrintPrettyTableImpl(const NJson::TJsonValue& plan, TStr
                 } else if (key == "E-Rows") {
                     eRows = JsonToString(value);
                 } else if (key != "Name") {
-                    info.emplace_back(TStringBuilder() << key << ": " << JsonToString(value));
+                    info.emplace_back(TStringBuilder() << colors.LightYellow() << key << colors.Default() << ": " << replaceAll(replaceAll(JsonToString(value), "item.", ""), "state.", ""));
                 }
             }
 
             TStringBuilder operation;
             if (info.empty()) {
-                operation << offset << color << " -> " << colors.Default() << op.GetMapSafe().at("Name").GetString();
+                operation << offset << color << " -> " << colors.LightCyan() << op.GetMapSafe().at("Name").GetString() << colors.Default();
             } else {
-                operation << offset << color << " -> " << colors.Default() << op.GetMapSafe().at("Name").GetString()
+                operation << offset << color << " -> " << colors.LightCyan() << op.GetMapSafe().at("Name").GetString() << colors.Default()
                      << " (" << JoinStrings(info, ", ") << ")";
             }
 
@@ -484,7 +498,7 @@ void TQueryPlanPrinter::PrintPrettyTableImpl(const NJson::TJsonValue& plan, TStr
         }
     } else {
         TStringBuilder operation;
-        operation << offset << color << " -> " << colors.Default() << node.at("Node Type").GetString();
+        operation << offset << color << " -> " << colors.LightCyan() << node.at("Node Type").GetString() << colors.Default();
         newRow.Column(0, std::move(operation));
     }
 
diff --git a/ydb/public/lib/ydb_cli/common/pretty_table.cpp b/ydb/public/lib/ydb_cli/common/pretty_table.cpp
index b32c8ae9d00..96c2d39a20f 100644
--- a/ydb/public/lib/ydb_cli/common/pretty_table.cpp
+++ b/ydb/public/lib/ydb_cli/common/pretty_table.cpp
@@ -1,6 +1,7 @@
 #include "pretty_table.h"
 #include "common.h"
 
+#include <library/cpp/colorizer/colors.h>
 #include <util/generic/algorithm.h>
 #include <util/generic/xrange.h>
 #include <util/stream/format.h>
@@ -15,131 +16,261 @@ TPrettyTable::TRow::TRow(size_t nColumns)
 {
 }
 
-size_t TotalAnsiEscapeCodeLen(TStringBuf text) {
+size_t TPrettyTable::TRow::ColumnWidth(size_t columnIndex) const {
+    Y_ABORT_UNLESS(columnIndex < Columns.size());
     enum {
         TEXT,
-        BEFORE_CODE,
-        IN_CODE,
+        COLOR,
+        UTF8,
     } 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;
-    for (char ch : data) {
-        size_t n = 0;
-        // if the first bit of the character is not 0, we met a multibyte
-        // counting the number of single bits at the beginning of a byte
-        while ((ch & 0x80) != 0) {
-            n++;
-            ch <<= 1;
-        }
-        // update counter
-        if (n != 0) {
-            extraBytes += n - 1;
-        }
-    }
-    // update counter with len of color
-    extraBytes += TotalAnsiEscapeCodeLen(data);
-    
-    return extraBytes;
-}
-
-size_t TPrettyTable::TRow::ColumnWidth(size_t columnIndex) const {
-    Y_ABORT_UNLESS(columnIndex < Columns.size());
-
     size_t width = 0;
     TStringBuf data;
     for (const auto& line : Columns.at(columnIndex)) {
         data = line;
 
-        size_t extraBytes = ExtraBytes(data);
+        // flag of first symbol in color
+        bool first = false;
+        // flag of enfing of color
+        bool endcolor = false;
+        int utf8Len = 0;
+        // count of visible chars
+        size_t curLen = 0;
+        for (char ch : data) {
+            switch (state) {
+                case TEXT:
+                    // begin of color
+                    if (ch == '\033') {
+                        state = COLOR;
+                        first = true;
+                        endcolor = false;
+                    // begin utf8
+                    } else  if ((ch & 0x80) != 0) {
+                        curLen++;
+                        utf8Len = 0;
+                        // if the first bit of the character is not 0, we met a multibyte
+                        // counting the number of single bits at the beginning of a byte
+                        while ((ch & 0x80) != 0) {
+                            utf8Len++;
+                            ch <<= 1;
+                        }
+                        state = UTF8;
+                    // common text
+                    } else {
+                        curLen++;
+                    }
+                    break;
+                case UTF8:
+                    // skip n chars
+                    utf8Len -= 1;
+                    if (utf8Len == 0) {
+                        curLen++;
+                        while ((ch & 0x80) != 0) {
+                            utf8Len++;
+                            ch <<= 1;
+                        }
+                        if (utf8Len != 0) {
+                            state = UTF8;
+                        } else {
+                            state = TEXT;
+                        }
+                    }
+                    break;
+                case COLOR:
+                    // first symbol must be [
+                    if (first) {
+                        if (ch != '[') {
+                            state = TEXT;
+                        }
+                        first = false;
+                    // at the end of color can be digits, m and ;
+                    } else if (endcolor) {
+                        if (ch != ';' && !isdigit(ch) && ch != 'm' ) {
+                            curLen++;
+                            state = TEXT;
+                        }
+                    // ending after ;
+                    } else {
+                        if (ch == ';') {
+                            endcolor = true;
+                        }
+                    }
+                    break;
+            }
+        }
 
-        width = Max(width, line.size() - extraBytes);
+        width = Max(width, curLen);
     }
 
     return width;
 }
 
-bool TPrettyTable::TRow::PrintColumns(IOutputStream& o, const TVector<size_t>& widths, size_t lineNumber) const {
+size_t TPrettyTable::TRow::PrintColumns(IOutputStream& o, const TVector<size_t>& widths, size_t offset, TString& oldColor) const {
     bool next = false;
-
+    size_t firstColLen = 0;
+    NColorizer::TColors colors = NColorizer::AutoColors(Cout);
+    
     for (size_t columnIndex : xrange(Columns.size())) {
+        enum {
+            TEXT,
+            COLOR,
+            UTF8,
+        } state = TEXT;
+        o << colors.Default();
         if (columnIndex == 0) {
             o << "│ ";
+            o << oldColor;
         } else {
             o << " │ ";
         }
-
+        
         if (size_t width = widths.at(columnIndex)) {
             const auto& column = Columns.at(columnIndex);
 
             TStringBuf data;
-            size_t extraBytes;
-            size_t l = 0;
+            // count of visible chars
+            size_t curLen = 0;
+            // count of chars
+            size_t absCurLen = 0;
+            // flag of first symbol in color
+            bool first = false;
+            // flag of enfing of color
+            bool endcolor = false;
+
+            absCurLen = 0;
             for (const auto& line : column) {
                 data = line;
-                extraBytes = ExtraBytes(data);
-                if (data && l < lineNumber) {
-                    data.Skip(extraBytes);
+                TString s = "";
+                
+                data.Skip(offset);
+                // len of utf8 symbol
+                int utf8Len = 0;
+                
+                for (char ch : data) {
+                    if (next && columnIndex == 0) {
+                        break;
+                    }
+                    switch (state) {
+                        case TEXT:
+                            // begin of color
+                            if (ch == '\033') {
+                                
+                                state = COLOR;
+                                first = true;
+                                endcolor = false;
+                                o << s;
+                                s = ch;
+                            // begin utf8
+                            } else  if ((ch & 0x80) != 0) {
+                                o << s;
+                                s = "";
+                                utf8Len = 0;
+                                curLen++;
+                                // if the first bit of the character is not 0, we met a multibyte
+                                // counting the number of single bits at the beginning of a byte
+                                while ((ch & 0x80) != 0) {
+                                    utf8Len++;
+                                    ch <<= 1;
+                                }
+                                if (curLen + utf8Len > width) {
+                                    next = true;
+                                    curLen -= 1;
+                                    break;
+                                }
+                                o << data.SubStr(absCurLen, utf8Len);
+                                state = UTF8;
+                            // common text
+                            } else {
+                                curLen++;
+                                if (curLen > width) {
+                                    next = true;
+                                    curLen -= 1;
+                                    break;
+                                }
+                                s += ch; 
+                            }
+                            break;
+                        case UTF8:
+                            // skip n chars
+                            utf8Len -= 1;
+                            if (utf8Len == 0) {
+                                while ((ch & 0x80) != 0) {
+                                    utf8Len++;
+                                    ch <<= 1;
+                                }
+                                if (curLen + utf8Len > width) {
+                                    next = true;
+                                    curLen -= 1;
+                                    break;
+                                }
+                                curLen++;
+                                if (utf8Len != 0) {
+                                    o << data.SubStr(absCurLen, utf8Len);
+                                    state = UTF8;
+                                } else {
+                                    s = ch;
+                                    state = TEXT;
+                                }
+                            }
+                            break;
+                        case COLOR:
+                            // first symbol must be [
+                            if (first) {
+                                if (ch != '[') {
+                                    o << s;
+                                    s = ch;
+                                    state = TEXT;
+                                } else {
+                                    s += ch;
+                                }
+                                first = false;
+                            // at the end of color can be digits, m and ;
+                            } else if (endcolor) {
+                                if (ch != ';' && !isdigit(ch) && ch != 'm' ) {
+                                    o << s;
+                                    oldColor = s;
+                                    curLen++;
+                                    if (curLen > width) {
+                                        next = true;
+                                        curLen -= 1;
+                                        break;
+                                    }
+                                    s = ch;
+                                    state = TEXT;
+                                } else {
+                                    s += ch;
+                                }
+                            // ending after ;
+                            } else {
+                                if (ch == ';') {
+                                    endcolor = true;
+                                }
+                                s += ch;
+                            }
+                            
+                            break;
+                    }
+                    absCurLen++;
                 }
-                while (data && l < lineNumber) {
-                    data.Skip(width);
-                    ++l;
+                
+                if (s != "") {
+                    o << s;
                 }
             }
-            extraBytes = ExtraBytes(data);
-            width += extraBytes;
-            
-
-            if (data) {
-                o << RightPad(data.SubStr(0, width), width);
-            } else {
-                o << RightPad(' ', width);
-            }
-
-            if (data.size() > width) {
-                next = true;
+            o << TString(width - curLen, ' ');
+            if (columnIndex == 0) {
+                firstColLen = absCurLen - 1;
             }
         }
     }
+    
+    o << colors.Default();
     o << " │" << Endl;
-
-    return next;
+    if (next) {
+        return firstColLen;
+    } else {
+        return 0;
+    }
 }
 
 bool TPrettyTable::TRow::HasFreeText() const {
@@ -194,9 +325,12 @@ void TPrettyTable::Print(IOutputStream& o) const {
     for (auto i : xrange(Rows.size())) {
         const auto& row = Rows.at(i);
 
-        size_t line = 0;
-        while (row.PrintColumns(o, widths, line)) {
-            ++line;
+        size_t res = 1;
+        size_t offset = 0;
+        TString oldColor = "";
+        while (res != 0) {
+            res = row.PrintColumns(o, widths, offset, oldColor);
+            offset += res;
         }
 
         if (row.HasFreeText()) {
diff --git a/ydb/public/lib/ydb_cli/common/pretty_table.h b/ydb/public/lib/ydb_cli/common/pretty_table.h
index 4637e7f13ed..36de1df6da4 100644
--- a/ydb/public/lib/ydb_cli/common/pretty_table.h
+++ b/ydb/public/lib/ydb_cli/common/pretty_table.h
@@ -76,7 +76,7 @@ public:
     private:
         size_t ColumnWidth(size_t columnIndex) const;
         size_t ExtraBytes(TStringBuf data) const;
-        bool PrintColumns(IOutputStream& o, const TVector<size_t>& widths, size_t lineNumber) const;
+        size_t PrintColumns(IOutputStream& o, const TVector<size_t>& widths, size_t offset, TString& oldColor) const;
         bool HasFreeText() const;
         void PrintFreeText(IOutputStream& o, size_t width) const;