Skip to content

Commit 27887b8

Browse files
VPolkaVPolka
authored
Merge 1e0bec3 into 3caf512 
2 parents 3caf512 + 1e0bec3 commit 27887b8

File tree

6 files changed

+127
-21
lines changed

6 files changed

+127
-21
lines changed

‎ydb/core/kqp/common/compilation/events.h‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ struct TEvRecompileRequest: public TEventLocal<TEvRecompileRequest, TKqpEvents::
7676
const TMaybe<TKqpQueryId>& query, bool isQueryActionPrepare, TInstant deadline,
7777
TKqpDbCountersPtr dbCounters, const TGUCSettings::TPtr& gUCSettings, const TMaybe<TString>& applicationName,
7878
std::shared_ptr<std::atomic<bool>> intrestedInResult, const TIntrusivePtr<TUserRequestContext>& userRequestContext,
79-
NLWTrace::TOrbit orbit = {}, TKqpTempTablesState::TConstPtr tempTablesState = nullptr)
79+
NLWTrace::TOrbit orbit = {}, TKqpTempTablesState::TConstPtr tempTablesState = nullptr, TMaybe<TQueryAst> queryAst = Nothing())
8080
: UserToken(userToken)
8181
, Uid(uid)
8282
, Query(query)
@@ -89,6 +89,7 @@ struct TEvRecompileRequest: public TEventLocal<TEvRecompileRequest, TKqpEvents::
8989
, Orbit(std::move(orbit))
9090
, TempTablesState(std::move(tempTablesState))
9191
, IntrestedInResult(std::move(intrestedInResult))
92+
, QueryAst(queryAst)
9293
{
9394
}
9495

@@ -107,6 +108,8 @@ struct TEvRecompileRequest: public TEventLocal<TEvRecompileRequest, TKqpEvents::
107108

108109
TKqpTempTablesState::TConstPtr TempTablesState;
109110
std::shared_ptr<std::atomic<bool>> IntrestedInResult;
111+
112+
TMaybe<TQueryAst> QueryAst;
110113
};
111114

