Skip to content

Commit 780b669

Browse files
GrigoriyPAGrigoriyPA
committed
YQ fixed script execution id parsing (#18462)
1 parent d883baf commit 780b669

File tree

6 files changed

+77
-31
lines changed

6 files changed

+77
-31
lines changed

‎ydb/core/grpc_services/query/rpc_fetch_script_results.cpp

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -125,18 +125,14 @@ class TFetchScriptResultsRPC : public TRpcRequestActor<TFetchScriptResultsRPC, T
125125
}
126126

127127
bool GetExecutionIdFromRequest() {
128-
try {
129-
TMaybe<TString> executionId = NKqp::ScriptExecutionIdFromOperation(GetProtoRequest()->operation_id());
130-
if (!executionId) {
131-
Reply(Ydb::StatusIds::BAD_REQUEST, "Invalid operation id");
132-
return false;
133-
}
134-
ExecutionId = *executionId;
135-
return true;
136-
} catch (const std::exception& ex) {
137-
Reply(Ydb::StatusIds::BAD_REQUEST, TStringBuilder() << "Invalid operation id: " << ex.what());
128+
TString error;
129+
TMaybe<TString> executionId = NKqp::ScriptExecutionIdFromOperation(GetProtoRequest()->operation_id(), error);
130+
if (!executionId) {
131+
Reply(Ydb::StatusIds::BAD_REQUEST, error);
138132
return false;
139133
}
134+
ExecutionId = *executionId;
135+
return true;
140136
}
141137

142138
private:

‎ydb/core/kqp/common/kqp_script_executions.cpp

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#include "kqp_script_executions.h"
22

3+
#include <util/string/builder.h>
4+
35
#include <ydb/public/sdk/cpp/src/library/operation_id/protos/operation_id.pb.h>
46

57
namespace NKikimr::NKqp {
@@ -11,21 +13,29 @@ TString ScriptExecutionOperationFromExecutionId(const TString& executionId) {
1113
return NOperationId::ProtoToString(operationId);
1214
}
1315

14-
TMaybe<TString> ScriptExecutionIdFromOperation(const TString& operationId) {
16+
TMaybe<TString> ScriptExecutionIdFromOperation(const TString& operationId, TString& error) try {
1517
NOperationId::TOperationId operation(operationId);
16-
return ScriptExecutionIdFromOperation(operation);
18+
return ScriptExecutionIdFromOperation(operation, error);
19+
} catch (const std::exception& ex) {
20+
error = TStringBuilder() << "Invalid operation id: " << ex.what();
21+
return Nothing();
1722
}
1823

19-
TMaybe<TString> ScriptExecutionIdFromOperation(const NOperationId::TOperationId& operationId) {
24+
TMaybe<TString> ScriptExecutionIdFromOperation(const NOperationId::TOperationId& operationId, TString& error) try {
2025
if (operationId.GetKind() != NOperationId::TOperationId::SCRIPT_EXECUTION) {
26+
error = TStringBuilder() << "Invalid operation id, expected SCRIPT_EXECUTION = " << static_cast<int>(NOperationId::TOperationId::SCRIPT_EXECUTION) << " kind, got " << static_cast<int>(operationId.GetKind());
2127
return Nothing();
2228
}
2329

2430
const auto& values = operationId.GetValue("id");
2531
if (values.empty() || !values[0]) {
32+
error = TStringBuilder() << "Invalid operation id, please specify key 'id'";
2633
return Nothing();
2734
}
2835
return TString{*values[0]};
36+
} catch (const std::exception& ex) {
37+
error = TStringBuilder() << "Invalid operation id: " << ex.what();
38+
return Nothing();
2939
}
3040

3141
} // namespace NKikimr::NKqp

‎ydb/core/kqp/common/kqp_script_executions.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
namespace NKikimr::NKqp {
99

1010
TString ScriptExecutionOperationFromExecutionId(const TString& executionId);
11-
TMaybe<TString> ScriptExecutionIdFromOperation(const TString& operationId);
12-
TMaybe<TString> ScriptExecutionIdFromOperation(const NOperationId::TOperationId& operationId);
11+
TMaybe<TString> ScriptExecutionIdFromOperation(const TString& operationId, TString& error);
12+
TMaybe<TString> ScriptExecutionIdFromOperation(const NOperationId::TOperationId& operationId, TString& error);
1313

1414
} // namespace NKikimr::NKqp

‎ydb/core/kqp/proxy_service/kqp_script_executions.cpp

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -990,9 +990,10 @@ class TForgetScriptExecutionOperationActor : public TActorBootstrapped<TForgetSc
990990
{}
991991

992992
void Bootstrap() {
993-
TMaybe<TString> executionId = NKqp::ScriptExecutionIdFromOperation(Request->Get()->OperationId);
993+
TString error;
994+
TMaybe<TString> executionId = NKqp::ScriptExecutionIdFromOperation(Request->Get()->OperationId, error);
994995
if (!executionId) {
995-
Reply(Ydb::StatusIds::BAD_REQUEST, "Incorrect operation id");
996+
Reply(Ydb::StatusIds::BAD_REQUEST, error);
996997
return;
997998
}
998999
ExecutionId = *executionId;
@@ -1258,9 +1259,10 @@ class TGetScriptExecutionOperationActor : public TCheckLeaseStatusActorBase {
12581259
{}
12591260

12601261
void OnBootstrap() override {
1261-
TMaybe<TString> executionId = ScriptExecutionIdFromOperation(Request->Get()->OperationId);
1262+
TString error;
1263+
TMaybe<TString> executionId = ScriptExecutionIdFromOperation(Request->Get()->OperationId, error);
12621264
if (!executionId) {
1263-
Reply(Ydb::StatusIds::BAD_REQUEST, "Incorrect operation id");
1265+
Reply(Ydb::StatusIds::BAD_REQUEST, error);
12641266
return;
12651267
}
12661268
ExecutionId = *executionId;
@@ -1600,9 +1602,10 @@ class TCancelScriptExecutionOperationActor : public NActors::TActorBootstrapped<
16001602
{}
16011603

16021604
void Bootstrap() {
1603-
const TMaybe<TString> executionId = NKqp::ScriptExecutionIdFromOperation(Request->Get()->OperationId);
1605+
TString error;
1606+
const TMaybe<TString> executionId = NKqp::ScriptExecutionIdFromOperation(Request->Get()->OperationId, error);
16041607
if (!executionId) {
1605-
return Reply(Ydb::StatusIds::BAD_REQUEST, "Incorrect operation id");
1608+
return Reply(Ydb::StatusIds::BAD_REQUEST, error);
16061609
}
16071610
ExecutionId = *executionId;
16081611

@@ -1616,14 +1619,15 @@ class TCancelScriptExecutionOperationActor : public NActors::TActorBootstrapped<
16161619
hFunc(TEvKqp::TEvCancelScriptExecutionResponse, Handle);
16171620
hFunc(NActors::TEvents::TEvUndelivered, Handle);
16181621
hFunc(NActors::TEvInterconnect::TEvNodeDisconnected, Handle);
1622+
IgnoreFunc(NActors::TEvInterconnect::TEvNodeConnected);
16191623
)
16201624

16211625
void Handle(TEvPrivate::TEvLeaseCheckResult::TPtr& ev) {
16221626
if (ev->Get()->Status == Ydb::StatusIds::SUCCESS) {
16231627
KQP_PROXY_LOG_D("[TCancelScriptExecutionOperationActor] ExecutionId: " << ExecutionId << ", check lease success");
16241628
RunScriptActor = ev->Get()->RunScriptActorId;
16251629
if (ev->Get()->OperationStatus) {
1626-
Reply(Ydb::StatusIds::PRECONDITION_FAILED); // Already finished.
1630+
Reply(Ydb::StatusIds::PRECONDITION_FAILED, "Script execution operation is already finished");
16271631
} else {
16281632
if (CancelSent) { // We have not found the actor, but after it status of the operation is not defined, something strage happened.
16291633
Reply(Ydb::StatusIds::INTERNAL_ERROR, "Failed to cancel script execution operation");

‎ydb/core/kqp/ut/federated_query/generic_ut/kqp_generic_provider_ut.cpp

Lines changed: 41 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
#include <ydb/library/yql/providers/generic/connector/libcpp/ut_helpers/database_resolver_mock.h>
77
#include <ydb/library/yql/providers/s3/actors/yql_s3_actors_factory_impl.h>
88
#include <ydb/public/api/protos/ydb_query.pb.h>
9+
#include <ydb/public/api/grpc/ydb_operation_v1.grpc.pb.h>
910
#include <ydb/public/api/grpc/ydb_query_v1.grpc.pb.h>
1011
#include <ydb/public/sdk/cpp/src/library/grpc/client/grpc_client_low.h>
1112
#include <ydb-cpp-sdk/client/operation/operation.h>
@@ -489,27 +490,60 @@ namespace NKikimr::NKqp {
489490

490491
// Create trash query
491492
NYdbGrpc::TGRpcClientLow clientLow;
493+
const auto channel = grpc::CreateChannel("localhost:" + ToString(kikimr->GetTestServer().GetGRpcServer().GetPort()), grpc::InsecureChannelCredentials());
494+
const auto queryServiceStub = Ydb::Query::V1::QueryService::NewStub(channel);
495+
const auto operationServiceStub = Ydb::Operation::V1::OperationService::NewStub(channel);
492496

493-
std::shared_ptr<grpc::Channel> channel;
494-
channel = grpc::CreateChannel("localhost:" + ToString(kikimr->GetTestServer().GetGRpcServer().GetPort()), grpc::InsecureChannelCredentials());
495497
{
496-
std::unique_ptr<Ydb::Query::V1::QueryService::Stub> stub;
497-
stub = Ydb::Query::V1::QueryService::NewStub(channel);
498498
grpc::ClientContext context;
499499
Ydb::Query::FetchScriptResultsRequest request;
500500
request.set_operation_id(operationId);
501501
request.set_fetch_token(fetchToken);
502502
Ydb::Query::FetchScriptResultsResponse response;
503-
grpc::Status st = stub->FetchScriptResults(&context, request, &response);
503+
grpc::Status st = queryServiceStub->FetchScriptResults(&context, request, &response);
504504
UNIT_ASSERT(st.ok());
505-
UNIT_ASSERT_EQUAL_C(response.status(), Ydb::StatusIds::BAD_REQUEST, response);
505+
UNIT_ASSERT_VALUES_EQUAL_C(response.status(), Ydb::StatusIds::BAD_REQUEST, response);
506+
}
507+
508+
{
509+
grpc::ClientContext context;
510+
Ydb::Operations::ForgetOperationRequest request;
511+
request.set_id(operationId);
512+
Ydb::Operations::ForgetOperationResponse response;
513+
grpc::Status st = operationServiceStub->ForgetOperation(&context, request, &response);
514+
UNIT_ASSERT(st.ok());
515+
UNIT_ASSERT_VALUES_EQUAL_C(response.status(), Ydb::StatusIds::BAD_REQUEST, response);
516+
}
517+
518+
{
519+
grpc::ClientContext context;
520+
Ydb::Operations::GetOperationRequest request;
521+
request.set_id(operationId);
522+
Ydb::Operations::GetOperationResponse response;
523+
grpc::Status st = operationServiceStub->GetOperation(&context, request, &response);
524+
UNIT_ASSERT(st.ok());
525+
UNIT_ASSERT_VALUES_EQUAL_C(response.operation().status(), Ydb::StatusIds::BAD_REQUEST, response);
526+
}
527+
528+
{
529+
grpc::ClientContext context;
530+
Ydb::Operations::CancelOperationRequest request;
531+
request.set_id(operationId);
532+
Ydb::Operations::CancelOperationResponse response;
533+
grpc::Status st = operationServiceStub->CancelOperation(&context, request, &response);
534+
UNIT_ASSERT(st.ok());
535+
UNIT_ASSERT_VALUES_EQUAL_C(response.status(), Ydb::StatusIds::BAD_REQUEST, response);
506536
}
507537
}
508538

509-
Y_UNIT_TEST(TestFailsOnIncorrectScriptExecutionOperationId) {
539+
Y_UNIT_TEST(TestFailsOnIncorrectScriptExecutionOperationId1) {
510540
TestFailsOnIncorrectScriptExecutionOperation("trash", "");
511541
}
512542

543+
Y_UNIT_TEST(TestFailsOnIncorrectScriptExecutionOperationId2) {
544+
TestFailsOnIncorrectScriptExecutionOperation("ydb://scriptexec/9?fd=b214872a-d040e60d-62a1b34-a9be3c3d", "trash");
545+
}
546+
513547
Y_UNIT_TEST(TestFailsOnIncorrectScriptExecutionFetchToken) {
514548
TestFailsOnIncorrectScriptExecutionOperation("", "trash");
515549
}

‎ydb/tests/tools/kqprun/src/ydb_setup.cpp

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -516,13 +516,15 @@ class TYdbSetup::TImpl {
516516
}
517517

518518
NKikimr::NKqp::TEvFetchScriptResultsResponse::TPtr FetchScriptExecutionResultsRequest(const TString& database, const TString& operation, i32 resultSetId) const {
519-
TString executionId = *NKikimr::NKqp::ScriptExecutionIdFromOperation(operation);
519+
TString error;
520+
const auto executionId = NKikimr::NKqp::ScriptExecutionIdFromOperation(operation, error);
521+
Y_ENSURE(executionId, error);
520522

521523
ui32 nodeIndex = GetNodeIndexForDatabase(database);
522524
NActors::TActorId edgeActor = GetRuntime()->AllocateEdgeActor(nodeIndex);
523525
auto rowsLimit = Settings_.AppConfig.GetQueryServiceConfig().GetScriptResultRowsLimit();
524526
auto sizeLimit = Settings_.AppConfig.GetQueryServiceConfig().GetScriptResultSizeLimit();
525-
NActors::IActor* fetchActor = NKikimr::NKqp::CreateGetScriptExecutionResultActor(edgeActor, GetDatabasePath(database), executionId, resultSetId, 0, rowsLimit, sizeLimit, TInstant::Max());
527+
NActors::IActor* fetchActor = NKikimr::NKqp::CreateGetScriptExecutionResultActor(edgeActor, GetDatabasePath(database), *executionId, resultSetId, 0, rowsLimit, sizeLimit, TInstant::Max());
526528

527529
GetRuntime()->Register(fetchActor, nodeIndex, GetRuntime()->GetAppData(nodeIndex).UserPoolId);
528530

0 commit comments

Comments
 (0)