Merge f51c50f into a18f18d

Author: what-the-fawk

ydb/core/fq/libs/http_api_client/http_client.py

@@ -149,6 +149,21 @@ def create_query(
 
         self._validate_http_error(response, expected_code=expected_code)
         return response.json()["id"]
+    
+    def restart_query(
+        self,
+        query_id: str,
+        idempotency_key: str | None = None,
+        expected_code: int = 204,
+    ):
+        response = self.session.get(
+            self._compose_api_url(f"/api/fq/v1/queries/{query_id}/start"),
+            headers=self._build_headers(idempotency_key=idempotency_key),
+            params=self._build_params(),
+        )
+
+        self._validate_http_error(response, expected_code)
+        return response.json()["status"]
 
     def get_query_status(self, query_id, request_id=None, expected_code=200) -> Any:
         response = self.session.get(

ydb/core/public_http/fq_handlers.h

@@ -249,6 +249,7 @@ void SetIdempotencyKey(T& dst, const TString& key) {
 
 template <typename GrpcProtoRequestType, typename HttpProtoRequestType, typename GrpcProtoResultType, typename HttpProtoResultType, typename GrpcProtoResponseType>
 class TGrpcCallWrapper : public TActorBootstrapped<TGrpcCallWrapper<GrpcProtoRequestType, HttpProtoRequestType, GrpcProtoResultType, HttpProtoResultType, GrpcProtoResponseType>> {
+protected:
     THttpRequestContext RequestContext;
 
     typedef std::function<std::unique_ptr<NGRpcService::TEvProxyRuntimeEvent>(TIntrusivePtr<NYdbGrpc::IRequestContextBase> ctx)> TGrpcProxyEventFactory;
@@ -344,20 +345,22 @@ class TGrpcCallWrapper : public TActorBootstrapped<TGrpcCallWrapper<GrpcProtoReq
         RequestContext.ResponseBadRequest(Ydb::StatusIds::BAD_REQUEST, error);
     }
 
+    #define CheckTypedResponse(typedResponse, ResultType)                                                   \
+    if (!typedResponse->operation().result().template Is<ResultType>()) {                                   \
+            TStringStream json;                                                                             \
+            auto* httpResult = google::protobuf::Arena::CreateMessage<FQHttp::Error>(resp->GetArena());     \
+            FqConvert(typedResponse->operation(), *httpResult);                                             \
+            FqPackToJson(json, *httpResult, jsonSettings);                                                  \
+            requestContext.ResponseBadRequestJson(typedResponse->operation().status(), json.Str());         \
+            return;                                                                                         \
+    }
+
     static void SendReply(const THttpRequestContext& requestContext, const TJsonSettings& jsonSettings, NProtoBuf::Message* resp, ui32 status) {
         Y_ABORT_UNLESS(resp);
         Y_ABORT_UNLESS(resp->GetArena());
         Y_UNUSED(status);
         auto* typedResponse = static_cast<TGrpcProtoResponseType*>(resp);
-        if (!typedResponse->operation().result().template Is<TGrpcProtoResultType>()) {
-            TStringStream json;
-            auto* httpResult = google::protobuf::Arena::CreateMessage<FQHttp::Error>(resp->GetArena());
-            FqConvert(typedResponse->operation(), *httpResult);
-            FqPackToJson(json, *httpResult, jsonSettings);
-
-            requestContext.ResponseBadRequestJson(typedResponse->operation().status(), json.Str());
-            return;
-        }
+        CheckTypedResponse(typedResponse, TGrpcProtoResultType);
 
         auto* grpcResult = google::protobuf::Arena::CreateMessage<TGrpcProtoResultType>(resp->GetArena());
         typedResponse->operation().result().UnpackTo(grpcResult);
@@ -396,4 +399,63 @@ DECLARE_YQ_GRPC_ACTOR(GetQueryStatus, GetQueryStatus);
 DECLARE_YQ_GRPC_ACTOR_WIHT_EMPTY_RESULT(StopQuery, ControlQuery);
 DECLARE_YQ_GRPC_ACTOR(GetResultData, GetResultData);
 
+#define TGrpcCallWrapperBase TGrpcCallWrapper<FederatedQuery::DescribeQueryRequest, FQHttp::GetQueryRequest, FederatedQuery::ModifyQueryResult, google::protobuf::Empty, FederatedQuery::ModifyQueryResponse>
+
+class TRestartQueryRequest : public TGrpcCallWrapperBase {
+public:
+    TRestartQueryRequest(const THttpRequestContext& ctx)
+    : TGrpcCallWrapperBase(ctx, &NGRpcService::CreateFederatedQueryDescribeQueryRequestOperationCall)
+    {}
+
+    void Bootstrap(const TActorContext& ctx) {
+        auto describeRequest = std::make_unique<FederatedQuery::DescribeQueryRequest>();
+        if (!Parse(*describeRequest)) {
+            this->Die(ctx);
+            return;
+        }
+
+        TIntrusivePtr<TGrpcRequestContextWrapper> requestContext = new TGrpcRequestContextWrapper(
+            RequestContext,
+            std::move(describeRequest),
+            [&, this, query_id = describeRequest->Getquery_id()](const THttpRequestContext& requestContext, const TJsonSettings& jsonSettings, NProtoBuf::Message* resp, ui32 status) {
+            Y_ABORT_UNLESS(resp);
+            Y_ABORT_UNLESS(resp->GetArena());
+
+            auto* typedResponse = static_cast<FederatedQuery::DescribeQueryResponse*>(resp);
+            CheckTypedResponse(typedResponse, FederatedQuery::DescribeQueryResult);
+
+            FederatedQuery::DescribeQueryResult* describeResult = google::protobuf::Arena::CreateMessage<FederatedQuery::DescribeQueryResult>(resp->GetArena());
+            typedResponse->operation().result().UnpackTo(describeResult);
+
+
+            // modify
+            auto modifyRequest = std::unique_ptr<FederatedQuery::ModifyQueryRequest>(google::protobuf::Arena::CreateMessage<FederatedQuery::ModifyQueryRequest>(resp->GetArena()));
+
+            modifyRequest->set_query_id(query_id);
+            auto content = describeResult->Getquery().content();
+            modifyRequest->set_allocated_content(&content);
+            modifyRequest->set_execute_mode(::FederatedQuery::ExecuteMode::RUN);
+            modifyRequest->set_allocated_disposition(nullptr);
+            modifyRequest->set_state_load_mode(::FederatedQuery::StateLoadMode::FROM_LAST_CHECKPOINT);
+            modifyRequest->set_previous_revision(describeResult->Getquery().meta().Getlast_job_query_revision());
+            modifyRequest->set_idempotency_key(requestContext.GetIdempotencyKey());
+
+            TIntrusivePtr<TGrpcRequestContextWrapper> requestContextModify = new TGrpcRequestContextWrapper(
+                this->RequestContext,
+                std::move(modifyRequest),
+                TGrpcCallWrapper<FederatedQuery::ModifyQueryRequest, int, FederatedQuery::ModifyQueryResult, google::protobuf::Empty, FederatedQuery::ModifyQueryResponse>::SendReply
+            );
+
+            // новый тип события, по идее нужна новая фабрика
+            EventFactory = &NGRpcService::CreateFederatedQueryModifyQueryRequestOperationCall;
+            ctx.Send(NGRpcService::CreateGRpcRequestProxyId(), EventFactory(requestContextModify).release());
+            this->Die(ctx);
+        });
+
+        ctx.Send(NGRpcService::CreateGRpcRequestProxyId(), EventFactory(requestContext).release());
+    }
+};
+
+#undef TGrpcCallWrapperBase
+
 } // namespace NKikimr::NPublicHttp

ydb/core/public_http/http_service.cpp

@@ -57,6 +57,7 @@ namespace {
             Router.RegisterHandler(HTTP_METHOD_GET,  "/api/fq/v1/queries/{query_id}/status", CreateHttpHandler<TJsonGetQueryStatus>());
             Router.RegisterHandler(HTTP_METHOD_GET,  "/api/fq/v1/queries/{query_id}/results/{result_set_index}", CreateHttpHandler<TJsonGetResultData>());
             Router.RegisterHandler(HTTP_METHOD_POST, "/api/fq/v1/queries/{query_id}/stop", CreateHttpHandler<TJsonStopQuery>());
+            Router.RegisterHandler(HTTP_METHOD_POST, "/api/fq/v1/queries/{query_id}/start", CreateHttpHandler<TRestartQueryRequest>());
         }
 
         void Bootstrap(const TActorContext& ctx) {

ydb/tests/fq/http_api/test_http_api.py

@@ -98,6 +98,12 @@ def test_simple_analytics_query(self):
             response = client.stop_query(query_id)
             assert response.status_code == 204
 
+            response = client.restart_query(query_id)
+            assert response.status_code == 204
+
+            response = client.stop_query(query_id)
+            assert response.status_code == 204
+
     def test_empty_query(self):
         with self.create_client() as client:
             with pytest.raises(
@@ -109,6 +115,14 @@ def test_empty_query(self):
             ):  # noqa
                 client.create_query()
 
+    def test_invalid_id(self):
+        with self.create_client() as client:
+            resp = client.stop_quert(query_id="nevalidno")
+            assert resp.status_code == 404
+
+            resp = client.restart_query(query_id="snova nevalidno")
+            assert resp.status_code == 404
+
     def test_warning(self):
         with self.create_client() as client:
             query_id = client.create_query(query_text="select 10000000000000000000+1")
@@ -228,6 +242,30 @@ def test_stop_idempotency(self):
         self.streaming_over_kikimr.compute_plane.start()
         c.wait_query_status(query_id, fq.QueryMeta.ABORTED_BY_USER)
 
+    def test_restart_idempotency(self):
+        c = FederatedQueryClient("my_folder", streaming_over_kikimr=self.streaming_over_kikimr)
+        self.streaming_over_kikimr.compute_plane.stop()
+        query_id = c.create_query("select1", "select 1").result.query_id
+        c.wait_query_status(query_id, fq.QueryMeta.STARTING)
+    
+        with self.create_client() as client:
+            response1 = client.stop_query(query_id, idempotency_key="Z")
+            assert response1.status_code == 204
+
+            response2 = client.restart_query(query_id, idempotency_key="Z")
+            assert response2.status_code == 204
+
+            response2 = client.restart_query(query_id, idempotency_key="Z")
+            assert response2.status_code == 204
+
+            response2 = client.restart_query(query_id)
+            assert response2.status_code == 400
+            
+        
+        self.streaming_over_kikimr.compute_plane.start()
+        c.wait_query_status(query_id, fq.QueryMeta.ABORTED_BY_USER)
+        pass
+
     def test_simple_streaming_query(self):
         self.init_topics("simple_streaming_query", create_output=False)
         c = FederatedQueryClient("my_folder", streaming_over_kikimr=self.streaming_over_kikimr)
@@ -259,6 +297,9 @@ def test_simple_streaming_query(self):
             response = client.stop_query(query_id)
             assert response.status_code == 204
 
+            response = client.restart_query(query_id)
+            assert response.status_code == 204
+
             wait_for_query_status(client, query_id, ["FAILED"])
 
             query_json2 = client.get_query(query_id)