112115
struct TEvCompileResponse: public TEventLocal<TEvCompileResponse, TKqpEvents::EvCompileResponse> {

‎ydb/core/kqp/compile_service/kqp_compile_service.cpp‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -731,13 +731,17 @@ class TKqpCompileService : public TActorBootstrapped<TKqpCompileService> {
731731
NWilson::TSpan compileServiceSpan(TWilsonKqp::CompileService, ev->Get() ? std::move(ev->TraceId) : NWilson::TTraceId(), "CompileService");
732732

733733
TKqpCompileSettings compileSettings(true, request.IsQueryActionPrepare, false, request.Deadline, TableServiceConfig.GetEnableAstCache() ? ECompileActorAction::PARSE : ECompileActorAction::COMPILE);
734-
TKqpCompileRequest compileRequest(ev->Sender, request.Uid, compileResult ? *compileResult->Query : *request.Query,
734+
TKqpCompileRequest compileRequest(ev->Sender, request.Uid, request.Query ? *request.Query : *compileResult->Query,
735735
compileSettings, request.UserToken, dbCounters, request.GUCSettings, request.ApplicationName,
736736
ev->Cookie, std::move(ev->Get()->IntrestedInResult),
737737
ev->Get()->UserRequestContext,
738738
ev->Get() ? std::move(ev->Get()->Orbit) : NLWTrace::TOrbit(),
739739
std::move(compileServiceSpan), std::move(ev->Get()->TempTablesState));
740740

741+
if (TableServiceConfig.GetEnableAstCache() && request.QueryAst) {
742+
return CompileByAst(*request.QueryAst, compileRequest, ctx);
743+
}
744+
741745
if (!RequestsQueue.Enqueue(std::move(compileRequest))) {
742746
Counters->ReportCompileRequestRejected(dbCounters);
743747

‎ydb/core/kqp/session_actor/kqp_query_state.cpp‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -244,8 +244,15 @@ std::unique_ptr<TEvKqp::TEvRecompileRequest> TKqpQueryState::BuildReCompileReque
244244
compileDeadline = Min(compileDeadline, QueryDeadlines.CancelAt);
245245
}
246246

247-
return std::make_unique<TEvKqp::TEvRecompileRequest>(UserToken, CompileResult->Uid, CompileResult->Query, isQueryActionPrepare,
248-
compileDeadline, DbCounters, gUCSettingsPtr, ApplicationName, std::move(cookie), UserRequestContext, std::move(Orbit), TempTablesState);
247+
TMaybe<TQueryAst> statementAst;
248+
if (!Statements.empty()) {
249+
YQL_ENSURE(CurrentStatementId < Statements.size());
250+
statementAst = Statements[CurrentStatementId];
251+
}
252+
253+
return std::make_unique<TEvKqp::TEvRecompileRequest>(UserToken, CompileResult->Uid, query, isQueryActionPrepare,
254+
compileDeadline, DbCounters, gUCSettingsPtr, ApplicationName, std::move(cookie), UserRequestContext, std::move(Orbit), TempTablesState,
255+
statementAst);
249256
}
250257

251258
std::unique_ptr<TEvKqp::TEvCompileRequest> TKqpQueryState::BuildSplitRequest(std::shared_ptr<std::atomic<bool>> cookie, const TGUCSettings::TPtr& gUCSettingsPtr) {
@@ -295,7 +302,6 @@ bool TKqpQueryState::PrepareNextStatementPart() {
295302
QueryData = {};
296303
PreparedQuery = {};
297304
CompileResult = {};
298-
TxCtx = {};
299305
CurrentTx = 0;
300306
TableVersions = {};
301307
MaxReadType = ETableReadType::Other;

‎ydb/core/kqp/session_actor/kqp_query_state.h‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,7 @@ class TKqpQueryState : public TNonCopyable {
113113
TTxId TxId; // User tx
114114
bool Commit = false;
115115
bool Commited = false;
116+
bool Begin = false;
116117

117118
NTopic::TTopicOperations TopicOperations;
118119
TDuration CpuTime;
@@ -408,15 +409,12 @@ class TKqpQueryState : public TNonCopyable {
408409
}
409410

410411
void PrepareCurrentStatement() {
411-
QueryData = {};
412+
QueryData = std::make_shared<TQueryData>(TxCtx->TxAlloc);
412413
PreparedQuery = {};
413414
CompileResult = {};
414-
TxCtx = {};
415415
CurrentTx = 0;
416416
TableVersions = {};
417417
MaxReadType = ETableReadType::Other;
418-
Commit = false;
419-
Commited = false;
420418
TopicOperations = {};
421419
ReplayMessage = {};
422420
}
@@ -435,7 +433,6 @@ class TKqpQueryState : public TNonCopyable {
435433
TxCtx->EffectiveIsolationLevel = NKikimrKqp::ISOLATION_LEVEL_UNDEFINED;
436434
break;
437435
default:
438-
Commit = true;
439436
TxCtx->EffectiveIsolationLevel = NKikimrKqp::ISOLATION_LEVEL_SERIALIZABLE;
440437
}
441438
}

‎ydb/core/kqp/session_actor/kqp_session_actor.cpp‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -678,6 +678,7 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
678678
}
679679

680680
void BeginTx(const Ydb::Table::TransactionSettings& settings) {
681+
QueryState->Begin = true;
681682
QueryState->TxId = UlidGen.Next();
682683
QueryState->TxCtx = MakeIntrusive<TKqpTransactionContext>(false, AppData()->FunctionRegistry,
683684
AppData()->TimeProvider, AppData()->RandomProvider, Config->EnableKqpImmediateEffects);
@@ -731,7 +732,7 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
731732
if (hasTxControl || QueryState->HasImpliedTx()) {
732733
const auto& txControl = hasTxControl ? QueryState->GetTxControl() : GetImpliedTxControl();
733734

734-
QueryState->Commit = txControl.commit_tx();
735+
QueryState->Commit = txControl.commit_tx() && QueryState->ProcessingLastStatement();
735736
switch (txControl.tx_selector_case()) {
736737
case Ydb::Table::TransactionControl::kTxId: {
737738
auto txId = TTxId::FromString(txControl.tx_id());
@@ -746,15 +747,17 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
746747
break;
747748
}
748749
case Ydb::Table::TransactionControl::kBeginTx: {
749-
BeginTx(txControl.begin_tx());
750+
if (!QueryState->Begin) {
751+
BeginTx(txControl.begin_tx());
752+
}
750753
break;
751-
}
752-
case Ydb::Table::TransactionControl::TX_SELECTOR_NOT_SET:
754+
}
755+
case Ydb::Table::TransactionControl::TX_SELECTOR_NOT_SET:
753756
ythrow TRequestFail(Ydb::StatusIds::BAD_REQUEST)
754757
<< "wrong TxControl: tx_selector must be set";
755758
break;
756759
}
757-
} else {
760+
} else if (QueryState->CurrentStatementId == 0) {
758761
QueryState->TxCtx = MakeIntrusive<TKqpTransactionContext>(false, AppData()->FunctionRegistry,
759762
AppData()->TimeProvider, AppData()->RandomProvider, Config->EnableKqpImmediateEffects);
760763
QueryState->QueryData = std::make_shared<TQueryData>(QueryState->TxCtx->TxAlloc);

‎ydb/core/kqp/ut/service/kqp_qs_queries_ut.cpp‎

Lines changed: 99 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2159,6 +2159,77 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
21592159
TKikimrRunner kikimr(serverSettings);
21602160
auto db = kikimr.GetQueryClient();
21612161

2162+
{
2163+
// DDl + DML with explicit transaction
2164+
auto result = db.ExecuteQuery(R"(
2165+
CREATE TABLE TestDdlDml1 (
2166+
Key Uint64,
2167+
Value1 String,
2168+
Value2 String,
2169+
PRIMARY KEY (Key)
2170+
);
2171+
UPSERT INTO TestDdlDml1 (Key, Value1, Value2) VALUES (1, "1", "2");
2172+
SELECT * FROM TestDdlDml1;
2173+
ALTER TABLE TestDdlDml1 DROP COLUMN Value2;
2174+
UPSERT INTO TestDdlDml1 (Key, Value1) VALUES (2, "2");
2175+
SELECT * FROM TestDdlDml1;
2176+
)", TTxControl::BeginTx().CommitTx()).ExtractValueSync();
2177+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::GENERIC_ERROR, result.GetIssues().ToString());
2178+
UNIT_ASSERT(result.GetIssues().ToOneLineString().Contains("Queries with mixed data and scheme operations are not supported."));
2179+
}
2180+
2181+
{
2182+
// DDl + DML with implicit transaction
2183+
auto result = db.ExecuteQuery(R"(
2184+
CREATE TABLE TestDdlDml2 (
2185+
Key Uint64,
2186+
Value1 String,
2187+
Value2 String,
2188+
PRIMARY KEY (Key)
2189+
);
2190+
UPSERT INTO TestDdlDml2 (Key, Value1, Value2) VALUES (1, "1", "2");
2191+
SELECT * FROM TestDdlDml2;
2192+
ALTER TABLE TestDdlDml2 DROP COLUMN Value2;
2193+
UPSERT INTO TestDdlDml2 (Key, Value1) VALUES (2, "2");
2194+
SELECT * FROM TestDdlDml2;
2195+
)", TTxControl::NoTx()).ExtractValueSync();
2196+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
2197+
UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 2);
2198+
CompareYson(R"([[[1u];["1"];["2"]]])", FormatResultSetYson(result.GetResultSet(0)));
2199+
CompareYson(R"([[[1u];["1"]];[[2u];["2"]]])", FormatResultSetYson(result.GetResultSet(1)));
2200+
UNIT_ASSERT_EQUAL_C(result.GetIssues().Size(), 0, result.GetIssues().ToString());
2201+
2202+
result = db.ExecuteQuery(R"(
2203+
SELECT * FROM TestDdlDml2;
2204+
)", TTxControl::BeginTx().CommitTx()).ExtractValueSync();
2205+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
2206+
UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 1);
2207+
CompareYson(R"([[[1u];["1"]];[[2u];["2"]]])", FormatResultSetYson(result.GetResultSet(0)));
2208+
2209+
result = db.ExecuteQuery(R"(
2210+
CREATE TABLE TestDdlDml4 (
2211+
Key Uint64,
2212+
Value1 String,
2213+
Value2 String,
2214+
PRIMARY KEY (Key)
2215+
);
2216+
UPSERT INTO TestDdlDml4 (Key, Value1, Value2) VALUES (1, "1", "2");
2217+
SELECT * FROM TestDdlDml4;
2218+
ALTER TABLE TestDdlDml4 DROP COLUMN Value2;
2219+
UPSERT INTO TestDdlDml4 (Key, Value1) VALUES (2, "2");
2220+
SELECT * FROM TestDdlDml5;
2221+
)", TTxControl::NoTx()).ExtractValueSync();
2222+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SCHEME_ERROR, result.GetIssues().ToString());
2223+
UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 0);
2224+
2225+
result = db.ExecuteQuery(R"(
2226+
SELECT * FROM TestDdlDml4;
2227+
)", TTxControl::NoTx()).ExtractValueSync();
2228+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
2229+
UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 1);
2230+
CompareYson(R"([])", FormatResultSetYson(result.GetResultSet(0)));
2231+
}
2232+
21622233
{
21632234
// Base test with ddl and dml statements
21642235
auto result = db.ExecuteQuery(R"(
@@ -2228,13 +2299,21 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
22282299
{
22292300
// Test with query with error
22302301
auto result = db.ExecuteQuery(R"(
2231-
UPSERT INTO TestDdl2 (Key, Value) VALUES (1, "One");
2302+
SELECT * FROM TestDdl2;
2303+
)", TTxControl::NoTx()).ExtractValueSync();
2304+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
2305+
UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 1);
2306+
CompareYson(R"([[[1u];["One"]];[[2u];["Two"]];[[3u];["Three"]]])", FormatResultSetYson(result.GetResultSet(0)));
2307+
UNIT_ASSERT_EQUAL_C(result.GetIssues().Size(), 0, result.GetIssues().ToString());
2308+
2309+
result = db.ExecuteQuery(R"(
2310+
UPSERT INTO TestDdl2 (Key, Value) VALUES (4, "Four");
22322311
CREATE TABLE TestDdl3 (
22332312
Key Uint64,
22342313
Value String,
22352314
PRIMARY KEY (Key)
22362315
);
2237-
UPSERT INTO TestDdl2 (Key, Value) VALUES (4, "Four");
2316+
UPSERT INTO TestDdl2 (Key, Value) VALUES (5, "Five");
22382317
CREATE TABLE TestDdl2 (
22392318
Key Uint64,
22402319
Value String,
@@ -2245,18 +2324,32 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
22452324
Value String,
22462325
PRIMARY KEY (Key)
22472326
);
2248-
UPSERT INTO TestDdl1 (Key, Value) VALUES (3, "Three");
22492327
)", TTxControl::NoTx()).ExtractValueSync();
22502328
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::GENERIC_ERROR, result.GetIssues().ToString());
22512329
UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 0);
2330+
UNIT_ASSERT(result.GetIssues().ToOneLineString().Contains("Check failed: path: '/Root/TestDdl2', error: path exist"));
22522331

