Skip to content

Commit 0577192

Browse files
GrigoriyPAGrigoriyPA
authored
YQ-4060 fix multi statement select (#13946)
1 parent f61d285 commit 0577192

File tree

3 files changed

+195
-87
lines changed

3 files changed

+195
-87
lines changed

‎ydb/core/kqp/provider/yql_kikimr_opt_build.cpp

Lines changed: 57 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -309,6 +309,11 @@ bool IsDqRead(const TExprBase& node, TExprContext& ctx, TTypeAnnotationContext&
309309
TExprBase providerArg(node.Ref().Child(1));
310310
if (auto maybeDataSource = providerArg.Maybe<TCoDataSource>()) {
311311
TStringBuf dataSourceCategory = maybeDataSource.Cast().Category();
312+
if (dataSourceCategory == NYql::PgProviderName) {
313+
// All pg reads should be replaced on TPgTableContent
314+
return false;
315+
}
316+
312317
auto dataSourceProviderIt = types.DataSourceMap.find(dataSourceCategory);
313318
if (dataSourceProviderIt != types.DataSourceMap.end()) {
314319
if (auto* dqIntegration = dataSourceProviderIt->second->GetDqIntegration()) {
@@ -362,7 +367,7 @@ bool IsDqWrite(const TExprBase& node, TExprContext& ctx, TTypeAnnotationContext&
362367
return false;
363368
}
364369

365-
bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, TKiExploreTxResults& txRes,
370+
bool ExploreNode(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, TKiExploreTxResults& txRes,
366371
TIntrusivePtr<TKikimrTablesData> tablesData, TTypeAnnotationContext& types) {
367372

368373
if (txRes.Ops.cend() != txRes.Ops.find(node.Raw())) {
@@ -374,9 +379,9 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
374379
return true;
375380
}
376381

377-
if (auto maybeLeft = node.Maybe<TCoLeft>()) {
382+
if (node.Maybe<TCoLeft>()) {
378383
txRes.Ops.insert(node.Raw());
379-
return ExploreTx(maybeLeft.Cast().Input(), ctx, dataSink, txRes, tablesData, types);
384+
return true;
380385
}
381386

382387
auto checkDataSource = [dataSink] (const TKiDataSource& ds) {
@@ -400,27 +405,24 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
400405
YQL_ENSURE(key.GetKeyType() == TKikimrKey::Type::Table);
401406
auto table = key.GetTablePath();
402407
txRes.Ops.insert(node.Raw());
403-
auto result = ExploreTx(maybeRead.Cast().World(), ctx, dataSink, txRes, tablesData, types);
404408

405409
YQL_ENSURE(tablesData);
406410
const auto& tableData = tablesData->ExistingTable(cluster, table);
407411
YQL_ENSURE(tableData.Metadata);
408412
auto readColumns = read.GetSelectColumns(ctx, tableData);
409413
txRes.AddReadOpToQueryBlock(key, readColumns, tableData.Metadata);
410414
txRes.AddTableOperation(BuildTableOpNode(cluster, table, TYdbOperation::Select, read.Pos(), ctx));
411-
return result;
415+
return true;
412416
}
413417

414418
if (IsDqRead(node, ctx, types, true, &txRes.HasErrors)) {
415419
txRes.Ops.insert(node.Raw());
416-
TExprNode::TPtr worldChild = node.Raw()->ChildPtr(0);
417-
return ExploreTx(TExprBase(worldChild), ctx, dataSink, txRes, tablesData, types);
420+
return true;
418421
}
419422

420423
if (IsPgRead(node, types)) {
421424
txRes.Ops.insert(node.Raw());
422-
TExprNode::TPtr worldChild = node.Raw()->ChildPtr(0);
423-
return ExploreTx(TExprBase(worldChild), ctx, dataSink, txRes, tablesData, types);
425+
return true;
424426
}
425427

426428
if (auto maybeWrite = node.Maybe<TKiWriteTable>()) {
@@ -431,7 +433,6 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
431433

432434
auto table = write.Table().Value();
433435
txRes.Ops.insert(node.Raw());
434-
auto result = ExploreTx(write.World(), ctx, dataSink, txRes, tablesData, types);
435436
auto tableOp = GetTableOp(write);
436437

437438
YQL_ENSURE(tablesData);
@@ -472,14 +473,13 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
472473
}
473474

474475
txRes.AddTableOperation(BuildTableOpNode(cluster, table, tableOp, write.Pos(), ctx));
475-
return result;
476+
return true;
476477
}
477478

478479
if (IsDqWrite(node, ctx, types)) {
479480
txRes.Ops.insert(node.Raw());
480481
txRes.AddEffect(node, THashMap<TString, TPrimitiveYdbOperations>{});
481-
TExprNode::TPtr worldChild = node.Raw()->ChildPtr(0);
482-
return ExploreTx(TExprBase(worldChild), ctx, dataSink, txRes, tablesData, types);
482+
return true;
483483
}
484484

485485
if (auto maybeUpdate = node.Maybe<TKiUpdateTable>()) {
@@ -490,7 +490,6 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
490490

491491
auto table = update.Table().Value();
492492
txRes.Ops.insert(node.Raw());
493-
auto result = ExploreTx(update.World(), ctx, dataSink, txRes, tablesData, types);
494493
const auto tableOp = TYdbOperation::Update;
495494

496495
YQL_ENSURE(tablesData);
@@ -525,7 +524,7 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
525524
}
526525

527526
txRes.AddTableOperation(BuildTableOpNode(cluster, table, tableOp, update.Pos(), ctx));
528-
return result;
527+
return true;
529528
}
530529

531530
if (auto maybeDelete = node.Maybe<TKiDeleteTable>()) {
@@ -536,7 +535,6 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
536535

537536
auto table = del.Table().Value();
538537
txRes.Ops.insert(node.Raw());
539-
auto result = ExploreTx(del.World(), ctx, dataSink, txRes, tablesData, types);
540538
const auto tableOp = TYdbOperation::Delete;
541539

542540
YQL_ENSURE(tablesData);
@@ -564,7 +562,7 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
564562
}
565563

566564
txRes.AddTableOperation(BuildTableOpNode(cluster, table, tableOp, del.Pos(), ctx));
567-
return result;
565+
return true;
568566
}
569567

570568
if (auto maybeCreate = node.Maybe<TKiCreateTable>()) {
@@ -575,9 +573,8 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
575573

576574
auto table = create.Table().Value();
577575
txRes.Ops.insert(node.Raw());
578-
auto result = ExploreTx(create.World(), ctx, dataSink, txRes, tablesData, types);
579576
txRes.AddTableOperation(BuildTableOpNode(cluster, table, TYdbOperation::CreateTable, create.Pos(), ctx));
580-
return result;
577+
return true;
581578
}
582579

583580
if (auto maybeDrop = node.Maybe<TKiDropTable>()) {
@@ -588,9 +585,8 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
588585

589586
auto table = drop.Table().Value();
590587
txRes.Ops.insert(node.Raw());
591-
auto result = ExploreTx(drop.World(), ctx, dataSink, txRes, tablesData, types);
592588
txRes.AddTableOperation(BuildTableOpNode(cluster, table, TYdbOperation::DropTable, drop.Pos(), ctx));
593-
return result;
589+
return true;
594590
}
595591

596592
if (auto maybeAlter = node.Maybe<TKiAlterTable>()) {
@@ -601,9 +597,8 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
601597

602598
auto table = alter.Table().Value();
603599
txRes.Ops.insert(node.Raw());
604-
auto result = ExploreTx(alter.World(), ctx, dataSink, txRes, tablesData, types);
605600
txRes.AddTableOperation(BuildTableOpNode(cluster, table, TYdbOperation::AlterTable, alter.Pos(), ctx));
606-
return result;
601+
return true;
607602
}
608603

609604
if (auto maybeCreateUser = node.Maybe<TKiCreateUser>()) {
@@ -613,9 +608,8 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
613608
}
614609

615610
txRes.Ops.insert(node.Raw());
616-
auto result = ExploreTx(createUser.World(), ctx, dataSink, txRes, tablesData, types);
617611
txRes.AddTableOperation(BuildYdbOpNode(cluster, TYdbOperation::CreateUser, createUser.Pos(), ctx));
618-
return result;
612+
return true;
619613
}
620614

621615
if (auto maybeAlterUser = node.Maybe<TKiAlterUser>()) {
@@ -625,9 +619,8 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
625619
}
626620

627621
txRes.Ops.insert(node.Raw());
628-
auto result = ExploreTx(alterUser.World(), ctx, dataSink, txRes, tablesData, types);
629622
txRes.AddTableOperation(BuildYdbOpNode(cluster, TYdbOperation::AlterUser, alterUser.Pos(), ctx));
630-
return result;
623+
return true;
631624
}
632625

633626
if (auto maybeDropUser = node.Maybe<TKiDropUser>()) {
@@ -637,9 +630,8 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
637630
}
638631

639632
txRes.Ops.insert(node.Raw());
640-
auto result = ExploreTx(dropUser.World(), ctx, dataSink, txRes, tablesData, types);
641633
txRes.AddTableOperation(BuildYdbOpNode(cluster, TYdbOperation::DropUser, dropUser.Pos(), ctx));
642-
return result;
634+
return true;
643635
}
644636

645637
if (auto maybeCreateGroup = node.Maybe<TKiCreateGroup>()) {
@@ -649,9 +641,8 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
649641
}
650642

651643
txRes.Ops.insert(node.Raw());
652-
auto result = ExploreTx(createGroup.World(), ctx, dataSink, txRes, tablesData, types);
653644
txRes.AddTableOperation(BuildYdbOpNode(cluster, TYdbOperation::CreateGroup, createGroup.Pos(), ctx));
654-
return result;
645+
return true;
655646
}
656647

