Skip to content

Commit 20664c6

Browse files
Zxc4wonderZxc4wonder
committed
fix(console): preserve execution and audit state
1 parent e02e92f commit 20664c6

6 files changed

Lines changed: 49 additions & 8 deletions

File tree

‎backend/framework/src/main/java/org/jumpserver/chen/framework/console/QueryConsole.java‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,10 @@ public class QueryConsole extends AbstractConsole {
7575
static final String QUERY_INSERT_NOT_SUPPORTED = "QUERY_INSERT_NOT_SUPPORTED";
7676
static final String QUERY_DELETE_NOT_SUPPORTED = "QUERY_DELETE_NOT_SUPPORTED";
7777
static final String CONSOLE_DATA_VIEW_EDIT_NOT_SUPPORTED = "CONSOLE_DATA_VIEW_EDIT_NOT_SUPPORTED";
78+
private static final String EXECUTION_STATUS_RUNNING = "running";
79+
private static final String EXECUTION_STATUS_SUCCESS = "success";
80+
private static final String EXECUTION_STATUS_ERROR = "error";
81+
private static final String EXECUTION_STATUS_CANCELLED = "cancelled";
7882

7983
private final Datasource datasource;
8084
private final boolean consoleMode;
@@ -453,7 +457,7 @@ private void onDataViewAction(DataViewAction action) {
453457
}
454458
try {
455459
var context = this.tableEditContextFactory.create(dataView, this.getDatasource().getDruidDbType());
456-
var result = this.tableChangesPreviewService.preview(context, action.getDataView(), request);
460+
var result = this.tableChangesPreviewService.preview(context, dataView.getTitle(), request);
457461
this.getPacketIO().sendPacket(PACKET_SAVE_CHANGES_PREVIEW_RESULT, result);
458462
} catch (IllegalArgumentException e) {
459463
this.getPacketIO().sendPacket(PACKET_SAVE_CHANGES_PREVIEW_RESULT, this.rejectedPreview(dataView, e.getMessage()));
@@ -462,7 +466,7 @@ private void onDataViewAction(DataViewAction action) {
462466
}
463467
if (DataViewAction.ACTION_SAVE_CHANGES.equals(action.getAction())) {
464468
var request = GSON.fromJson(GSON.toJson(action.getData()), SaveChangesRequest.class);
465-
SaveChangesResult result = this.saveQueryChanges(dataView, action.getDataView(), request);
469+
SaveChangesResult result = this.saveQueryChanges(dataView, dataView.getTitle(), request);
466470
this.getPacketIO().sendPacket(PACKET_SAVE_CHANGES_RESULT, result);
467471
return;
468472
}
@@ -610,6 +614,7 @@ private SaveChangesResult rejectedSave(DataView dataView, String reason) {
610614
}
611615

612616
public void onCancel() {
617+
this.getState().setExecutionStatus(EXECUTION_STATUS_CANCELLED);
613618
try {
614619
var plan = this.currentPlan;
615620
if (plan != null && plan.getStatement() != null) {
@@ -740,6 +745,7 @@ private Path resolveSQLFileInSessionTemp(String filename) {
740745

741746
public void onSQL(String sql) {
742747
this.getState().setInQuery(true);
748+
this.getState().setExecutionStatus(EXECUTION_STATUS_RUNNING);
743749
this.stateManager.commit();
744750
var session = SessionManager.getCurrentSession();
745751

@@ -763,12 +769,19 @@ public void onSQL(String sql) {
763769
}
764770
this.ensureCurrentSchema();
765771
} catch (ParserException e) {
772+
this.getState().setExecutionStatus(EXECUTION_STATUS_ERROR);
766773
this.getConsoleLogger().error("%s: %s", MessageUtils.get("ParseError"), e.getMessage());
767774
this.getPacketIO().sendPacket("message", Message.error(MessageUtils.get("ParseError"), e.getMessage()));
768775
} catch (SQLException e) {
776+
if (!StringUtils.equals(this.getState().getExecutionStatus(), EXECUTION_STATUS_CANCELLED)) {
777+
this.getState().setExecutionStatus(EXECUTION_STATUS_ERROR);
778+
}
769779
this.getConsoleLogger().error("%s: %s", MessageUtils.get("ExecuteError"), e.getMessage());
770780
this.getPacketIO().sendPacket("message", Message.error(MessageUtils.get("ExecuteError"), e.getMessage()));
771781
} finally {
782+
if (StringUtils.equals(this.getState().getExecutionStatus(), EXECUTION_STATUS_RUNNING)) {
783+
this.getState().setExecutionStatus(EXECUTION_STATUS_SUCCESS);
784+
}
772785
this.getState().setInQuery(false);
773786
this.getState().setCanCancel(false);
774787
this.stateManager.commit();
@@ -782,6 +795,11 @@ private boolean canExecuteStatement(Session session, String sql, ACLResult aclRe
782795
if (aclResult.getRiskLevel() == Common.RiskLevel.Reject ||
783796
aclResult.getRiskLevel() == Common.RiskLevel.ReviewReject ||
784797
aclResult.getRiskLevel() == Common.RiskLevel.ReviewCancel) {
798+
this.getState().setExecutionStatus(
799+
aclResult.getRiskLevel() == Common.RiskLevel.ReviewCancel
800+
? EXECUTION_STATUS_CANCELLED
801+
: EXECUTION_STATUS_ERROR
802+
);
785803
this.getConsoleLogger().error("%s", MessageUtils.get("ACLRejectError"));
786804
CommandRecord commandRecord = new CommandRecord(sql);
787805
commandRecord.setRiskLevel(aclResult.getRiskLevel());
@@ -800,6 +818,7 @@ private boolean confirmStatementWarning(Session session) {
800818
dialog.addButton(new Button(MessageUtils.get("Submit"), "submit", countDownLatch::countDown));
801819
dialog.addButton(new Button(MessageUtils.get("Cancel"), "cancel", () -> {
802820
hasNext.set(false);
821+
this.getState().setExecutionStatus(EXECUTION_STATUS_CANCELLED);
803822
countDownLatch.countDown();
804823
this.getConsoleLogger().warn(MessageUtils.get("ExecutionCanceled"));
805824
}));
@@ -813,12 +832,14 @@ private boolean confirmStatementWarning(Session session) {
813832
try {
814833
if (!countDownLatch.await(WARNING_DIALOG_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
815834
hasNext.set(false);
835+
this.getState().setExecutionStatus(EXECUTION_STATUS_CANCELLED);
816836
dialogHandle.cancel();
817837
this.getConsoleLogger().warn(MessageUtils.get("ExecutionCanceled"));
818838
}
819839
return hasNext.get();
820840
} catch (InterruptedException e) {
821841
Thread.currentThread().interrupt();
842+
this.getState().setExecutionStatus(EXECUTION_STATUS_ERROR);
822843
this.getConsoleLogger().error("获取结果失败!");
823844
return false;
824845
} finally {

‎backend/framework/src/main/java/org/jumpserver/chen/framework/console/state/QueryConsoleState.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ public class QueryConsoleState extends State {
1515
private int timeout;
1616
private boolean editorLoading;
1717
private boolean canCancel;
18+
private volatile String executionStatus;
1819

1920
public QueryConsoleState(String title) {
2021
super(title);
@@ -23,6 +24,7 @@ public QueryConsoleState(String title) {
2324
this.timeout = 30;
2425
this.editorLoading = false;
2526
this.canCancel = false;
27+
this.executionStatus = "";
2628
}
2729

2830
}

‎backend/framework/src/main/java/org/jumpserver/chen/framework/session/Session.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,8 @@ public interface Session {
8080

8181
boolean isActive();
8282

83+
boolean isClosing();
84+
8385
void close();
8486

8587
void close(String message, Object... args);

‎backend/framework/src/main/java/org/jumpserver/chen/framework/session/SessionManager.java‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import java.util.Map;
77
import java.util.UUID;
88
import java.util.concurrent.ConcurrentHashMap;
9+
import java.util.concurrent.atomic.AtomicBoolean;
910

1011

1112
@Slf4j
@@ -33,10 +34,17 @@ public static void unregisterSession(String token) {
3334
}
3435

3536
public static boolean registerConsole(String token, String consoleId, Console console) {
36-
return instance.store.computeIfPresent(token, (ignored, session) -> {
37-
session.getConsoles().put(consoleId, console);
37+
var registered = new AtomicBoolean(false);
38+
instance.store.computeIfPresent(token, (ignored, session) -> {
39+
synchronized (session) {
40+
if (!session.isClosing()) {
41+
session.getConsoles().put(consoleId, console);
42+
registered.set(true);
43+
}
44+
}
3845
return session;
39-
}) != null;
46+
});
47+
return registered.get();
4048
}
4149

4250
public int getCurrentSessionCount() {

‎backend/framework/src/main/java/org/jumpserver/chen/framework/session/impl/BaseSession.java‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,11 @@ public boolean isActive() {
202202
return this.getPacketIO() != null && this.getPacketIO().getWsSession().isOpen();
203203
}
204204

205+
@Override
206+
public boolean isClosing() {
207+
return this.closeStarted.get();
208+
}
209+
205210
@Override
206211
public void close() {
207212
if (!this.beginClose()) {
@@ -211,15 +216,17 @@ public void close() {
211216
}
212217

213218
protected final boolean beginClose() {
214-
return this.closeStarted.compareAndSet(false, true);
219+
synchronized (this) {
220+
return this.closeStarted.compareAndSet(false, true);
221+
}
215222
}
216223

217224
protected final void closeSessionResources() {
218225
if (this.getController() != null) {
219226
this.getController().cancelAllDialogs();
220227
}
221-
SessionManager.unregisterSession(this.getWebToken());
222228
this.closeConsoles();
229+
SessionManager.unregisterSession(this.getWebToken());
223230
this.getDatasource().close();
224231
this.getPacketIO().close();
225232
var path = this.getTempPath();
@@ -228,7 +235,7 @@ protected final void closeSessionResources() {
228235
}
229236
}
230237

231-
private void closeConsoles() {
238+
protected final void closeConsoles() {
232239
var detached = new ArrayList<Map.Entry<String, Console>>();
233240
while (!this.consoles.isEmpty()) {
234241
for (var entry : this.consoles.entrySet()) {

‎backend/framework/src/main/java/org/jumpserver/chen/framework/session/impl/JMSSession.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -277,6 +277,7 @@ private void closeJmsSessionResources() {
277277
if (this.getController() != null) {
278278
this.getController().cancelAllDialogs();
279279
}
280+
this.closeConsoles();
280281
try {
281282
this.replayHandler.release();
282283
this.finishedJmsSession();

0 commit comments

Comments
 (0)