22532332
result = db.ExecuteQuery(R"(
22542333
SELECT * FROM TestDdl2;
22552334
)", TTxControl::NoTx()).ExtractValueSync();
22562335
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
22572336
UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 1);
2258-
CompareYson(R"([[[1u];["One"]];[[2u];["Two"]];[[3u];["Three"]];[[4u];["Four"]]])", FormatResultSetYson(result.GetResultSet(0)));
2337+
CompareYson(R"([[[1u];["One"]];[[2u];["Two"]];[[3u];["Three"]]])", FormatResultSetYson(result.GetResultSet(0)));
2338+
UNIT_ASSERT_EQUAL_C(result.GetIssues().Size(), 0, result.GetIssues().ToString());
2339+
2340+
result = db.ExecuteQuery(R"(
2341+
SELECT * FROM TestDdl3;
2342+
)", TTxControl::NoTx()).ExtractValueSync();
2343+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
2344+
UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 1);
2345+
CompareYson(R"([])", FormatResultSetYson(result.GetResultSet(0)));
22592346
UNIT_ASSERT_EQUAL_C(result.GetIssues().Size(), 0, result.GetIssues().ToString());
2347+
2348+
result = db.ExecuteQuery(R"(
2349+
SELECT * FROM TestDdl4;
2350+
)", TTxControl::NoTx()).ExtractValueSync();
2351+
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SCHEME_ERROR, result.GetIssues().ToString());
2352+
UNIT_ASSERT(result.GetIssues().ToOneLineString().Contains("Cannot find table 'db.[/Root/TestDdl4]'"));
22602353
}
22612354

22622355
{
@@ -2329,7 +2422,7 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
23292422
)", TTxControl::NoTx()).ExtractValueSync();
23302423
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
23312424
UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 1);
2332-
CompareYson(R"([[[1u];[1u]];[[2u];[2u]]])", FormatResultSetYson(result.GetResultSet(0)));
2425+
CompareYson(R"([[[1u];[1u]]])", FormatResultSetYson(result.GetResultSet(0)));
23332426

23342427
result = db.ExecuteQuery(R"(
23352428
CREATE TABLE TestDdl5 (
@@ -2348,7 +2441,7 @@ Y_UNIT_TEST_SUITE(KqpQueryService) {
23482441
)", TTxControl::NoTx()).ExtractValueSync();
23492442
UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::SUCCESS, result.GetIssues().ToString());
23502443
UNIT_ASSERT_VALUES_EQUAL(result.GetResultSets().size(), 1);
2351-
CompareYson(R"([[[1u];[1u]]])", FormatResultSetYson(result.GetResultSet(0)));
2444+
CompareYson(R"([])", FormatResultSetYson(result.GetResultSet(0)));
23522445
}
23532446
}
23542447

0 commit comments

Comments
 (0)