@OnMessage(maxMessageSize = 1024000)
    public void onMessage(String message, Session session) {
        CompletableFuture.runAsync(() -> {
            WsMessage wsMessage = JSONObject.parseObject(message, WsMessage.class);
            // Process your messages here
            try {
                String actionType = wsMessage.getActionType();
                if (WsMessage.ActionType.PING.equalsIgnoreCase(actionType)) {
                    WsResult wsResult = new WsResult();
                    ActionResult actionResult = new ActionResult();
                    actionResult.setSuccess(true);
                    wsResult.setActionType(WsMessage.ActionType.PING);
                    wsResult.setUuid(wsMessage.getUuid());
                    wsResult.setMessage(actionResult);
                    SendMessage(session, wsResult);
                    timer.cancel();
                    heartBeat(session);
                } else {
                    ContextUtils.setContext(Context.builder()
                            .loginUser(loginUser)
                            .build());
                    Dbutils.setSession();
                    JSONObject jsonObject = wsMessage.getMessage();
                    Long dataSourceId = jsonObject.getLong("dataSourceId");
                    String databaseName = jsonObject.getString("databaseName");
                    String schemaName = jsonObject.getString("schemaName");
                    Long consoleId = jsonObject.getLong("consoleId");
                    String key = connectInfoKey(dataSourceId, databaseName, schemaName, consoleId);
                    ConnectInfo connectInfo = connectInfoMap.get(key);
                    if (connectInfo == null) {
                        connectInfo = wsService.toInfo(dataSourceId, databaseName, consoleId, schemaName);
                        connectInfoMap.put(key, connectInfo);
                    }
                    Chat2DBContext.putContext(connectInfo);
                    if (WsMessage.ActionType.EXECUTE.equalsIgnoreCase(actionType)) {
                        DmlRequest request = jsonObject.toJavaObject(DmlRequest.class);
                        ListResult<ExecuteResultVO> result = wsService.execute(request);
                        WsResult resultMessage = new WsResult();
                        resultMessage.setUuid(wsMessage.getUuid());
                        resultMessage.setActionType(wsMessage.getActionType());
                        resultMessage.setMessage(result);
                        SendMessage(session, resultMessage);
                    }
                }
            } catch (Exception e) {
                WsResult wsResult = new WsResult();
                ActionResult actionResult = new ActionResult();
                actionResult.setSuccess(false);
                actionResult.setErrorCode(e.getMessage());
                wsResult.setActionType(WsMessage.ActionType.ERROR);
                wsResult.setUuid(wsMessage.getUuid());
                wsResult.setMessage(actionResult);
                SendMessage(session, wsResult);
            } finally {
                Chat2DBContext.removeContext();
                ContextUtils.removeContext();
                Dbutils.removeSession();
            }
        });

    }
Method called after receiving client message

@param message message sent by the client
chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/ws/WsServer.java