Skip to content

Commit dff4634

Browse files
gridnevvvitgridnevvvit
authored
Fix returning if input is a parameter (#28414)
1 parent 06b1da6 commit dff4634

File tree

2 files changed

+307
-2
lines changed

2 files changed

+307
-2
lines changed

‎ydb/core/kqp/opt/physical/effects/kqp_opt_phy_returning.cpp‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,7 @@ TExprBase KqpBuildReturning(TExprBase node, TExprContext& ctx, const TTypeAnnota
173173
return TExprBase(ctx.ChangeChild(*returning.Raw(), TKqlReturningList::idx_Update, inputExpr.Ptr()));
174174
};
175175

176+
176177
if (auto maybeList = returning.Update().Maybe<TExprList>()) {
177178
for (auto item : maybeList.Cast()) {
178179
if (auto upsert = item.Maybe<TKqlUpsertRows>()) {
@@ -224,7 +225,7 @@ TExprBase KqpRewriteReturningUpsert(TExprBase node, TExprContext& ctx, const TKq
224225
return node;
225226
}
226227

227-
if (upsert.Input().Maybe<TDqPrecompute>() || upsert.Input().Maybe<TDqPhyPrecompute>()) {
228+
if (upsert.Input().Maybe<TDqPrecompute>() || upsert.Input().Maybe<TDqPhyPrecompute>() || upsert.Input().Maybe<TCoParameter>()) {
228229
return node;
229230
}
230231

@@ -247,7 +248,7 @@ TExprBase KqpRewriteReturningDelete(TExprBase node, TExprContext& ctx, const TKq
247248
return node;
248249
}
249250

250-
if (del.Input().Maybe<TDqPrecompute>() || del.Input().Maybe<TDqPhyPrecompute>()) {
251+
if (del.Input().Maybe<TDqPrecompute>() || del.Input().Maybe<TDqPhyPrecompute>() || del.Input().Maybe<TCoParameter>()) {
251252
return node;
252253
}
253254

‎ydb/core/kqp/ut/opt/kqp_returning_ut.cpp‎

Lines changed: 304 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -377,6 +377,25 @@ TString ExecuteReturningQuery(TKikimrRunner& kikimr, bool queryService, TString
377377
return FormatResultSetYson(result.GetResultSet(0));
378378
}
379379

380+
TString ExecuteReturningQueryWithParams(TKikimrRunner& kikimr, bool queryService, TString query, const TParams& params) {
381+
if (queryService) {
382+
auto qdb = kikimr.GetQueryClient();
383+
auto qSession = qdb.GetSession().GetValueSync().GetSession();
384+
auto settings = NYdb::NQuery::TExecuteQuerySettings()
385+
.Syntax(NYdb::NQuery::ESyntax::YqlV1);
386+
auto result = qSession.ExecuteQuery(
387+
query, NYdb::NQuery::TTxControl::BeginTx().CommitTx(), params, settings).ExtractValueSync();
388+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
389+
return FormatResultSetYson(result.GetResultSet(0));
390+
}
391+
392+
auto db = kikimr.GetTableClient();
393+
auto session = db.CreateSession().GetValueSync().GetSession();
394+
auto result = session.ExecuteDataQuery(query, TTxControl::BeginTx().CommitTx(), params).ExtractValueSync();
395+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
396+
return FormatResultSetYson(result.GetResultSet(0));
397+
}
398+
380399
Y_UNIT_TEST_TWIN(ReturningWorks, QueryService) {
381400
auto kikimr = DefaultKikimrRunner();
382401
auto db = kikimr.GetTableClient();
@@ -663,6 +682,291 @@ Y_UNIT_TEST(ReturningTypes) {
663682
}
664683
}
665684

685+
Y_UNIT_TEST_TWIN(ReturningUpsertAsTableListNotNullOnly, QueryService) {
686+
// Test for issue #27021: Query fails when using RETURNING CLAUSE with UPSERT
687+
// to table with only NOT NULL fields and query parameters of type List
688+
auto kikimr = DefaultKikimrRunner();
689+
690+
auto client = kikimr.GetTableClient();
691+
auto session = client.CreateSession().GetValueSync().GetSession();
692+
693+
// Create table with only NOT NULL fields
694+
const auto queryCreate = Q_(R"(
695+
--!syntax_v1
696+
CREATE TABLE test_table (
697+
id Uint64 NOT NULL,
698+
value Utf8 NOT NULL,
699+
PRIMARY KEY (id)
700+
);
701+
)");
702+
703+
auto resultCreate = session.ExecuteSchemeQuery(queryCreate).GetValueSync();
704+
UNIT_ASSERT_C(resultCreate.IsSuccess(), resultCreate.GetIssues().ToString());
705+
706+
{
707+
// Test case from issue #27021
708+
const auto query = Q_(R"(
709+
--!syntax_v1
710+
DECLARE $data AS List<Struct<id: UInt64, value: Utf8>>;
711+
712+
UPSERT INTO test_table
713+
SELECT * FROM AS_TABLE($data)
714+
RETURNING *;
715+
)");
716+
717+
auto paramsBuilder = TParamsBuilder();
718+
auto& dataParam = paramsBuilder.AddParam("$data");
719+
720+
dataParam.BeginList();
721+
dataParam.AddListItem()
722+
.BeginStruct()
723+
.AddMember("id")
724+
.Uint64(1)
725+
.AddMember("value")
726+
.Utf8("test1")
727+
.EndStruct();
728+
dataParam.AddListItem()
729+
.BeginStruct()
730+
.AddMember("id")
731+
.Uint64(2)
732+
.AddMember("value")
733+
.Utf8("test2")
734+
.EndStruct();
735+
dataParam.EndList();
736+
dataParam.Build();
737+
738+
auto params = paramsBuilder.Build();
739+
740+
// This should succeed, but currently fails with infinite loop error
741+
CompareYson(R"([[1u;"test1"];[2u;"test2"]])",
742+
ExecuteReturningQueryWithParams(kikimr, QueryService, query, params));
743+
}
744+
745+
{
746+
// Test with explicit field names in SELECT clause
747+
const auto query = Q_(R"(
748+
--!syntax_v1
749+
DECLARE $data AS List<Struct<id: UInt64, value: Utf8>>;
750+
751+
UPSERT INTO test_table
752+
SELECT id, value FROM AS_TABLE($data)
753+
RETURNING *;
754+
)");
755+
756+
auto paramsBuilder = TParamsBuilder();
757+
auto& dataParam = paramsBuilder.AddParam("$data");
758+
759+
dataParam.BeginList();
760+
dataParam.AddListItem()
761+
.BeginStruct()
762+
.AddMember("id")
763+
.Uint64(3)
764+
.AddMember("value")
765+
.Utf8("test3")
766+
.EndStruct();
767+
dataParam.EndList();
768+
dataParam.Build();
769+
770+
auto params = paramsBuilder.Build();
771+
772+
CompareYson(R"([[3u;"test3"]])",
773+
ExecuteReturningQueryWithParams(kikimr, QueryService, query, params));
774+
}
775+
776+
{
777+
// Test DELETE with RETURNING using AS_TABLE with List parameter (same issue #27021)
778+
// First insert some data without RETURNING
779+
const auto insertQuery = Q_(R"(
780+
--!syntax_v1
781+
DECLARE $data AS List<Struct<id: UInt64, value: Utf8>>;
782+
783+
UPSERT INTO test_table
784+
SELECT * FROM AS_TABLE($data);
785+
)");
786+
787+
auto insertParamsBuilder = TParamsBuilder();
788+
auto& insertDataParam = insertParamsBuilder.AddParam("$data");
789+
790+
insertDataParam.BeginList();
791+
insertDataParam.AddListItem()
792+
.BeginStruct()
793+
.AddMember("id")
794+
.Uint64(10)
795+
.AddMember("value")
796+
.Utf8("delete1")
797+
.EndStruct();
798+
insertDataParam.EndList();
799+
insertDataParam.Build();
800+
801+
auto insertParams = insertParamsBuilder.Build();
802+
803+
// Execute insert without RETURNING
804+
if (QueryService) {
805+
auto qdb = kikimr.GetQueryClient();
806+
auto qSession = qdb.GetSession().GetValueSync().GetSession();
807+
auto settings = NYdb::NQuery::TExecuteQuerySettings()
808+
.Syntax(NYdb::NQuery::ESyntax::YqlV1);
809+
auto insertResult = qSession.ExecuteQuery(
810+
insertQuery, NYdb::NQuery::TTxControl::BeginTx().CommitTx(), insertParams, settings).ExtractValueSync();
811+
UNIT_ASSERT_VALUES_EQUAL_C(insertResult.GetStatus(), EStatus::SUCCESS, insertResult.GetIssues().ToString());
812+
} else {
813+
auto db = kikimr.GetTableClient();
814+
auto session = db.CreateSession().GetValueSync().GetSession();
815+
auto insertResult = session.ExecuteDataQuery(insertQuery, TTxControl::BeginTx().CommitTx(), insertParams).ExtractValueSync();
816+
UNIT_ASSERT_VALUES_EQUAL_C(insertResult.GetStatus(), EStatus::SUCCESS, insertResult.GetIssues().ToString());
817+
}
818+
819+
// Now DELETE with RETURNING using AS_TABLE
820+
const auto deleteQuery = Q_(R"(
821+
--!syntax_v1
822+
DECLARE $data AS List<Struct<id: UInt64>>;
823+
824+
DELETE FROM test_table ON
825+
SELECT * FROM AS_TABLE($data)
826+
RETURNING *;
827+
)");
828+
829+
auto deleteParamsBuilder = TParamsBuilder();
830+
auto& deleteDataParam = deleteParamsBuilder.AddParam("$data");
831+
832+
deleteDataParam.BeginList();
833+
deleteDataParam.AddListItem()
834+
.BeginStruct()
835+
.AddMember("id")
836+
.Uint64(10)
837+
.EndStruct();
838+
deleteDataParam.EndList();
839+
deleteDataParam.Build();
840+
841+
auto deleteParams = deleteParamsBuilder.Build();
842+
843+
// This should succeed, but currently fails with infinite loop error
844+
CompareYson(R"([[10u;"delete1"]])",
845+
ExecuteReturningQueryWithParams(kikimr, QueryService, deleteQuery, deleteParams));
846+
}
847+
}
848+
849+
Y_UNIT_TEST_TWIN(ReturningUpsertAsTableListWithNullable, QueryService) {
850+
// Test that nullable columns work correctly (this should work even before the fix)
851+
auto kikimr = DefaultKikimrRunner();
852+
853+
auto client = kikimr.GetTableClient();
854+
auto session = client.CreateSession().GetValueSync().GetSession();
855+
856+
const auto queryCreate = Q_(R"(
857+
--!syntax_v1
858+
CREATE TABLE test_table_nullable (
859+
id Uint64 NOT NULL,
860+
value Utf8,
861+
PRIMARY KEY (id)
862+
);
863+
)");
864+
865+
auto resultCreate = session.ExecuteSchemeQuery(queryCreate).GetValueSync();
866+
UNIT_ASSERT_C(resultCreate.IsSuccess(), resultCreate.GetIssues().ToString());
867+
868+
{
869+
const auto query = Q_(R"(
870+
--!syntax_v1
871+
DECLARE $data AS List<Struct<id: UInt64, value: Utf8?>>;
872+
873+
UPSERT INTO test_table_nullable
874+
SELECT * FROM AS_TABLE($data)
875+
RETURNING *;
876+
)");
877+
878+
auto paramsBuilder = TParamsBuilder();
879+
auto& dataParam = paramsBuilder.AddParam("$data");
880+
881+
dataParam.BeginList();
882+
dataParam.AddListItem()
883+
.BeginStruct()
884+
.AddMember("id")
885+
.Uint64(1)
886+
.AddMember("value")
887+
.OptionalUtf8("test1")
888+
.EndStruct();
889+
dataParam.EndList();
890+
dataParam.Build();
891+
892+
auto params = paramsBuilder.Build();
893+
894+
CompareYson(R"([[1u;["test1"]]])",
895+
ExecuteReturningQueryWithParams(kikimr, QueryService, query, params));
896+
}
897+
898+
{
899+
// Test DELETE with RETURNING using AS_TABLE with List parameter (with nullable columns)
900+
// First insert some data without RETURNING
901+
const auto insertQuery = Q_(R"(
902+
--!syntax_v1
903+
DECLARE $data AS List<Struct<id: UInt64, value: Utf8?>>;
904+
905+
UPSERT INTO test_table_nullable
906+
SELECT * FROM AS_TABLE($data);
907+
)");
908+
909+
auto insertParamsBuilder = TParamsBuilder();
910+
auto& insertDataParam = insertParamsBuilder.AddParam("$data");
911+
912+
insertDataParam.BeginList();
913+
insertDataParam.AddListItem()
914+
.BeginStruct()
915+
.AddMember("id")
916+
.Uint64(20)
917+
.AddMember("value")
918+
.OptionalUtf8("delete_nullable1")
919+
.EndStruct();
920+
insertDataParam.EndList();
921+
insertDataParam.Build();
922+
923+
auto insertParams = insertParamsBuilder.Build();
924+
925+
// Execute insert without RETURNING
926+
if (QueryService) {
927+
auto qdb = kikimr.GetQueryClient();
928+
auto qSession = qdb.GetSession().GetValueSync().GetSession();
929+
auto settings = NYdb::NQuery::TExecuteQuerySettings()
930+
.Syntax(NYdb::NQuery::ESyntax::YqlV1);
931+
auto insertResult = qSession.ExecuteQuery(
932+
insertQuery, NYdb::NQuery::TTxControl::BeginTx().CommitTx(), insertParams, settings).ExtractValueSync();
933+
UNIT_ASSERT_VALUES_EQUAL_C(insertResult.GetStatus(), EStatus::SUCCESS, insertResult.GetIssues().ToString());
934+
} else {
935+
auto db = kikimr.GetTableClient();
936+
auto session = db.CreateSession().GetValueSync().GetSession();
937+
auto insertResult = session.ExecuteDataQuery(insertQuery, TTxControl::BeginTx().CommitTx(), insertParams).ExtractValueSync();
938+
UNIT_ASSERT_VALUES_EQUAL_C(insertResult.GetStatus(), EStatus::SUCCESS, insertResult.GetIssues().ToString());
939+
}
940+
941+
// Now DELETE with RETURNING using AS_TABLE
942+
const auto deleteQuery = Q_(R"(
943+
--!syntax_v1
944+
DECLARE $data AS List<Struct<id: UInt64>>;
945+
946+
DELETE FROM test_table_nullable ON
947+
SELECT * FROM AS_TABLE($data)
948+
RETURNING *;
949+
)");
950+
951+
auto deleteParamsBuilder = TParamsBuilder();
952+
auto& deleteDataParam = deleteParamsBuilder.AddParam("$data");
953+
954+
deleteDataParam.BeginList();
955+
deleteDataParam.AddListItem()
956+
.BeginStruct()
957+
.AddMember("id")
958+
.Uint64(20)
959+
.EndStruct();
960+
deleteDataParam.EndList();
961+
deleteDataParam.Build();
962+
963+
auto deleteParams = deleteParamsBuilder.Build();
964+
965+
CompareYson(R"([[20u;["delete_nullable1"]]])",
966+
ExecuteReturningQueryWithParams(kikimr, QueryService, deleteQuery, deleteParams));
967+
}
968+
}
969+
666970
}
667971

668972
} // namespace NKikimr::NKqp

0 commit comments

Comments
 (0)