KIKIMR-20115: add MKQLTerminate to kqp QueryData

uzhastik · Nov 29, 2023 · dcda303 · dcda303
1 parent 0fa5e01
commit dcda303
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 1 deletion.
diff --git a/ydb/core/kqp/query_data/kqp_query_data.cpp b/ydb/core/kqp/query_data/kqp_query_data.cpp
@@ -309,6 +309,7 @@ bool TQueryData::AddUVParam(const TString& name, NKikimr::NMiniKQL::TType* type,
 
 bool TQueryData::AddTypedValueParam(const TString& name, const Ydb::TypedValue& param) {
     auto guard = TypeEnv().BindAllocator();
+    const TBindTerminator bind(this);
     auto [typeFromProto, value] = ImportValueFromProto(
         param.type(), param.value(), TypeEnv(), AllocState->HolderFactory);
     return AddUVParam(name, typeFromProto, value);
@@ -492,4 +493,16 @@ void TQueryData::Clear() {
     }
 }
 
+void TQueryData::Terminate(const char* message) const {
+    TStringBuf reason = (message ? TStringBuf(message) : TStringBuf("(unknown)"));
+    TString fullMessage = TStringBuilder() <<
+        "Terminate was called, reason(" << reason.size() << "): " << reason << Endl;
+    AllocState->HolderFactory.CleanupModulesOnTerminate();
+    if (std::current_exception()) {
+        throw;
+    }
+
+    ythrow yexception() << fullMessage;
+}
+
 } // namespace NKikimr::NKqp
diff --git a/ydb/core/kqp/query_data/kqp_query_data.h b/ydb/core/kqp/query_data/kqp_query_data.h
@@ -176,7 +176,7 @@ class TTxAllocatorState: public TTimeAndRandomProvider {
     std::pair<NKikimr::NMiniKQL::TType*, NUdf::TUnboxedValue> GetInternalBindingValue(const NKqpProto::TKqpPhyParamBinding& paramBinding);
 };
 
-class TQueryData {
+class TQueryData : NMiniKQL::ITerminator {
 private:
     using TTypedUnboxedValue = std::pair<NKikimr::NMiniKQL::TType*, NUdf::TUnboxedValue>;
     using TNamedUnboxedValue = std::pair<const TString, TTypedUnboxedValue>;
@@ -277,6 +277,8 @@ class TQueryData {
             return false;
         };
     }
+
+    void Terminate(const char* message) const final;
 };
 
 

diff --git a/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp b/ydb/core/kqp/ut/pg/kqp_pg_ut.cpp
@@ -3062,6 +3062,46 @@ Y_UNIT_TEST_SUITE(KqpPg) {
             }
         }
     }
+
+    Y_UNIT_TEST(MkqlTerminate) {
+        NKikimrConfig::TAppConfig appConfig;
+        appConfig.MutableTableServiceConfig()->SetEnablePreparedDdl(true);
+        auto setting = NKikimrKqp::TKqpSetting();
+        auto serverSettings = TKikimrSettings()
+            .SetAppConfig(appConfig)
+            .SetKqpSettings({setting});
+        TKikimrRunner kikimr(serverSettings.SetWithSampleTables(false));
+        auto db = kikimr.GetQueryClient();
+        auto settings = NYdb::NQuery::TExecuteQuerySettings().Syntax(NYdb::NQuery::ESyntax::Pg);
+         {
+            auto result = db.ExecuteQuery(R"(
+                CREATE TABLE t (id INT PRIMARY KEY, data1 UUID[], data2 UUID[][]);
+            )", NYdb::NQuery::TTxControl::NoTx(), settings).ExtractValueSync();
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+        {
+            TString binval;
+            for (int i = 0; i < 4; i++) {
+                binval.push_back(0);
+            }
+            auto p1Value = TPgValue(TPgValue::VK_TEXT, "1", TPgType("pgunknown"));
+            auto p2Value = TPgValue(TPgValue::VK_BINARY, binval, TPgType("pgunknown"));
+            auto params = TParamsBuilder()
+                .AddParam("$p1")
+                    .Pg(p1Value)
+                    .Build()
+                .AddParam("$p2")
+                    .Pg(p2Value)
+                    .Build()
+                .Build();
+
+            auto result = db.ExecuteQuery(R"(
+                INSERT INTO t (id, data2) VALUES ($1, $2);
+            )", NYdb::NQuery::TTxControl::NoTx(), params, settings).ExtractValueSync();
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::BAD_REQUEST, result.GetIssues().ToString());
+            UNIT_ASSERT(result.GetIssues().ToString().Contains("invalid byte sequence for encoding \"UTF8\": 0x00"));
+        }
+    }
 }
 
 } // namespace NKqp