657648
if (auto maybeAlterGroup = node.Maybe<TKiAlterGroup>()) {
@@ -661,9 +652,8 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
661652
}
662653

663654
txRes.Ops.insert(node.Raw());
664-
auto result = ExploreTx(alterGroup.World(), ctx, dataSink, txRes, tablesData, types);
665655
txRes.AddTableOperation(BuildYdbOpNode(cluster, TYdbOperation::AlterGroup, alterGroup.Pos(), ctx));
666-
return result;
656+
return true;
667657
}
668658

669659
if (auto maybeRenameGroup = node.Maybe<TKiRenameGroup>()) {
@@ -673,9 +663,8 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
673663
}
674664

675665
txRes.Ops.insert(node.Raw());
676-
auto result = ExploreTx(renameGroup.World(), ctx, dataSink, txRes, tablesData, types);
677666
txRes.AddTableOperation(BuildYdbOpNode(cluster, TYdbOperation::RenameGroup, renameGroup.Pos(), ctx));
678-
return result;
667+
return true;
679668
}
680669

681670
if (auto maybeDropGroup = node.Maybe<TKiDropGroup>()) {
@@ -685,9 +674,8 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
685674
}
686675

687676
txRes.Ops.insert(node.Raw());
688-
auto result = ExploreTx(dropGroup.World(), ctx, dataSink, txRes, tablesData, types);
689677
txRes.AddTableOperation(BuildYdbOpNode(cluster, TYdbOperation::DropGroup, dropGroup.Pos(), ctx));
690-
return result;
678+
return true;
691679
}
692680

