KIKIMR-20156: validate equal kqp request columns vs program for CS

1 files changed, 53 insertions, 15 deletions

diff --git a/ydb/core/tx/columnshard/columnshard__read_base.cpp b/ydb/core/tx/columnshard/columnshard__read_base.cpp
index 862a089688..dbd0d638f6 100644
--- a/ydb/core/tx/columnshard/columnshard__read_base.cpp
+++ b/ydb/core/tx/columnshard/columnshard__read_base.cpp
@@ -31,29 +31,67 @@ TTxReadBase::PrepareReadMetadata(const NOlap::TReadDescription& read,
 
 bool TTxReadBase::ParseProgram(NKikimrSchemeOp::EOlapProgramType programType,
     TString serializedProgram, NOlap::TReadDescription& read, const NOlap::IColumnResolver& columnResolver) {
+
+    AFL_VERIFY(!read.ColumnIds.size() || !read.ColumnNames.size());
+    std::vector<TString> names;
+    std::set<TString> namesChecker;
+    for (auto&& i : read.ColumnIds) {
+        names.emplace_back(columnResolver.GetColumnName(i));
+        AFL_VERIFY(namesChecker.emplace(names.back()).second);
+    }
     if (serializedProgram.empty()) {
-        AFL_VERIFY(!read.ColumnIds.size() || !read.ColumnNames.size());
-        NOlap::TProgramContainer container;
-        std::vector<TString> names;
-        for (auto&& i : read.ColumnIds) {
-            names.emplace_back(columnResolver.GetColumnName(i));
-        }
         for (auto&& i : read.ColumnNames) {
             names.emplace_back(i);
+            AFL_VERIFY(namesChecker.emplace(names.back()).second);
         }
+        NOlap::TProgramContainer container;
         container.OverrideProcessingColumns(std::vector<TString>(names.begin(), names.end()));
         read.SetProgram(std::move(container));
         return true;
-    }
-    NOlap::TProgramContainer ssaProgram;
-    TString error;
-    if (!ssaProgram.Init(columnResolver, programType, serializedProgram, error)) {
-        ErrorDescription = TStringBuilder() << "Can't parse SsaProgram at " << Self->TabletID() << " / " << error;
-        return false;
-    }
-    read.SetProgram(std::move(ssaProgram));
+    } else {
+        NOlap::TProgramContainer ssaProgram;
+        TString error;
+        if (!ssaProgram.Init(columnResolver, programType, serializedProgram, error)) {
+            ErrorDescription = TStringBuilder() << "Can't parse SsaProgram at " << Self->TabletID() << " / " << error;
+            return false;
+        }
 
-    return true;
+        if (names.size()) {
+            std::set<TString> programColumns;
+            for (auto&& i : ssaProgram.GetSourceColumns()) {
+                if (!i.second.IsGenerated()) {
+                    programColumns.emplace(i.second.GetColumnName());
+                }
+            }
+            //its possible dont use columns from filter where pk field compare with null and remove from PKFilter and program, but stay in kqp columns request
+            if (Self->TablesManager.HasPrimaryIndex()) {
+                for (auto&& i : Self->TablesManager.GetIndexInfo(read.GetSnapshot()).GetReplaceKey()->field_names()) {
+                    const TString cId(i.data(), i.size());
+                    namesChecker.erase(cId);
+                    programColumns.erase(cId);
+                }
+            }
+
+            const auto getDiffColumnsMessage = [&]() {
+                return TStringBuilder() << "ssa program has different columns with kqp request: kqp_columns=" << JoinSeq(",", namesChecker) << " vs program_columns=" << JoinSeq(",", programColumns);
+            };
+
+            if (namesChecker.size() != programColumns.size()) {
+                ErrorDescription = getDiffColumnsMessage();
+                return false;
+            }
+            for (auto&& i : namesChecker) {
+                if (!programColumns.contains(i)) {
+                    ErrorDescription = getDiffColumnsMessage();
+                    return false;
+                }
+            }
+        }
+
+        read.SetProgram(std::move(ssaProgram));
+
+        return true;
+    }
 }
 
 }