Skip to content

Commit

Permalink
[KYUUBI #5873] Support to get operation progress with RESTful API
Browse files Browse the repository at this point in the history
# 🔍 Description
## Issue References 🔗

Since #2493, with jdbc conneciton, user can get TProgressUpdateResp to check the operation progress percentage, but for RESTful api, it is not supported yet.

As title, this PR supports to get operation progress with RESTful api.
This pull request fixes #5873

## Describe Your Solution 🔧

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

## Types of changes 🔖

- [ ] Bugfix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)

## Test Plan 🧪

#### Behavior Without This Pull Request ⚰️

#### Behavior With This Pull Request 🎉

#### Related Unit Tests

---

# Checklists
## 📝 Author Self Checklist

- [ ] My code follows the [style guidelines](https://kyuubi.readthedocs.io/en/master/contributing/code/style.html) of this project
- [ ] I have performed a self-review
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] This patch was not authored or co-authored using [Generative Tooling](https://www.apache.org/legal/generative-tooling.html)

## 📝 Committer Pre-Merge Checklist

- [ ] Pull request title is okay.
- [ ] No license issues.
- [ ] Milestone correctly set?
- [ ] Test coverage is ok
- [ ] Assignees are selected.
- [ ] Minimum number of approvals
- [ ] No changes are requested

**Be nice. Be informative.**

Closes #5875 from turboFei/progress_resp.

Closes #5873

f177da0 [Fei Wang] using dto event
0dbd450 [Fei Wang] ut
c0758a5 [Fei Wang] save
29acb12 [Fei Wang] save
9a92cd0 [Fei Wang] save

Authored-by: Fei Wang <[email protected]>
Signed-off-by: Fei Wang <[email protected]>
  • Loading branch information
turboFei committed Dec 19, 2023
1 parent 7e96dc7 commit bd379c5
Show file tree
Hide file tree
Showing 8 changed files with 250 additions and 30 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,7 @@ abstract class AbstractOperation(session: Session) extends Operation with Loggin
this.operationException = opEx
}

