[Stable-25-3-1] Fix upsert to table with unique index (#26508)

Author: nikvas0

ydb/core/kqp/opt/physical/effects/kqp_opt_phy_upsert_index.cpp

@@ -339,11 +339,7 @@ RewriteInputForConstraint(const TExprBase& inputRows, const THashSet<TStringBuf>
         hasUniqIndex |= (indexDesc->Type == TIndexDescription::EType::GlobalSyncUnique);
         for (const auto& indexKeyCol : indexDesc->KeyColumns) {
             if (inputColumns.contains(indexKeyCol)) {
-                if (!usedIndexes.contains(indexDesc->Name) &&
-                    std::find(mainPk.begin(), mainPk.end(), indexKeyCol) == mainPk.end())
-                {
-                    usedIndexes.insert(indexDesc->Name);
-                }
+                usedIndexes.insert(indexDesc->Name);
             } else {
                 // input always contains key columns
                 YQL_ENSURE(std::find(mainPk.begin(), mainPk.end(), indexKeyCol) == mainPk.end());
@@ -352,6 +348,8 @@ RewriteInputForConstraint(const TExprBase& inputRows, const THashSet<TStringBuf>
         }
     }
 
+    AFL_ENSURE(!hasUniqIndex || !usedIndexes.empty());
+
     if (!hasUniqIndex) {
         missedKeyInput.clear();
     }

ydb/core/kqp/ut/indexes/kqp_indexes_ut.cpp

@@ -5972,6 +5972,120 @@ R"([[#;#;["Primary1"];[41u]];[["Secondary2"];[2u];["Primary2"];[42u]];[["Seconda
             );
         }
     }
+
+    Y_UNIT_TEST_TWIN(IndexUpsert, Uniq) {
+        auto setting = NKikimrKqp::TKqpSetting();
+        auto serverSettings = TKikimrSettings().SetKqpSettings({setting});
+        TKikimrRunner kikimr(serverSettings);
+
+        auto client = kikimr.GetQueryClient();
+
+        {
+            const TString query(Q_(std::format(R"(
+                CREATE TABLE `/Root/TestTable` (
+                    a Int32,
+                    b Int32,
+                    PRIMARY KEY(a,b),
+                    INDEX ix_b GLOBAL {} SYNC ON (b)
+                );
+            )", Uniq ? "UNIQUE" : "")));
+
+            auto result = client.ExecuteQuery(
+                    query,
+                    NQuery::TTxControl::NoTx())
+                .ExtractValueSync();
+            UNIT_ASSERT(result.IsSuccess());
+        }
+
+        {
+            const TString query(Q_(R"(
+                $v=[<|a:10,b:20|>,<|a:30,b:20|>];
+                UPSERT INTO `/Root/TestTable` SELECT * FROM AS_TABLE($v);
+            )"));
+
+            auto result = client.ExecuteQuery(
+                    query,
+                    NQuery::TTxControl::NoTx())
+                .ExtractValueSync();
+            if (Uniq) {
+                UNIT_ASSERT_C(!result.IsSuccess(), result.GetIssues().ToString());
+                UNIT_ASSERT_STRING_CONTAINS_C(
+                    result.GetIssues().ToString(),
+                    "Duplicated keys found.",
+                    result.GetIssues().ToString());
+            } else {
+                UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+            }
+        }
+
+        {
+            const TString query(Q_(R"(
+                $v=[<|a:10,b:20|>,<|a:10,b:20|>];
+                UPSERT INTO `/Root/TestTable` SELECT * FROM AS_TABLE($v);
+            )"));
+
+            auto result = client.ExecuteQuery(
+                    query,
+                    NQuery::TTxControl::NoTx())
+                .ExtractValueSync();
+            if (Uniq) {
+                UNIT_ASSERT_C(!result.IsSuccess(), result.GetIssues().ToString());
+                UNIT_ASSERT_STRING_CONTAINS_C(
+                    result.GetIssues().ToString(),
+                    "Duplicated keys found.",
+                    result.GetIssues().ToString());
+            } else {
+                UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+            }
+        }
+
+        {
+            const TString query(Q_(R"(
+                $v=[<|a:10,b:10|>,<|a:30,b:30|>];
+                UPSERT INTO `/Root/TestTable` SELECT * FROM AS_TABLE($v);
+            )"));
+
+            auto result = client.ExecuteQuery(
+                    query,
+                    NQuery::TTxControl::NoTx())
+                .ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+
+        {
+            const TString query(Q_(R"(
+                $v=[<|a:20,b:10|>,<|a:20,b:30|>];
+                UPSERT INTO `/Root/TestTable` SELECT * FROM AS_TABLE($v);
+            )"));
+
+            auto result = client.ExecuteQuery(
+                    query,
+                    NQuery::TTxControl::NoTx())
+                .ExtractValueSync();
+            if (Uniq) {
+                UNIT_ASSERT_C(!result.IsSuccess(), result.GetIssues().ToString());
+                UNIT_ASSERT_STRING_CONTAINS_C(
+                    result.GetIssues().ToString(),
+                    "Conflict with existing key.",
+                    result.GetIssues().ToString());
+            } else {
+                UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+            }
+        }
+
+        {
+            const TString query(Q_(R"(
+                $v=[<|a:20,b:40|>,<|a:20,b:50|>];
+                UPSERT INTO `/Root/TestTable` SELECT * FROM AS_TABLE($v);
+            )"));
+
+            auto result = client.ExecuteQuery(
+                    query,
+                    NQuery::TTxControl::NoTx())
+                .ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+    }
 }
 
 }