Skip to content

Commit 09278b9

Browse files
vhsu14Presto CUDF CI
authored andcommitted
[native] Enable getTaskInfo and deleteTask (prestodb#25535)
## Description - Native changes to enable thrift for getTaskInfo and deleteTask endpoints - Depends on prestodb#25464 ## Motivation and Context prestodb/rfcs#38 ## Test Plan Build package and test in verifier: 230498 ## Release Notes Please follow [release notes guidelines](https://github.com/prestodb/presto/wiki/Release-Notes-Guidelines) and fill in the release notes below. ``` == NO RELEASE NOTE == ```
1 parent b418e82 commit 09278b9

File tree

1 file changed

+31
-8
lines changed

1 file changed

+31
-8
lines changed

presto-native-execution/presto_cpp/main/TaskResource.cpp

Lines changed: 31 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -392,9 +392,13 @@ proxygen::RequestHandler* TaskResource::deleteTask(
392392
message->getQueryParam(protocol::PRESTO_ABORT_TASK_URL_PARAM) == "true";
393393
}
394394
bool summarize = message->hasQueryParam("summarize");
395+
auto& headers = message->getHeaders();
396+
const auto& acceptHeader = headers.getSingleOrEmpty(proxygen::HTTP_HEADER_ACCEPT);
397+
const auto sendThrift =
398+
acceptHeader.find(http::kMimeTypeApplicationThrift) != std::string::npos;
395399

396400
return new http::CallbackRequestHandler(
397-
[this, taskId, abort, summarize](
401+
[this, taskId, abort, summarize, sendThrift](
398402
proxygen::HTTPMessage* /*message*/,
399403
const std::vector<std::unique_ptr<folly::IOBuf>>& /*body*/,
400404
proxygen::ResponseHandler* downstream,
@@ -407,12 +411,19 @@ proxygen::RequestHandler* TaskResource::deleteTask(
407411
return std::move(taskInfo);
408412
})
409413
.via(folly::EventBaseManager::get()->getEventBase())
410-
.thenValue([taskId, downstream, handlerState](auto&& taskInfo) {
414+
.thenValue([taskId, downstream, handlerState, sendThrift](auto&& taskInfo) {
411415
if (!handlerState->requestExpired()) {
412416
if (taskInfo == nullptr) {
413417
sendTaskNotFound(downstream, taskId);
414418
}
415-
http::sendOkResponse(downstream, json(*taskInfo));
419+
if (sendThrift) {
420+
thrift::TaskInfo thriftTaskInfo;
421+
toThrift(*taskInfo, thriftTaskInfo);
422+
http::sendOkThriftResponse(
423+
downstream, thriftWrite(thriftTaskInfo));
424+
} else {
425+
http::sendOkResponse(downstream, json(*taskInfo));
426+
}
416427
}
417428
})
418429
.thenError(
@@ -604,8 +615,13 @@ proxygen::RequestHandler* TaskResource::getTaskInfo(
604615
auto maxWait = getMaxWait(message);
605616
bool summarize = message->hasQueryParam("summarize");
606617

618+
auto& headers = message->getHeaders();
619+
const auto& acceptHeader = headers.getSingleOrEmpty(proxygen::HTTP_HEADER_ACCEPT);
620+
const auto sendThrift =
621+
acceptHeader.find(http::kMimeTypeApplicationThrift) != std::string::npos;
622+
607623
return new http::CallbackRequestHandler(
608-
[this, taskId, currentState, maxWait, summarize](
624+
[this, taskId, currentState, maxWait, summarize, sendThrift](
609625
proxygen::HTTPMessage* /*message*/,
610626
const std::vector<std::unique_ptr<folly::IOBuf>>& /*body*/,
611627
proxygen::ResponseHandler* downstream,
@@ -619,16 +635,23 @@ proxygen::RequestHandler* TaskResource::getTaskInfo(
619635
maxWait,
620636
summarize,
621637
handlerState,
622-
downstream]() {
638+
downstream,
639+
sendThrift]() {
623640
taskManager_
624641
.getTaskInfo(
625642
taskId, summarize, currentState, maxWait, handlerState)
626643
.via(evb)
627-
.thenValue([downstream, taskId, handlerState](
644+
.thenValue([downstream, taskId, handlerState, sendThrift](
628645
std::unique_ptr<protocol::TaskInfo> taskInfo) {
629646
if (!handlerState->requestExpired()) {
630-
json taskInfoJson = *taskInfo;
631-
http::sendOkResponse(downstream, taskInfoJson);
647+
if (sendThrift) {
648+
thrift::TaskInfo thriftTaskInfo;
649+
toThrift(*taskInfo, thriftTaskInfo);
650+
http::sendOkThriftResponse(
651+
downstream, thriftWrite(thriftTaskInfo));
652+
} else {
653+
http::sendOkResponse(downstream, json(*taskInfo));
654+
}
632655
}
633656
})
634657
.thenError(

0 commit comments

Comments
 (0)