Skip to content

Add timeout StatementState #2724

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ykmr1224 merged 3 commits into from
Jun 7, 2024
Merged

Add timeout StatementState #2724

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d48630e
Add timeout StatementState
ykmr1224 Jun 6, 2024
e661352
Fix code style
ykmr1224 Jun 6, 2024
f4e8768
Fix coverage
ykmr1224 Jun 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions spark/src/main/java/org/opensearch/sql/spark/execution/statement/Statement.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ public void cancel() {

if (statementState.equals(StatementState.SUCCESS)
|| statementState.equals(StatementState.FAILED)
|| statementState.equals(StatementState.TIMEOUT)
|| statementState.equals(StatementState.CANCELLED)) {
String errorMsg =
String.format(
Expand Down
1 change: 1 addition & 0 deletions spark/src/main/java/org/opensearch/sql/spark/execution/statement/StatementState.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ public enum StatementState {
RUNNING("running"),
SUCCESS("success"),
FAILED("failed"),
TIMEOUT("timeout"),
CANCELLED("cancelled");

private final String state;
Expand Down
53 changes: 16 additions & 37 deletions spark/src/test/java/org/opensearch/sql/spark/execution/statement/StatementTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,7 @@ public void openFailedBecauseConflict() {
}

@Test
public void cancelNotExistStatement() {
public void cancelNotExistStatement_throwsException() {
StatementId stId = new StatementId("statementId");
Statement st = buildStatement(stId);
st.open();
Expand All @@ -141,8 +141,6 @@ public void cancelFailedBecauseOfConflict() {
statementStorageService.updateStatementState(st.getStatementModel(), CANCELLED);

assertEquals(StatementState.CANCELLED, running.getStatementState());

// cancel conflict
IllegalStateException exception = assertThrows(IllegalStateException.class, st::cancel);
assertEquals(
String.format(
Expand All @@ -151,55 +149,36 @@ public void cancelFailedBecauseOfConflict() {
}

@Test
public void cancelSuccessStatementFailed() {
StatementId stId = new StatementId("statementId");
Statement st = createStatement(stId);

// update to running state
StatementModel model = st.getStatementModel();
st.setStatementModel(
StatementModel.copyWithState(
st.getStatementModel(), StatementState.SUCCESS, model.getMetadata()));

// cancel conflict
IllegalStateException exception = assertThrows(IllegalStateException.class, st::cancel);
assertEquals(
String.format("can't cancel statement in success state. statement: %s.", stId),
exception.getMessage());
public void cancelCancelledStatement_throwsException() {
testCancelThrowsExceptionGivenStatementState(StatementState.CANCELLED);
}

@Test
public void cancelFailedStatementFailed() {
StatementId stId = new StatementId("statementId");
Statement st = createStatement(stId);

// update to running state
StatementModel model = st.getStatementModel();
st.setStatementModel(
StatementModel.copyWithState(
st.getStatementModel(), StatementState.FAILED, model.getMetadata()));
public void cancelSuccessStatement_throwsException() {
testCancelThrowsExceptionGivenStatementState(StatementState.SUCCESS);
}

// cancel conflict
IllegalStateException exception = assertThrows(IllegalStateException.class, st::cancel);
assertEquals(
String.format("can't cancel statement in failed state. statement: %s.", stId),
exception.getMessage());
@Test
public void cancelFailedStatement_throwsException() {
testCancelThrowsExceptionGivenStatementState(StatementState.FAILED);
}

@Test
public void cancelCancelledStatementFailed() {
public void cancelTimeoutStatement_throwsException() {
testCancelThrowsExceptionGivenStatementState(StatementState.TIMEOUT);
}

private void testCancelThrowsExceptionGivenStatementState(StatementState state) {
StatementId stId = new StatementId("statementId");
Statement st = createStatement(stId);

// update to running state
StatementModel model = st.getStatementModel();
st.setStatementModel(
StatementModel.copyWithState(st.getStatementModel(), CANCELLED, model.getMetadata()));
StatementModel.copyWithState(st.getStatementModel(), state, model.getMetadata()));

// cancel conflict
IllegalStateException exception = assertThrows(IllegalStateException.class, st::cancel);
assertEquals(
String.format("can't cancel statement in cancelled state. statement: %s.", stId),
String.format("can't cancel statement in %s state. statement: %s.", state.getState(), stId),
exception.getMessage());
}

Expand Down