693681
if (auto maybeExecQuery = node.Maybe<TKiExecDataQuery>()) {
@@ -700,39 +688,29 @@ bool ExploreTx(TExprBase node, TExprContext& ctx, const TKiDataSink& dataSink, T
700688
return true;
701689
}
702690

703-
if (auto maybeCommit = node.Maybe<TCoCommit>()) {
704-
auto commit = maybeCommit.Cast();
705-
706-
if (commit.DataSink().Maybe<TKiDataSink>() && checkDataSink(commit.DataSink().Cast<TKiDataSink>())) {
707-
txRes.Sync.push_back(commit);
708-
return true;
709-
}
710-
711-
return ExploreTx(commit.World(), ctx, dataSink, txRes, tablesData, types);
691+
if (node.Maybe<TCoCommit>()) {
692+
return true;
712693
}
713694

714-
if (auto maybeSync = node.Maybe<TCoSync>()) {
695+
if (node.Maybe<TCoSync>()) {
715696
txRes.Ops.insert(node.Raw());
716-
for (auto child : maybeSync.Cast()) {
717-
if (!ExploreTx(child, ctx, dataSink, txRes, tablesData, types)) {
718-
return false;
719-
}
720-
}
721697
return true;
722698
}
723699

724700
if (node.Maybe<TResWrite>() ||
725701
node.Maybe<TResPull>())
726702
{
727703
txRes.Ops.insert(node.Raw());
728-
bool result = ExploreTx(TExprBase(node.Ref().ChildPtr(0)), ctx, dataSink, txRes, tablesData, types);
729-
// Cerr << KqpExprToPrettyString(*node.Raw(), ctx) << Endl;
730704
txRes.AddResult(node);
731-
return result;
705+
return true;
732706
}
733707

734708
if (node.Ref().IsCallable(ConfigureName)) {
735-
txRes.Sync.push_back(node);
709+
return true;
710+
}
711+
712+
if (node.Maybe<TCoCons>()) {
713+
txRes.Ops.insert(node.Raw());
736714
return true;
737715
}
738716

@@ -757,32 +735,42 @@ bool IsKikimrPureNode(const TExprNode::TPtr& node) {
757735
return true;
758736
}
759737

760-
bool CheckTx(TExprBase txStart, const TKiDataSink& dataSink, const THashSet<const TExprNode*>& txOps,
761-
const THashSet<const TExprNode*>& txSync)
738+
bool ExploreTx(TExprBase root, TExprContext& ctx, const TKiDataSink& dataSink, TKiExploreTxResults& txRes,
739+
TIntrusivePtr<TKikimrTablesData> tablesData, TTypeAnnotationContext& types)
762740
{
763-
bool hasErrors = false;
764-
VisitExpr(txStart.Ptr(), [&txOps, &txSync, &hasErrors, dataSink] (const TExprNode::TPtr& node) {
765-
if (hasErrors) {
766-
return false;
741+
const auto preFunc = [&dataSink, &txRes](const TExprNode::TPtr& node) {
742+
if (const auto maybeCommit = TExprBase(node).Maybe<TCoCommit>()) {
743+
const auto commit = maybeCommit.Cast();
744+
if (commit.DataSink().Maybe<TKiDataSink>() && commit.DataSink().Cast<TKiDataSink>().Raw() == dataSink.Raw()) {
745+
txRes.Sync.push_back(commit);
746+
return false;
747+
}
748+
return true;
767749
}
768750

769-
if (txSync.find(node.Get()) != txSync.cend()) {
751+
if (node->IsCallable(ConfigureName)) {
752+
txRes.Sync.push_back(TExprBase(node));
770753
return false;
771754
}
772755

773-
if (auto maybeCommit = TMaybeNode<TCoCommit>(node)) {
774-
if (maybeCommit.Cast().DataSink().Raw() != dataSink.Raw()) {
775-
return true;
776-
}
756+
return true;
757+
};
758+
759+
bool hasErrors = false;
760+
const auto postFunc = [&hasErrors, &ctx, &dataSink, &txRes, tablesData, &types](const TExprNode::TPtr& node) {
761+
if (hasErrors) {
762+
return false;
777763
}
778764

779-
if (!IsKikimrPureNode(node) && txOps.find(node.Get()) == txOps.cend()) {
765+
if (!ExploreNode(TExprBase(node), ctx, dataSink, txRes, tablesData, types) && !IsKikimrPureNode(node)) {
780766
hasErrors = true;
781767
return false;
782768
}
783769

784770
return true;
785-
});
771+
};
772+
773+
VisitExpr(root.Ptr(), preFunc, postFunc);
786774

787775
return !hasErrors;
788776
}
@@ -964,11 +952,6 @@ TExprNode::TPtr KiBuildQuery(TExprBase node, TExprContext& ctx, TStringBuf datab
964952
return node.Ptr();
965953
}
966954

967-
auto txSyncSet = txExplore.GetSyncSet();
968-
if (!CheckTx(commit.World(), kiDataSink, txExplore.Ops, txSyncSet)) {
969-
return node.Ptr();
970-
}
971-
972955
bool hasScheme;
973956
bool hasData;
974957
txExplore.GetTableOperations(hasScheme, hasData);

0 commit comments

Comments
 (0)