def getOperationJobProgress: TProgressUpdateResp = operationJobProgress
def setOperationJobProgress(opJobProgress: TProgressUpdateResp): Unit = {
this.operationJobProgress = opJobProgress
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ public class KyuubiOperationEvent {

private Map<String, String> metrics;

private OperationProgress progress;

public KyuubiOperationEvent() {}

public KyuubiOperationEvent(
Expand All @@ -68,7 +70,8 @@ public KyuubiOperationEvent(
String sessionUser,
String sessionType,
String kyuubiInstance,
Map<String, String> metrics) {
Map<String, String> metrics,
OperationProgress progress) {
this.statementId = statementId;
this.remoteId = remoteId;
this.statement = statement;
Expand All @@ -84,6 +87,7 @@ public KyuubiOperationEvent(
this.sessionType = sessionType;
this.kyuubiInstance = kyuubiInstance;
this.metrics = metrics;
this.progress = progress;
}

public static KyuubiOperationEvent.KyuubiOperationEventBuilder builder() {
Expand Down Expand Up @@ -121,6 +125,8 @@ public static class KyuubiOperationEventBuilder {

private Map<String, String> metrics;

private OperationProgress progress;

public KyuubiOperationEventBuilder() {}

public KyuubiOperationEvent.KyuubiOperationEventBuilder statementId(final String statementId) {
Expand Down Expand Up @@ -201,6 +207,12 @@ public KyuubiOperationEvent.KyuubiOperationEventBuilder metrics(
return this;
}

public KyuubiOperationEvent.KyuubiOperationEventBuilder progress(
final OperationProgress progress) {
this.progress = progress;
return this;
}

public KyuubiOperationEvent build() {
return new KyuubiOperationEvent(
statementId,
Expand All @@ -217,7 +229,8 @@ public KyuubiOperationEvent build() {
sessionUser,
sessionType,
kyuubiInstance,
metrics);
metrics,
progress);
}
}

Expand Down Expand Up @@ -340,4 +353,12 @@ public Map<String, String> getMetrics() {
public void setMetrics(Map<String, String> metrics) {
this.metrics = metrics;
}

public OperationProgress getProgress() {
return progress;
}

public void setProgress(OperationProgress progress) {
this.progress = progress;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ public class OperationData {
private String sessionType;
private String kyuubiInstance;
private Map<String, String> metrics;
private OperationProgress progress;

public OperationData() {}

Expand All @@ -53,7 +54,8 @@ public OperationData(
String sessionUser,
String sessionType,
String kyuubiInstance,
Map<String, String> metrics) {
Map<String, String> metrics,
OperationProgress progress) {
this.identifier = identifier;
this.remoteId = remoteId;
this.statement = statement;
Expand All @@ -67,6 +69,7 @@ public OperationData(
this.sessionType = sessionType;
this.kyuubiInstance = kyuubiInstance;
this.metrics = metrics;
this.progress = progress;
}

public String getIdentifier() {
Expand Down Expand Up @@ -176,6 +179,14 @@ public void setMetrics(Map<String, String> metrics) {
this.metrics = metrics;
}

public OperationProgress getProgress() {
return progress;
}

public void setProgress(OperationProgress progress) {
this.progress = progress;
}

@Override
public boolean equals(Object o) {
if (this == o) return true;
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,133 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package org.apache.kyuubi.client.api.v1.dto;

import java.util.Collections;
import java.util.List;
import java.util.Objects;
import org.apache.commons.lang3.builder.ReflectionToStringBuilder;
import org.apache.commons.lang3.builder.ToStringStyle;

public class OperationProgress {
private List<String> headerNames;
private List<List<String>> rows;
private double progressedPercentage;
private String status;
private String footerSummary;
private long startTime;

public OperationProgress() {}

public OperationProgress(
List<String> headerNames,
List<List<String>> rows,
double progressedPercentage,
String status,
String footerSummary,
long startTime) {
this.headerNames = headerNames;
this.rows = rows;
this.progressedPercentage = progressedPercentage;
this.status = status;
this.footerSummary = footerSummary;
this.startTime = startTime;
}

public List<String> getHeaderNames() {
if (headerNames == null) {
return Collections.emptyList();
}
return headerNames;
}

public void setHeaderNames(List<String> headerNames) {
this.headerNames = headerNames;
}

public List<List<String>> getRows() {
if (rows == null) {
return Collections.emptyList();
}
return rows;
}

public void setRows(List<List<String>> rows) {
this.rows = rows;
}

public double getProgressedPercentage() {
return progressedPercentage;
}

public void setProgressedPercentage(double progressedPercentage) {
this.progressedPercentage = progressedPercentage;
}

public String getStatus() {
return status;
}

public void setStatus(String status) {
this.status = status;
}

public String getFooterSummary() {
return footerSummary;
}

public void setFooterSummary(String footerSummary) {
this.footerSummary = footerSummary;
}

public long getStartTime() {
return startTime;
}

public void setStartTime(long startTime) {
this.startTime = startTime;
}

@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
OperationProgress that = (OperationProgress) o;
return Double.compare(getProgressedPercentage(), that.getProgressedPercentage()) == 0
&& getStartTime() == that.getStartTime()
&& Objects.equals(getHeaderNames(), that.getHeaderNames())
&& Objects.equals(getRows(), that.getRows())
&& Objects.equals(getStatus(), that.getStatus())
&& Objects.equals(getFooterSummary(), that.getFooterSummary());
}

@Override
public int hashCode() {
return Objects.hash(
getHeaderNames(),
getRows(),
getProgressedPercentage(),
getStatus(),
getFooterSummary(),
getStartTime());
}

@Override
public String toString() {
return ReflectionToStringBuilder.toString(this, ToStringStyle.JSON_STYLE);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,35 @@ package org.apache.kyuubi.server.api
import scala.collection.JavaConverters._

import org.apache.kyuubi.{Logging, Utils}
import org.apache.kyuubi.client.api.v1.dto.{OperationData, ServerData, SessionData}
import org.apache.kyuubi.client.api.v1.dto
import org.apache.kyuubi.client.api.v1.dto.{OperationData, OperationProgress, ServerData, SessionData}
import org.apache.kyuubi.events.KyuubiOperationEvent
import org.apache.kyuubi.ha.client.ServiceNodeInfo
import org.apache.kyuubi.operation.KyuubiOperation
import org.apache.kyuubi.session.KyuubiSession

object ApiUtils extends Logging {
def sessionEvent(session: KyuubiSession): dto.KyuubiSessionEvent = {
session.getSessionEvent.map(event =>
dto.KyuubiSessionEvent.builder()
.sessionId(event.sessionId)
.clientVersion(event.clientVersion)
.sessionType(event.sessionType)
.sessionName(event.sessionName)
.user(event.user)
.clientIp(event.clientIP)
.serverIp(event.serverIP)
.conf(event.conf.asJava)
.remoteSessionId(event.remoteSessionId)
.engineId(event.engineId)
.eventTime(event.eventTime)
.openedTime(event.openedTime)
.startTime(event.startTime)
.endTime(event.endTime)
.totalOperations(event.totalOperations)
.exception(event.exception.orNull)
.build()).orNull
}

def sessionData(session: KyuubiSession): SessionData = {
val sessionEvent = session.getSessionEvent
Expand All @@ -45,6 +67,40 @@ object ApiUtils extends Logging {
sessionEvent.map(_.engineId).getOrElse(""))
}

private def operationProgress(operation: KyuubiOperation): OperationProgress = {
Option(operation.getOperationJobProgress).map { jobProgress =>
new OperationProgress(
jobProgress.getHeaderNames,
jobProgress.getRows,
jobProgress.getProgressedPercentage,
jobProgress.getStatus.toString,
jobProgress.getFooterSummary,
jobProgress.getStartTime)
}.orNull
}

def operationEvent(operation: KyuubiOperation): dto.KyuubiOperationEvent = {
val opEvent = KyuubiOperationEvent(operation)
dto.KyuubiOperationEvent.builder()
.statementId(opEvent.statementId)
.remoteId(opEvent.remoteId)
.statement(opEvent.statement)
.shouldRunAsync(opEvent.shouldRunAsync)
.state(opEvent.state)
.eventTime(opEvent.eventTime)
.createTime(opEvent.createTime)
.startTime(opEvent.startTime)
.completeTime(opEvent.completeTime)
.exception(opEvent.exception.orNull)
.sessionId(opEvent.sessionId)
.sessionUser(opEvent.sessionUser)
.sessionType(opEvent.sessionType)
.kyuubiInstance(opEvent.kyuubiInstance)
.metrics(opEvent.metrics.asJava)
.progress(operationProgress(operation))
.build()
}

def operationData(operation: KyuubiOperation): OperationData = {
val opEvent = KyuubiOperationEvent(operation)
new OperationData(
Expand All @@ -60,7 +116,8 @@ object ApiUtils extends Logging {
opEvent.sessionUser,
opEvent.sessionType,
operation.getSession.asInstanceOf[KyuubiSession].connectionUrl,
operation.metrics.asJava)
operation.metrics.asJava,
operationProgress(operation))
}

def serverData(nodeInfo: ServiceNodeInfo): ServerData = {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@ import io.swagger.v3.oas.annotations.tags.Tag

import org.apache.kyuubi.{KyuubiSQLException, Logging}
import org.apache.kyuubi.client.api.v1.dto._
import org.apache.kyuubi.events.KyuubiOperationEvent
import org.apache.kyuubi.operation.{FetchOrientation, KyuubiOperation, OperationHandle}
import org.apache.kyuubi.server.api.{ApiRequestContext, ApiUtils}
import org.apache.kyuubi.shaded.hive.service.rpc.thrift._
Expand All @@ -54,7 +53,7 @@ private[v1] class OperationsResource extends ApiRequestContext with Logging {
try {
val opHandle = OperationHandle(operationHandleStr)
val operation = fe.be.sessionManager.operationManager.getOperation(opHandle)
KyuubiOperationEvent(operation.asInstanceOf[KyuubiOperation])
ApiUtils.operationEvent(operation.asInstanceOf[KyuubiOperation])
} catch {
case NonFatal(e) =>
val errorMsg = "Error getting an operation event"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -69,26 +69,7 @@ private[v1] class SessionsResource extends ApiRequestContext with Logging {
@Path("{sessionHandle}")
def sessionInfo(@PathParam("sessionHandle") sessionHandleStr: String): dto.KyuubiSessionEvent = {
try {
sessionManager.getSession(sessionHandleStr)
.asInstanceOf[KyuubiSession].getSessionEvent.map(event =>
dto.KyuubiSessionEvent.builder
.sessionId(event.sessionId)
.clientVersion(event.clientVersion)
.sessionType(event.sessionType)
.sessionName(event.sessionName)
.user(event.user)
.clientIp(event.clientIP)
.serverIp(event.serverIP)
.conf(event.conf.asJava)
.remoteSessionId(event.remoteSessionId)
.engineId(event.engineId)
.eventTime(event.eventTime)
.openedTime(event.openedTime)
.startTime(event.startTime)
.endTime(event.endTime)
.totalOperations(event.totalOperations)
.exception(event.exception.orNull)
.build).get
ApiUtils.sessionEvent(sessionManager.getSession(sessionHandleStr).asInstanceOf[KyuubiSession])
} catch {
case NonFatal(e) =>
val errorMsg = s"Invalid $sessionHandleStr"
Expand Down
Loading

0 comments on commit bd379c5

Please sign in to comment.