Use a better way to replace nodes in RewriteReadFromView function

VisitExpr is unconventional. + Better error printout in unit tests of views.
jepett0 · Jan 18, 2024 · 504dad8 · 504dad8
1 parent 03b2738
commit 504dad8
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 26 deletions.
diff --git a/ydb/core/kqp/provider/rewrite_io_utils.cpp b/ydb/core/kqp/provider/rewrite_io_utils.cpp
@@ -52,14 +52,6 @@ void AddChild(const TExprNode::TPtr& parent, const TExprNode::TPtr& newChild) {
     parent->ChangeChildrenInplace(std::move(childrenToChange));
 }
 
-void ChangeChild(const TExprNode::TPtr& parent, ui32 index, const TExprNode::TPtr& newChild) {
-    YQL_ENSURE(parent->ChildrenSize() > index);
-
-    auto childrenToChange = parent->ChildrenList();
-    childrenToChange[index] = newChild;
-    parent->ChangeChildrenInplace(std::move(childrenToChange));
-}
-
 TExprNode::TPtr FindSavedQueryGraph(const TExprNode::TPtr& carrier) {
     if (carrier->ChildrenSize() == 0) {
         return nullptr;
@@ -72,17 +64,18 @@ void SaveQueryGraph(const TExprNode::TPtr& carrier, TExprContext& ctx, const TEx
     AddChild(carrier, ctx.NewCallable(payload->Pos(), QueryGraphNodeSignature, {payload}));
 }
 
-void InsertExecutionOrderDependencies(const TExprNode::TPtr& queryGraph, const TExprNode::TPtr& worldBefore) {
-    VisitExpr(
-        queryGraph,
-        nullptr,
-        [&worldBefore](const TExprNode::TPtr& node) {
-            if (node->ChildrenSize() > 0 && node->Child(0)->IsWorld()) {
-                ChangeChild(node, 0, worldBefore);
-            }
-            return true;
-        }
-    );
+void InsertExecutionOrderDependencies(
+    TExprNode::TPtr& queryGraph,
+    const TExprNode::TPtr& worldBefore,
+    TExprContext& ctx
+) {
+    const auto initialWorldOfTheQuery = FindNode(queryGraph, [](const TExprNode::TPtr& node) {
+        return node->IsWorld();
+    });
+    if (!initialWorldOfTheQuery) {
+        return;
+    }
+    queryGraph = ctx.ReplaceNode(std::move(queryGraph), *initialWorldOfTheQuery, worldBefore);
 }
 
 bool CheckTopLevelness(const TExprNode::TPtr& candidateRead, const TExprNode::TPtr& queryGraph) {
@@ -146,7 +139,7 @@ TExprNode::TPtr RewriteReadFromView(
         YQL_CLOG(TRACE, ProviderKqp) << "Expression graph of the query stored in the view:\n"
                                      << NCommon::ExprToPrettyString(ctx, *queryGraph);
 
-        InsertExecutionOrderDependencies(queryGraph, worldBeforeThisRead);
+        InsertExecutionOrderDependencies(queryGraph, worldBeforeThisRead, ctx);
         SaveQueryGraph(readNode.Ptr(), ctx, queryGraph);
     }
 

diff --git a/ydb/core/kqp/ut/view/view_ut.cpp b/ydb/core/kqp/ut/view/view_ut.cpp
@@ -55,20 +55,18 @@ TString ReadWholeFile(const TString& path) {
 }
 
 void ExecuteDataDefinitionQuery(TSession& session, const TString& script) {
-    Cerr << "Executing the following DDL script:\n" << script;
-
     const auto result = session.ExecuteSchemeQuery(script).ExtractValueSync();
-    UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+    UNIT_ASSERT_C(result.IsSuccess(), "Failed to execute the following DDL script:\n"
+                                          << script << "\nThe issues:\n" << result.GetIssues().ToString());
 }
 
 TDataQueryResult ExecuteDataModificationQuery(TSession& session, const TString& script) {
-    Cerr << "Executing the following DML script:\n" << script;
-
     const auto result = session.ExecuteDataQuery(
             script,
             TTxControl::BeginTx(TTxSettings::SerializableRW()).CommitTx()
         ).ExtractValueSync();
-    UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+    UNIT_ASSERT_C(result.IsSuccess(), "Failed to execute the following DML script:\n"
+                                          << script << "\nThe issues:\n" << result.GetIssues().ToString());
 
     return result;
 }