From 663dfcc31432b87ec611193559e43ffb7d2eabf3 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Wed, 17 Apr 2024 19:24:02 +0800 Subject: [PATCH 01/16] add AsyncExecuteContext for sql async execution --- .../session/ConnectionSessionConstants.java | 2 + .../core/session/ConnectionSessionUtil.java | 50 ++++ .../oceanbase/odc/core/sql/util/OBUtils.java | 61 +++++ .../v2/ConnectSessionController.java | 15 ++ .../connection/util/ConnectionInfoUtil.java | 23 ++ .../session/ConnectConsoleService.java | 123 +++++++++- .../service/session/OdcStatementCallBack.java | 30 ++- .../session/model/AsyncExecuteContext.java | 217 ++++++++++++++++++ .../session/model/AsyncExecuteResultResp.java | 43 ++++ 9 files changed, 560 insertions(+), 4 deletions(-) create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java diff --git a/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionConstants.java b/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionConstants.java index 1f0c90c330..81884b6052 100644 --- a/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionConstants.java +++ b/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionConstants.java @@ -80,11 +80,13 @@ public class ConnectionSessionConstants { */ public static final String QUERY_CACHE_KEY = "QUERY_CACHE"; public static final String FUTURE_JDBC_RESULT_KEY = "FUTURE_JDBC_RESULT"; + public static final String ASYNC_EXECUTE_CONTEXT_KEY = "ASYNC_EXECUTE_CONTEXT"; /** * The connection_id current database session needs to be stored in the database session in the form * of attributes, this is the key */ public static final String CONNECTION_ID_KEY = "CONNECTION_ID"; + public static final String CONNECTION_PROXY_SESSID_KEY = "PROXY_SESSID"; /** * The {@link BinaryDataManager} session needs to be stored in the database session in the form of * attributes, this is the key diff --git a/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionUtil.java b/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionUtil.java index 77d599262e..34c88167d1 100644 --- a/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionUtil.java +++ b/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionUtil.java @@ -200,6 +200,52 @@ public static void removeFutureJdbc(@NonNull ConnectionSession connectionSes id2FutureResult.remove(requestId); } + @SuppressWarnings("all") + public static String setExecuteContext(@NonNull ConnectionSession connectionSession, + @NonNull Object context) { + Object value = connectionSession.getAttribute(ConnectionSessionConstants.ASYNC_EXECUTE_CONTEXT_KEY); + Map id2ExecuteContext; + if (value instanceof Map) { + id2ExecuteContext = (Map) value; + } else { + id2ExecuteContext = new ConcurrentHashMap<>(); + connectionSession.setAttribute(ConnectionSessionConstants.ASYNC_EXECUTE_CONTEXT_KEY, id2ExecuteContext); + } + String uuid = UUID.randomUUID().toString(); + id2ExecuteContext.putIfAbsent(uuid, context); + return uuid; + } + + @SuppressWarnings("all") + public static Object getExecuteContext(@NonNull ConnectionSession connectionSession, + @NonNull String requestId) { + Object value = connectionSession.getAttribute(ConnectionSessionConstants.ASYNC_EXECUTE_CONTEXT_KEY); + Map id2ExecuteContext; + if (value instanceof Map) { + id2ExecuteContext = (Map) value; + } else { + throw new NotFoundException(ResourceType.ODC_ASYNC_SQL_RESULT, "session id", connectionSession.getId()); + } + Object context = id2ExecuteContext.get(requestId); + if (context == null) { + throw new NotFoundException(ResourceType.ODC_ASYNC_SQL_RESULT, "request id", requestId); + } + return context; + } + + @SuppressWarnings("all") + public static void removeExecuteContext(@NonNull ConnectionSession connectionSession, + @NonNull String requestId) { + Object value = connectionSession.getAttribute(ConnectionSessionConstants.ASYNC_EXECUTE_CONTEXT_KEY); + Map id2ExecuteContext; + if (!(value instanceof Map)) { + throw new NullPointerException("Result not found by session id " + connectionSession.getId()); + } else { + id2ExecuteContext = (Map) value; + } + id2ExecuteContext.remove(requestId); + } + public static void setUserId(@NonNull ConnectionSession connectionSession, @NonNull Long userId) { connectionSession.setAttribute(ConnectionSessionConstants.USER_ID_KEY, userId); } @@ -522,6 +568,10 @@ public static String getConsoleConnectionId(@NonNull ConnectionSession connectio return (String) connectionSession.getAttribute(ConnectionSessionConstants.CONNECTION_ID_KEY); } + public static String getConsoleConnectionProxySessId(@NonNull ConnectionSession connectionSession) { + return (String) connectionSession.getAttribute(ConnectionSessionConstants.CONNECTION_PROXY_SESSID_KEY); + } + public static String getVersion(@NonNull ConnectionSession connectionSession) { return (String) connectionSession.getAttribute(ConnectionSessionConstants.OB_VERSION); } diff --git a/server/odc-core/src/main/java/com/oceanbase/odc/core/sql/util/OBUtils.java b/server/odc-core/src/main/java/com/oceanbase/odc/core/sql/util/OBUtils.java index fc794bf869..31ef191eab 100644 --- a/server/odc-core/src/main/java/com/oceanbase/odc/core/sql/util/OBUtils.java +++ b/server/odc-core/src/main/java/com/oceanbase/odc/core/sql/util/OBUtils.java @@ -21,6 +21,8 @@ import java.sql.SQLException; import java.sql.Statement; import java.sql.Types; +import java.util.ArrayList; +import java.util.List; import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.Validate; @@ -380,4 +382,63 @@ public static String queryDBMSOutput(@NonNull Connection connection, Integer max } } + public static List querySessionIdsByProxySessId(@NonNull Statement statement, + @NonNull String proxySessId, ConnectType connectType) throws SQLException { + DialectType dialectType = connectType.getDialectType(); + SqlBuilder sqlBuilder = getBuilder(connectType) + .append("select id from ") + .append(dialectType.isMysql() ? "oceanbase" : "sys") + .append(".v$ob_processlist where proxy_sessid = ") + .append(proxySessId); + List ids = new ArrayList<>(); + try (ResultSet rs = statement.executeQuery(sqlBuilder.toString())) { + while (rs.next()) { + ids.add(rs.getString(1)); + } + } + return ids; + } + + /** + * OceanBase only supports ASH views in versions higher than 4.0. Therefore, this method is not + * applicable to earlier versions, please use sql_audit instead. + */ + public static String queryTraceIdFromASH(@NonNull Statement statement, + @NonNull List sessionIds, ConnectType connectType) throws SQLException { + DialectType dialectType = connectType.getDialectType(); + SqlBuilder sqlBuilder = getBuilder(connectType) + .append("select trace_id from ") + .append(dialectType.isMysql() ? "oceanbase" : "sys") + .append(".v$active_session_history where session_id in (") + .append(String.join(",", sessionIds)) + .append(")") + .append(dialectType.isMysql() ? " limit 1" : " and rownum=1"); + try (ResultSet rs = statement.executeQuery(sqlBuilder.toString())) { + if (!rs.next()) { + throw new SQLException("No result found in ASH."); + } + return rs.getString(1); + } + } + + /** + * OceanBase only supports ASH views in versions higher than 4.0. Therefore, this method is not + * applicable to earlier versions, please use sql_audit instead. + */ + public static String queryPlanIdByTraceId(@NonNull Statement statement, String traceId, ConnectType connectType) + throws SQLException { + DialectType dialectType = connectType.getDialectType(); + SqlBuilder sqlBuilder = getBuilder(connectType) + .append("select plan_id from ") + .append(dialectType.isMysql() ? "oceanbase" : "sys") + .append(".active_session_history where trace_id=") + .value(traceId); + try (ResultSet rs = statement.executeQuery(sqlBuilder.toString())) { + if (!rs.next()) { + throw new SQLException("No result found in ASH."); + } + return rs.getString(1); + } + } + } diff --git a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/ConnectSessionController.java b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/ConnectSessionController.java index 2002f61ddf..b886ee6e30 100644 --- a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/ConnectSessionController.java +++ b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/ConnectSessionController.java @@ -55,6 +55,7 @@ import com.oceanbase.odc.service.partitionplan.model.PartitionPlanPreviewReq; import com.oceanbase.odc.service.session.ConnectConsoleService; import com.oceanbase.odc.service.session.ConnectSessionService; +import com.oceanbase.odc.service.session.model.AsyncExecuteResultResp; import com.oceanbase.odc.service.session.model.BinaryContent; import com.oceanbase.odc.service.session.model.QueryTableOrViewDataReq; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteReq; @@ -118,6 +119,13 @@ public SuccessResponse asyncSqlExecute(@PathVariable String return Responses.success(consoleService.execute(SidUtils.getSessionId(sessionId), req)); } + @RequestMapping(value = {"/sessions/{sessionId}/sqls/asyncExecuteV2"}, method = RequestMethod.POST) + @StatefulRoute(stateName = StateName.DB_SESSION, stateIdExpression = "#sessionId") + public SuccessResponse asyncSqlExecuteV2(@PathVariable String sessionId, + @RequestBody SqlAsyncExecuteReq req) throws Exception { + return Responses.success(consoleService.executeV2(SidUtils.getSessionId(sessionId), req, true)); + } + /** * 获取异步执行sql的结果 Todo 这里的sqlIds后续需要改成一个string类型的requestId，异步api请求需要有超时机制 * @@ -132,6 +140,13 @@ public SuccessResponse> getAsyncSqlExecute(@PathVariable return Responses.success(consoleService.getAsyncResult(SidUtils.getSessionId(sessionId), requestId, null)); } + @RequestMapping(value = "/sessions/{sessionId}/sqls/getResultV2", method = RequestMethod.GET) + @StatefulRoute(stateName = StateName.DB_SESSION, stateIdExpression = "#sessionId") + public SuccessResponse getAsyncSqlExecuteV2(@PathVariable String sessionId, + @RequestParam String requestId) { + return Responses.success(consoleService.getAsyncResultV2(SidUtils.getSessionId(sessionId), requestId, null)); + } + /** * 对 sql 脚本的内容进行静态检查 * diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java index 48d59042cb..9b31b99adb 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java @@ -15,6 +15,7 @@ */ package com.oceanbase.odc.service.connection.util; +import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Statement; @@ -23,9 +24,11 @@ import org.springframework.jdbc.core.ConnectionCallback; import org.springframework.jdbc.core.StatementCallback; +import com.oceanbase.odc.common.util.VersionUtils; import com.oceanbase.odc.core.datasource.DataSourceFactory; import com.oceanbase.odc.core.session.ConnectionSession; import com.oceanbase.odc.core.session.ConnectionSessionConstants; +import com.oceanbase.odc.core.session.ConnectionSessionUtil; import com.oceanbase.odc.core.shared.Verify; import com.oceanbase.odc.core.shared.constant.DialectType; import com.oceanbase.odc.core.sql.execute.GeneralSyncJdbcExecutor; @@ -61,6 +64,11 @@ public static void initConnectionId(@NonNull Statement statement, @NonNull Conne String sessionId = queryConnectionId(statement, connectionSession.getDialectType()); Verify.notNull(sessionId, "SessionId"); connectionSession.setAttribute(ConnectionSessionConstants.CONNECTION_ID_KEY, sessionId); + if (connectionSession.getDialectType().isOceanbase() && VersionUtils.isGreaterThanOrEqualsTo( + ConnectionSessionUtil.getVersion(connectionSession), "4.2")) { + String proxySessId = queryProxySessId(statement, connectionSession.getDialectType(), sessionId); + connectionSession.setAttribute(ConnectionSessionConstants.CONNECTION_PROXY_SESSID_KEY, proxySessId); + } } catch (Exception exception) { log.warn("Failed to get database session ID, session={}", connectionSession, exception); } @@ -71,6 +79,21 @@ public static String queryConnectionId(@NonNull Statement statement, @NonNull Di return ConnectionPluginUtil.getSessionExtension(dialectType).getConnectionId(statement.getConnection()); } + public static String queryProxySessId(@NonNull Statement statement, @NonNull DialectType dialectType, + @NonNull String connectionId) throws SQLException { + String proxySessId = null; + String sql = "select proxy_sessid from " + + (dialectType.isMysql() ? "oceanbase" : "sys") + + ".v$ob_processlist where id = " + + connectionId; + try (ResultSet rs = statement.executeQuery(sql)) { + if (rs.next()) { + proxySessId = rs.getString(1); + } + } + return proxySessId; + } + public static void initSessionVersion(@NonNull ConnectionSession connectionSession) { InformationExtensionPoint point = ConnectionPluginUtil.getInformationExtension(connectionSession.getDialectType()); diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java index 331191d38f..41045d80a3 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java @@ -15,6 +15,8 @@ */ package com.oceanbase.odc.service.session; +import static com.oceanbase.odc.service.session.model.AsyncExecuteContext.SHOW_TABLE_COLUMN_INFO; + import java.io.IOException; import java.io.InputStream; import java.util.Arrays; @@ -24,6 +26,8 @@ import java.util.Map; import java.util.Objects; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; @@ -72,6 +76,7 @@ import com.oceanbase.odc.core.sql.execute.cache.table.VirtualElement; import com.oceanbase.odc.core.sql.execute.cache.table.VirtualTable; import com.oceanbase.odc.core.sql.execute.model.JdbcGeneralResult; +import com.oceanbase.odc.core.sql.execute.model.SqlExecuteStatus; import com.oceanbase.odc.core.sql.execute.model.SqlTuple; import com.oceanbase.odc.core.sql.parser.AbstractSyntaxTree; import com.oceanbase.odc.core.sql.parser.AbstractSyntaxTreeFactories; @@ -88,11 +93,14 @@ import com.oceanbase.odc.service.db.session.KillSessionResult; import com.oceanbase.odc.service.dml.ValueEncodeType; import com.oceanbase.odc.service.feature.AllFeatures; +import com.oceanbase.odc.service.iam.auth.AuthenticationFacade; import com.oceanbase.odc.service.permission.database.model.DatabasePermissionType; import com.oceanbase.odc.service.permission.database.model.UnauthorizedDatabase; import com.oceanbase.odc.service.session.interceptor.SqlCheckInterceptor; import com.oceanbase.odc.service.session.interceptor.SqlConsoleInterceptor; import com.oceanbase.odc.service.session.interceptor.SqlExecuteInterceptorService; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; +import com.oceanbase.odc.service.session.model.AsyncExecuteResultResp; import com.oceanbase.odc.service.session.model.BinaryContent; import com.oceanbase.odc.service.session.model.OdcResultSetMetaData.OdcTable; import com.oceanbase.odc.service.session.model.QueryTableOrViewDataReq; @@ -124,7 +132,9 @@ public class ConnectConsoleService { public static final int DEFAULT_GET_RESULT_TIMEOUT_SECONDS = 3; - private static final String SHOW_TABLE_COLUMN_INFO = "SHOW_TABLE_COLUMN_INFO"; + public static final int DEFAULT_GET_RESULT_TIMEOUT_SECONDS_V2 = 1; + private final Executor queryProfileMonitor = + Executors.newFixedThreadPool(5, r -> new Thread(r, "query-profile-monitor-" + r.hashCode())); @Autowired private ConnectSessionService sessionService; @Autowired @@ -139,6 +149,8 @@ public class ConnectConsoleService { private ConnectionService connectionService; @Autowired private UserConfigFacade userConfigFacade; + @Autowired + private AuthenticationFacade authenticationFacade; public SqlExecuteResult queryTableOrViewData(@NotNull String sessionId, @NotNull @Valid QueryTableOrViewDataReq req) throws Exception { @@ -306,6 +318,93 @@ public SqlAsyncExecuteResp execute(@NotNull String sessionId, return response; } + public SqlAsyncExecuteResp executeV2(@NotNull String sessionId, + @NotNull @Valid SqlAsyncExecuteReq request, boolean needSqlRuleCheck) throws Exception { + ConnectionSession connectionSession = sessionService.nullSafeGet(sessionId, true); + + long maxSqlLength = sessionProperties.getMaxSqlLength(); + if (maxSqlLength > 0) { + PreConditions.lessThanOrEqualTo("sqlLength", LimitMetric.SQL_LENGTH, + StringUtils.length(request.getSql()), maxSqlLength); + } + SqlAsyncExecuteResp result = filterKillSession(connectionSession, request); + if (result != null) { + return result; + } + List sqls = request.ifSplitSqls() + ? SqlUtils.splitWithOffset(connectionSession, request.getSql(), + sessionProperties.isOracleRemoveCommentPrefix()) + : Collections.singletonList(new OffsetString(0, request.getSql())); + if (sqls.size() == 0) { + /** + * if a sql only contains delimiter setting(eg. delimiter $$), code will do this + */ + SqlTuple sqlTuple = SqlTuple.newTuple(request.getSql()); + String id = ConnectionSessionUtil.setFutureJdbc(connectionSession, + FutureResult.successResultList(JdbcGeneralResult.successResult(sqlTuple)), null); + return SqlAsyncExecuteResp.newSqlAsyncExecuteResp(id, Collections.singletonList(sqlTuple)); + } + + long maxSqlStatementCount = sessionProperties.getMaxSqlStatementCount(); + if (maxSqlStatementCount > 0) { + PreConditions.lessThanOrEqualTo("sqlStatementCount", + LimitMetric.SQL_STATEMENT_COUNT, sqls.size(), maxSqlStatementCount); + } + + List sqlTuples = generateSqlTuple(sqls, connectionSession, request); + SqlAsyncExecuteResp response = SqlAsyncExecuteResp.newSqlAsyncExecuteResp(sqlTuples); + Map context = new HashMap<>(); + context.put(SHOW_TABLE_COLUMN_INFO, request.getShowTableColumnInfo()); + context.put(SqlCheckInterceptor.NEED_SQL_CHECK_KEY, needSqlRuleCheck); + context.put(SqlConsoleInterceptor.NEED_SQL_CONSOLE_CHECK, needSqlRuleCheck); + List stages = sqlTuples.stream() + .map(s -> s.getSqlWatch().start(SqlExecuteStages.SQL_PRE_CHECK)) + .collect(Collectors.toList()); + try { + if (!sqlInterceptService.preHandle(request, response, connectionSession, context)) { + return response; + } + } finally { + for (TraceStage stage : stages) { + try { + stage.close(); + } catch (Exception e) { + // eat exception + } + } + } + Integer queryLimit = checkQueryLimit(request.getQueryLimit()); + boolean continueExecutionOnError = + Objects.nonNull(request.getContinueExecutionOnError()) ? request.getContinueExecutionOnError() + : userConfigFacade.isContinueExecutionOnError(); + boolean stopOnError = !continueExecutionOnError; + AsyncExecuteContext executeContext = + new AsyncExecuteContext(connectionSession, sqlTuples, sqlInterceptService, + authenticationFacade.currentUser()); + executeContext.getContextMap().putAll(context); + OdcStatementCallBack statementCallBack = new OdcStatementCallBack(sqlTuples, connectionSession, + request.getAutoCommit(), queryLimit, stopOnError, executeContext); + + statementCallBack.setDbmsoutputMaxRows(sessionProperties.getDbmsOutputMaxRows()); + + boolean fullLinkTraceEnabled = + Objects.nonNull(request.getFullLinkTraceEnabled()) ? request.getFullLinkTraceEnabled() + : userConfigFacade.isFullLinkTraceEnabled(); + statementCallBack.setUseFullLinkTrace(fullLinkTraceEnabled); + + statementCallBack.setFullLinkTraceTimeout(sessionProperties.getFullLinkTraceTimeoutSeconds()); + statementCallBack.setMaxCachedSize(sessionProperties.getResultSetMaxCachedSize()); + statementCallBack.setMaxCachedLines(sessionProperties.getResultSetMaxCachedLines()); + statementCallBack.setLocale(LocaleContextHolder.getLocale()); + + Future> futureResult = connectionSession.getAsyncJdbcExecutor( + ConnectionSessionConstants.CONSOLE_DS_KEY).execute(statementCallBack); + executeContext.setFuture(futureResult); + String id = ConnectionSessionUtil.setExecuteContext(connectionSession, executeContext); + response.setRequestId(id); + return response; + } + public List getAsyncResult(@NotNull String sessionId, @NotNull String requestId) { return getAsyncResult(sessionId, requestId, DEFAULT_GET_RESULT_TIMEOUT_SECONDS); } @@ -341,6 +440,28 @@ public List getAsyncResult(@NotNull String sessionId, String r } } + public AsyncExecuteResultResp getAsyncResultV2(@NotNull String sessionId, String requestId, + Integer timeoutSeconds) { + PreConditions.validArgumentState(Objects.nonNull(requestId), ErrorCodes.SqlRegulationRuleBlocked, null, null); + ConnectionSession connectionSession = sessionService.nullSafeGet(sessionId); + Long sessionUserId = ConnectionSessionUtil.getUserId(connectionSession); + if (sessionUserId != authenticationFacade.currentUserId()) { + throw new AccessDeniedException(); + } + AsyncExecuteContext executeContext = + (AsyncExecuteContext) ConnectionSessionUtil.getExecuteContext(connectionSession, requestId); + int timeout = Objects.isNull(timeoutSeconds) ? DEFAULT_GET_RESULT_TIMEOUT_SECONDS_V2 : timeoutSeconds; + if (executeContext.await(timeout, TimeUnit.SECONDS)) { + ConnectionSessionUtil.removeExecuteContext(connectionSession, requestId); + return new AsyncExecuteResultResp(SqlExecuteStatus.SUCCESS, executeContext); + } else { + if (log.isDebugEnabled()) { + log.debug("Get sql execution result timed out, sessionId={}, requestId={}", sessionId, requestId); + } + return new AsyncExecuteResultResp(SqlExecuteStatus.RUNNING, executeContext); + } + } + public BinaryContent getBinaryContent(@NotNull String sessionId, @NotNull String sqlId, @NotNull Long rowNum, @NotNull Integer colNum, @NotNull Long skip, @NotNull Integer len, @NotNull ValueEncodeType format) throws IOException { diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java index 9ef4e7c960..824201e83c 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java @@ -35,6 +35,7 @@ import java.util.Locale; import java.util.Objects; import java.util.Optional; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.function.BiPredicate; import java.util.regex.Matcher; @@ -78,6 +79,7 @@ import com.oceanbase.odc.core.sql.util.FullLinkTraceUtil; import com.oceanbase.odc.core.sql.util.OBUtils; import com.oceanbase.odc.service.plugin.ConnectionPluginUtil; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; import lombok.Getter; import lombok.NonNull; @@ -114,6 +116,7 @@ public class OdcStatementCallBack implements StatementCallback sqls, @NonNull ConnectionSes public OdcStatementCallBack(@NonNull List sqls, @NonNull ConnectionSession connectionSession, Boolean autoCommit, Integer queryLimit, boolean stopWhenError) { + this(sqls, connectionSession, autoCommit, queryLimit, stopWhenError, null); + } + + public OdcStatementCallBack(@NonNull List sqls, @NonNull ConnectionSession connectionSession, + Boolean autoCommit, Integer queryLimit, boolean stopWhenError, AsyncExecuteContext executeContext) { this.sqls = sqls; this.autoCommit = autoCommit == null ? connectionSession.getDefaultAutoCommit() : autoCommit; this.connectType = connectionSession.getConnectType(); @@ -149,6 +157,7 @@ public OdcStatementCallBack(@NonNull List sqls, @NonNull ConnectionSes this.stopWhenError = stopWhenError; this.binaryDataManager = ConnectionSessionUtil.getBinaryDataManager(connectionSession); Validate.notNull(this.binaryDataManager, "BinaryDataManager can not be null"); + this.executeContext = executeContext; } @Override @@ -158,11 +167,15 @@ public List doInStatement(Statement statement) throws SQLExce } if (ConnectionSessionUtil.isConsoleSessionReset(connectionSession)) { ConnectionSessionUtil.setConsoleSessionResetFlag(connectionSession, false); - return this.sqls.stream().map(sqlTuple -> { + List results = this.sqls.stream().map(sqlTuple -> { JdbcGeneralResult result = JdbcGeneralResult.canceledResult(sqlTuple); result.setConnectionReset(true); return result; }).collect(Collectors.toList()); + if (executeContext != null) { + executeContext.finishQuery(results); + } + return results; } List returnVal = new LinkedList<>(); boolean currentAutoCommit = statement.getConnection().getAutoCommit(); @@ -179,9 +192,13 @@ public List doInStatement(Statement statement) throws SQLExce log.warn("Init driver statistic collect failed, reason={}", e.getMessage()); } List executeResults; + CountDownLatch latch = new CountDownLatch(1); + if (executeContext != null) { + executeContext.startQuery(latch); + } if (returnVal.stream().noneMatch(r -> r.getStatus() == SqlExecuteStatus.FAILED) || !stopWhenError) { try { - executeResults = doExecuteSql(statement, sqlTuple); + executeResults = doExecuteSql(statement, sqlTuple, latch); } catch (Exception exception) { executeResults = Collections.singletonList(JdbcGeneralResult.failedResult(sqlTuple, exception)); } @@ -189,6 +206,9 @@ public List doInStatement(Statement statement) throws SQLExce executeResults = Collections.singletonList(JdbcGeneralResult.canceledResult(sqlTuple)); } returnVal.addAll(executeResults); + if (executeContext != null) { + executeContext.finishQuery(executeResults); + } } Optional failed = returnVal .stream().filter(r -> r.getStatus() == SqlExecuteStatus.FAILED).findFirst(); @@ -304,7 +324,8 @@ private List consumeStatement(Statement statement, SqlTuple s return executeResults; } - protected List doExecuteSql(Statement statement, SqlTuple sqlTuple) throws Exception { + protected List doExecuteSql(Statement statement, SqlTuple sqlTuple, + CountDownLatch latch) throws Exception { String sql = sqlTuple.getExecutedSql(); if (!ifFunctionCallExists(sql)) { // use text protocal @@ -313,8 +334,10 @@ protected List doExecuteSql(Statement statement, SqlTuple sql try { isResultSet = statement.execute(sql); } catch (Exception e) { + latch.countDown(); return handleException(e, statement, sqlTuple); } + latch.countDown(); return consumeStatement(statement, sqlTuple, isResultSet); } } @@ -345,6 +368,7 @@ protected List doExecuteSql(Statement statement, SqlTuple sql } } boolean isResult = preparedStatement.execute(); + latch.countDown(); return consumeStatement(preparedStatement, sqlTuple, isResult); } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java new file mode 100644 index 0000000000..1368991a1f --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java @@ -0,0 +1,217 @@ +/* + * Copyright (c) 2023 OceanBase. + * + * Licensed 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 com.oceanbase.odc.service.session.model; + +import static com.oceanbase.odc.core.session.ConnectionSessionConstants.BACKEND_DS_KEY; +import static com.oceanbase.odc.core.session.ConnectionSessionConstants.CONSOLE_DS_KEY; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; + +import org.springframework.jdbc.core.StatementCallback; + +import com.oceanbase.odc.common.util.StringUtils; +import com.oceanbase.odc.common.util.TraceStage; +import com.oceanbase.odc.common.util.TraceWatch; +import com.oceanbase.odc.common.util.VersionUtils; +import com.oceanbase.odc.core.session.ConnectionSession; +import com.oceanbase.odc.core.session.ConnectionSessionUtil; +import com.oceanbase.odc.core.sql.execute.SqlExecuteStages; +import com.oceanbase.odc.core.sql.execute.model.JdbcGeneralResult; +import com.oceanbase.odc.core.sql.execute.model.SqlTuple; +import com.oceanbase.odc.core.sql.util.OBUtils; +import com.oceanbase.odc.service.db.browser.DBSchemaAccessors; +import com.oceanbase.odc.service.iam.model.User; +import com.oceanbase.odc.service.iam.util.SecurityContextUtils; +import com.oceanbase.odc.service.session.interceptor.SqlExecuteInterceptorService; +import com.oceanbase.odc.service.session.model.OdcResultSetMetaData.OdcTable; +import com.oceanbase.tools.dbbrowser.schema.DBSchemaAccessor; + +import lombok.Getter; +import lombok.NonNull; +import lombok.Setter; +import lombok.extern.slf4j.Slf4j; + +/** + * @author liuyizhuo.lyz + * @date 2024/4/15 + */ +@Getter +@Slf4j +public class AsyncExecuteContext { + public static final String SHOW_TABLE_COLUMN_INFO = "SHOW_TABLE_COLUMN_INFO"; + + private final ConnectionSession session; + private final List sqlTuples; + private final List results = new ArrayList<>(); + private final Map contextMap = new HashMap<>(); + private final SqlExecuteInterceptorService sqlInterceptService; + private final ExecutorService executor = Executors.newSingleThreadExecutor(); + private final User user; + private final List sessionIds; + private final boolean queryRuntimeTrace; + + @Setter + private Future> future; + @Setter + private String currentTraceId; + private Future handle; + private int count = 0; + + public AsyncExecuteContext(ConnectionSession session, List sqlTuples, + SqlExecuteInterceptorService sqlInterceptService, User user) { + this.session = session; + this.sqlTuples = sqlTuples; + this.sqlInterceptService = sqlInterceptService; + this.user = user; + this.queryRuntimeTrace = VersionUtils.isGreaterThanOrEqualsTo(ConnectionSessionUtil.getVersion(session), "4.2") + && sqlTuples.size() <= 10; + this.sessionIds = initSessionIds(session); + } + + public boolean await(long timeout, TimeUnit unit) { + if (future == null) { + return false; + } + try { + future.get(timeout, unit); + return true; + } catch (TimeoutException e) { + return false; + } catch (Exception e) { + throw new IllegalStateException(e); + } + } + + public int getTotal() { + return sqlTuples.size(); + } + + public List getResults() { + synchronized (results) { + ArrayList copiedResults = new ArrayList<>(results); + results.clear(); + return copiedResults; + } + } + + public void startQuery(CountDownLatch latch) { + count++; + if (handle != null && !handle.isDone()) { + handle.cancel(true); + } + if (queryRuntimeTrace) { + handle = executor.submit(() -> { + try { + if (!latch.await(1100, TimeUnit.MILLISECONDS)) { + String traceId = session.getSyncJdbcExecutor(BACKEND_DS_KEY) + .execute((StatementCallback) stmt -> OBUtils + .queryTraceIdFromASH(stmt, sessionIds, session.getConnectType())); + if (traceId != null) { + setCurrentTraceId(traceId); + } + } + } catch (InterruptedException e) { + log.warn("Failed to get trace id.", e); + } + return null; + }); + } + } + + public void finishQuery(List results) { + if (handle != null && !handle.isDone()) { + handle.cancel(true); + } + synchronized (this.results) { + SecurityContextUtils.setCurrentUser(user); + try { + for (JdbcGeneralResult result : results) { + this.results.add(mapResult(result)); + } + } finally { + SecurityContextUtils.clear(); + } + } + } + + private List initSessionIds(ConnectionSession session) { + if (!queryRuntimeTrace) { + return null; + } + String proxySessId = ConnectionSessionUtil.getConsoleConnectionProxySessId(session); + if (StringUtils.isEmpty(proxySessId)) { + return Collections.singletonList(ConnectionSessionUtil.getConsoleConnectionId(session)); + } + return session.getSyncJdbcExecutor(CONSOLE_DS_KEY).execute((StatementCallback>) stmt -> OBUtils + .querySessionIdsByProxySessId(stmt, proxySessId, session.getConnectType())); + } + + private SqlExecuteResult mapResult(JdbcGeneralResult jdbcResult) { + SqlExecuteResult res = generateResult(session, jdbcResult, contextMap); + TraceWatch traceWatch = res.getSqlTuple().getSqlWatch(); + try (TraceStage stage = traceWatch.start(SqlExecuteStages.SQL_AFTER_CHECK)) { + sqlInterceptService.afterCompletion(res, session, contextMap); + } catch (Exception e) { + throw new IllegalStateException(e); + } + if (!traceWatch.isClosed()) { + traceWatch.close(); + } + return res; + } + + private SqlExecuteResult generateResult(@NonNull ConnectionSession connectionSession, + @NonNull JdbcGeneralResult generalResult, @NonNull Map cxt) { + SqlExecuteResult result = new SqlExecuteResult(generalResult); + TraceWatch watch = generalResult.getSqlTuple().getSqlWatch(); + OdcTable resultTable = null; + DBSchemaAccessor schemaAccessor = DBSchemaAccessors.create(connectionSession); + try (TraceStage s = watch.start(SqlExecuteStages.INIT_SQL_TYPE)) { + result.initSqlType(connectionSession.getDialectType()); + } catch (Exception e) { + log.warn("Failed to init sql type", e); + } + try (TraceStage s = watch.start(SqlExecuteStages.INIT_EDITABLE_INFO)) { + resultTable = result.initEditableInfo(); + } catch (Exception e) { + log.warn("Failed to init editable info", e); + } + if (Boolean.TRUE.equals(cxt.get(SHOW_TABLE_COLUMN_INFO))) { + try (TraceStage s = watch.start(SqlExecuteStages.INIT_COLUMN_INFO)) { + result.initColumnInfo(connectionSession, resultTable, schemaAccessor); + } catch (Exception e) { + log.warn("Failed to init column comment", e); + } + } + try (TraceStage s = watch.start(SqlExecuteStages.INIT_WARNING_MESSAGE)) { + result.initWarningMessage(connectionSession); + } catch (Exception e) { + log.warn("Failed to init warning message", e); + } + return result; + } + +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java new file mode 100644 index 0000000000..78524e3e30 --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java @@ -0,0 +1,43 @@ +/* + * Copyright (c) 2023 OceanBase. + * + * Licensed 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 com.oceanbase.odc.service.session.model; + +import java.util.List; + +import com.oceanbase.odc.core.sql.execute.model.SqlExecuteStatus; + +import lombok.Data; + +/** + * @author liuyizhuo.lyz + * @date 2024/4/15 + */ +@Data +public class AsyncExecuteResultResp { + List results; + private String traceId; + private int total; + private int count; + private SqlExecuteStatus status; + + public AsyncExecuteResultResp(SqlExecuteStatus status, AsyncExecuteContext context) { + this.status = status; + results = context.getResults(); + traceId = context.getCurrentTraceId(); + total = context.getTotal(); + count = context.getCount(); + } +} From 88c1060a3f4b2a39239a564eb107b0d9871802d6 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Wed, 17 Apr 2024 19:45:16 +0800 Subject: [PATCH 02/16] use queue instead of list --- .../service/session/OdcStatementCallBack.java | 1 + .../session/model/AsyncExecuteContext.java | 36 ++++++++++++------- 2 files changed, 25 insertions(+), 12 deletions(-) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java index 824201e83c..102dac81d0 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java @@ -200,6 +200,7 @@ public List doInStatement(Statement statement) throws SQLExce try { executeResults = doExecuteSql(statement, sqlTuple, latch); } catch (Exception exception) { + latch.countDown(); executeResults = Collections.singletonList(JdbcGeneralResult.failedResult(sqlTuple, exception)); } } else { diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java index 1368991a1f..9d36731266 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java @@ -23,12 +23,15 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Queue; +import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.function.BiConsumer; import org.springframework.jdbc.core.StatementCallback; @@ -65,10 +68,11 @@ public class AsyncExecuteContext { private final ConnectionSession session; private final List sqlTuples; - private final List results = new ArrayList<>(); + private final Queue results = new ConcurrentLinkedQueue<>(); private final Map contextMap = new HashMap<>(); private final SqlExecuteInterceptorService sqlInterceptService; private final ExecutorService executor = Executors.newSingleThreadExecutor(); + private final List> traceIdHooks = new ArrayList<>(); private final User user; private final List sessionIds; private final boolean queryRuntimeTrace; @@ -110,10 +114,17 @@ public int getTotal() { } public List getResults() { - synchronized (results) { - ArrayList copiedResults = new ArrayList<>(results); - results.clear(); - return copiedResults; + List copiedResults = new ArrayList<>(); + while (!results.isEmpty()) { + copiedResults.add(results.poll()); + } + return copiedResults; + } + + public void setCurrentTraceId(String traceId) { + currentTraceId = traceId; + for (BiConsumer hook : traceIdHooks) { + hook.accept(session, traceId); } } @@ -145,15 +156,16 @@ public void finishQuery(List results) { if (handle != null && !handle.isDone()) { handle.cancel(true); } - synchronized (this.results) { - SecurityContextUtils.setCurrentUser(user); - try { - for (JdbcGeneralResult result : results) { - this.results.add(mapResult(result)); + SecurityContextUtils.setCurrentUser(user); + try { + for (JdbcGeneralResult result : results) { + this.results.add(mapResult(result)); + for (BiConsumer hook : traceIdHooks) { + hook.accept(session, result.getTraceId()); } - } finally { - SecurityContextUtils.clear(); } + } finally { + SecurityContextUtils.clear(); } } From f0e8ee40c69be73548b38dab4cf15f09057ea753 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Thu, 18 Apr 2024 21:12:19 +0800 Subject: [PATCH 03/16] response to CR comments --- .../session/ConnectionSessionConstants.java | 2 +- .../core/session/ConnectionSessionUtil.java | 3 +- .../v2/ConnectSessionController.java | 2 +- .../connection/util/ConnectionInfoUtil.java | 6 +- .../session/ConnectConsoleService.java | 28 +-- .../service/session/OdcStatementCallBack.java | 160 ++++++++++-------- .../session/model/AsyncExecuteContext.java | 138 +++------------ .../session/model/AsyncExecuteResultResp.java | 2 +- 8 files changed, 139 insertions(+), 202 deletions(-) diff --git a/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionConstants.java b/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionConstants.java index 81884b6052..417c45e072 100644 --- a/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionConstants.java +++ b/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionConstants.java @@ -86,7 +86,7 @@ public class ConnectionSessionConstants { * of attributes, this is the key */ public static final String CONNECTION_ID_KEY = "CONNECTION_ID"; - public static final String CONNECTION_PROXY_SESSID_KEY = "PROXY_SESSID"; + public static final String OB_PROXY_SESSID_KEY = "PROXY_SESSID"; /** * The {@link BinaryDataManager} session needs to be stored in the database session in the form of * attributes, this is the key diff --git a/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionUtil.java b/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionUtil.java index 34c88167d1..6d0a2a4df6 100644 --- a/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionUtil.java +++ b/server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionUtil.java @@ -569,7 +569,8 @@ public static String getConsoleConnectionId(@NonNull ConnectionSession connectio } public static String getConsoleConnectionProxySessId(@NonNull ConnectionSession connectionSession) { - return (String) connectionSession.getAttribute(ConnectionSessionConstants.CONNECTION_PROXY_SESSID_KEY); + Object proxySessId = connectionSession.getAttribute(ConnectionSessionConstants.OB_PROXY_SESSID_KEY); + return proxySessId == null ? null : (String) proxySessId; } public static String getVersion(@NonNull ConnectionSession connectionSession) { diff --git a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/ConnectSessionController.java b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/ConnectSessionController.java index b886ee6e30..2343b5f941 100644 --- a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/ConnectSessionController.java +++ b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/ConnectSessionController.java @@ -144,7 +144,7 @@ public SuccessResponse> getAsyncSqlExecute(@PathVariable @StatefulRoute(stateName = StateName.DB_SESSION, stateIdExpression = "#sessionId") public SuccessResponse getAsyncSqlExecuteV2(@PathVariable String sessionId, @RequestParam String requestId) { - return Responses.success(consoleService.getAsyncResultV2(SidUtils.getSessionId(sessionId), requestId, null)); + return Responses.success(consoleService.getAsyncResultV2(SidUtils.getSessionId(sessionId), requestId)); } /** diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java index 9b31b99adb..a2b638d359 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java @@ -66,8 +66,8 @@ public static void initConnectionId(@NonNull Statement statement, @NonNull Conne connectionSession.setAttribute(ConnectionSessionConstants.CONNECTION_ID_KEY, sessionId); if (connectionSession.getDialectType().isOceanbase() && VersionUtils.isGreaterThanOrEqualsTo( ConnectionSessionUtil.getVersion(connectionSession), "4.2")) { - String proxySessId = queryProxySessId(statement, connectionSession.getDialectType(), sessionId); - connectionSession.setAttribute(ConnectionSessionConstants.CONNECTION_PROXY_SESSID_KEY, proxySessId); + String proxySessId = queryOBProxySessId(statement, connectionSession.getDialectType(), sessionId); + connectionSession.setAttribute(ConnectionSessionConstants.OB_PROXY_SESSID_KEY, proxySessId); } } catch (Exception exception) { log.warn("Failed to get database session ID, session={}", connectionSession, exception); @@ -79,7 +79,7 @@ public static String queryConnectionId(@NonNull Statement statement, @NonNull Di return ConnectionPluginUtil.getSessionExtension(dialectType).getConnectionId(statement.getConnection()); } - public static String queryProxySessId(@NonNull Statement statement, @NonNull DialectType dialectType, + public static String queryOBProxySessId(@NonNull Statement statement, @NonNull DialectType dialectType, @NonNull String connectionId) throws SQLException { String proxySessId = null; String sql = "select proxy_sessid from " diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java index 41045d80a3..d41f84a6ca 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java @@ -132,7 +132,6 @@ public class ConnectConsoleService { public static final int DEFAULT_GET_RESULT_TIMEOUT_SECONDS = 3; - public static final int DEFAULT_GET_RESULT_TIMEOUT_SECONDS_V2 = 1; private final Executor queryProfileMonitor = Executors.newFixedThreadPool(5, r -> new Thread(r, "query-profile-monitor-" + r.hashCode())); @Autowired @@ -378,10 +377,17 @@ public SqlAsyncExecuteResp executeV2(@NotNull String sessionId, Objects.nonNull(request.getContinueExecutionOnError()) ? request.getContinueExecutionOnError() : userConfigFacade.isContinueExecutionOnError(); boolean stopOnError = !continueExecutionOnError; - AsyncExecuteContext executeContext = - new AsyncExecuteContext(connectionSession, sqlTuples, sqlInterceptService, - authenticationFacade.currentUser()); - executeContext.getContextMap().putAll(context); + AsyncExecuteContext executeContext = + new AsyncExecuteContext<>(connectionSession, sqlTuples, jdbcResult -> { + SqlExecuteResult res = generateResult(connectionSession, jdbcResult, context); + try (TraceStage stage = + res.getSqlTuple().getSqlWatch().start(SqlExecuteStages.SQL_AFTER_CHECK)) { + sqlInterceptService.afterCompletion(res, connectionSession, context); + } catch (Exception e) { + throw new IllegalStateException(e); + } + return res; + }, authenticationFacade.currentUser(), 1100L); OdcStatementCallBack statementCallBack = new OdcStatementCallBack(sqlTuples, connectionSession, request.getAutoCommit(), queryLimit, stopOnError, executeContext); @@ -440,18 +446,12 @@ public List getAsyncResult(@NotNull String sessionId, String r } } - public AsyncExecuteResultResp getAsyncResultV2(@NotNull String sessionId, String requestId, - Integer timeoutSeconds) { + public AsyncExecuteResultResp getAsyncResultV2(@NotNull String sessionId, String requestId) { PreConditions.validArgumentState(Objects.nonNull(requestId), ErrorCodes.SqlRegulationRuleBlocked, null, null); ConnectionSession connectionSession = sessionService.nullSafeGet(sessionId); - Long sessionUserId = ConnectionSessionUtil.getUserId(connectionSession); - if (sessionUserId != authenticationFacade.currentUserId()) { - throw new AccessDeniedException(); - } - AsyncExecuteContext executeContext = + AsyncExecuteContext executeContext = (AsyncExecuteContext) ConnectionSessionUtil.getExecuteContext(connectionSession, requestId); - int timeout = Objects.isNull(timeoutSeconds) ? DEFAULT_GET_RESULT_TIMEOUT_SECONDS_V2 : timeoutSeconds; - if (executeContext.await(timeout, TimeUnit.SECONDS)) { + if (executeContext.getFuture().isDone()) { ConnectionSessionUtil.removeExecuteContext(connectionSession, requestId); return new AsyncExecuteResultResp(SqlExecuteStatus.SUCCESS, executeContext); } else { diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java index 102dac81d0..f372b3cac6 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java @@ -36,6 +36,9 @@ import java.util.Objects; import java.util.Optional; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.function.BiPredicate; import java.util.regex.Matcher; @@ -80,6 +83,7 @@ import com.oceanbase.odc.core.sql.util.OBUtils; import com.oceanbase.odc.service.plugin.ConnectionPluginUtil; import com.oceanbase.odc.service.session.model.AsyncExecuteContext; +import com.oceanbase.odc.service.session.model.SqlExecuteResult; import lombok.Getter; import lombok.NonNull; @@ -116,7 +120,8 @@ public class OdcStatementCallBack implements StatementCallback context; + private final ExecutorService executor = Executors.newSingleThreadExecutor(); @Setter private boolean useFullLinkTrace = false; @Setter @@ -137,13 +142,15 @@ public OdcStatementCallBack(@NonNull List sqls, @NonNull ConnectionSes this(sqls, connectionSession, autoCommit, queryLimit, true); } + public OdcStatementCallBack(@NonNull List sqls, @NonNull ConnectionSession connectionSession, Boolean autoCommit, Integer queryLimit, boolean stopWhenError) { - this(sqls, connectionSession, autoCommit, queryLimit, stopWhenError, null); + this(sqls, connectionSession, autoCommit, queryLimit, true, null); } public OdcStatementCallBack(@NonNull List sqls, @NonNull ConnectionSession connectionSession, - Boolean autoCommit, Integer queryLimit, boolean stopWhenError, AsyncExecuteContext executeContext) { + Boolean autoCommit, Integer queryLimit, boolean stopWhenError, + AsyncExecuteContext context) { this.sqls = sqls; this.autoCommit = autoCommit == null ? connectionSession.getDefaultAutoCommit() : autoCommit; this.connectType = connectionSession.getConnectType(); @@ -157,7 +164,7 @@ public OdcStatementCallBack(@NonNull List sqls, @NonNull ConnectionSes this.stopWhenError = stopWhenError; this.binaryDataManager = ConnectionSessionUtil.getBinaryDataManager(connectionSession); Validate.notNull(this.binaryDataManager, "BinaryDataManager can not be null"); - this.executeContext = executeContext; + this.context = context; } @Override @@ -172,12 +179,11 @@ public List doInStatement(Statement statement) throws SQLExce result.setConnectionReset(true); return result; }).collect(Collectors.toList()); - if (executeContext != null) { - executeContext.finishQuery(results); + if (context != null) { + context.onQueryFinish(results); } return results; } - List returnVal = new LinkedList<>(); boolean currentAutoCommit = statement.getConnection().getAutoCommit(); try { applyStatementSettings(statement); @@ -185,44 +191,58 @@ public List doInStatement(Statement statement) throws SQLExce if (this.autoCommit ^ currentAutoCommit) { statement.getConnection().setAutoCommit(this.autoCommit); } + List returnVal = new LinkedList<>(); + Future handle = null; for (SqlTuple sqlTuple : this.sqls) { + if (handle != null) { + try { + handle.get(); + } catch (Exception e) { + // eat exception + } + } + if (context != null) { + context.onQueryStart(); + } try { applyConnectionSettings(statement); } catch (Exception e) { log.warn("Init driver statistic collect failed, reason={}", e.getMessage()); } List executeResults; - CountDownLatch latch = new CountDownLatch(1); - if (executeContext != null) { - executeContext.startQuery(latch); - } if (returnVal.stream().noneMatch(r -> r.getStatus() == SqlExecuteStatus.FAILED) || !stopWhenError) { - try { - executeResults = doExecuteSql(statement, sqlTuple, latch); - } catch (Exception exception) { - latch.countDown(); - executeResults = Collections.singletonList(JdbcGeneralResult.failedResult(sqlTuple, exception)); + CountDownLatch latch = new CountDownLatch(1); + if (context != null) { + handle = executor.submit(() -> { + try { + if (!latch.await(1100, TimeUnit.MILLISECONDS)) { + context.onQueryExecuting(); + } + } catch (InterruptedException e) { + log.warn("Failed to call back.", e); + } + return null; + }); } + executeResults = doExecuteSql(statement, sqlTuple, latch); } else { executeResults = Collections.singletonList(JdbcGeneralResult.canceledResult(sqlTuple)); } returnVal.addAll(executeResults); - if (executeContext != null) { - executeContext.finishQuery(executeResults); + if (context != null) { + context.onQueryFinish(executeResults); } } Optional failed = returnVal .stream().filter(r -> r.getStatus() == SqlExecuteStatus.FAILED).findFirst(); if (failed.isPresent()) { - throw failed.get().getThrown(); - } - - } catch (Exception e) { - try { - ConnectionSessionUtil.logSocketInfo(statement.getConnection(), "console error"); - } catch (Exception exception) { - log.warn("Failed to execute abnormal replenishment logic", exception); + try { + ConnectionSessionUtil.logSocketInfo(statement.getConnection(), "console error"); + } catch (Exception exception) { + log.warn("Failed to execute abnormal replenishment logic", exception); + } } + return returnVal; } finally { if (this.autoCommit ^ currentAutoCommit) { statement.getConnection().setAutoCommit(currentAutoCommit); @@ -234,7 +254,6 @@ public List doInStatement(Statement statement) throws SQLExce } } } - return returnVal; } private void applyConnectionSettings(Statement statement) throws SQLException { @@ -325,56 +344,65 @@ private List consumeStatement(Statement statement, SqlTuple s return executeResults; } - protected List doExecuteSql(Statement statement, SqlTuple sqlTuple, - CountDownLatch latch) throws Exception { - String sql = sqlTuple.getExecutedSql(); - if (!ifFunctionCallExists(sql)) { - // use text protocal - try (TraceStage stage = sqlTuple.getSqlWatch().start(SqlExecuteStages.EXECUTE)) { + protected List doExecuteSql(Statement statement, SqlTuple sqlTuple, CountDownLatch latch) { + try { + String sql = sqlTuple.getExecutedSql(); + if (!ifFunctionCallExists(sql)) { + // use text protocal + try (TraceStage stage = sqlTuple.getSqlWatch().start(SqlExecuteStages.EXECUTE)) { + boolean isResultSet; + try { + isResultSet = statement.execute(sql); + } catch (SQLException e) { + return handleException(e, statement, sqlTuple); + } + latch.countDown(); + return consumeStatement(statement, sqlTuple, isResultSet); + } + } + // use ps protocal + String preparedSql = OBJECT_VALUE_PATTERN.matcher(sql).replaceAll("?"); + log.info( + "Load_file call is detected in sql, use ps protocol to rewrite " + + "the original sql, originalSql={}, modifiedSql={}", + sql, preparedSql); + List definitions = retrieveFunctionCalls(sql); + log.info("There is a function call in sql, functions={}", definitions); + try (PreparedStatement preparedStatement = statement.getConnection().prepareStatement(preparedSql); + TraceStage stage = sqlTuple.getSqlWatch().start(SqlExecuteStages.EXECUTE)) { + for (int i = 0; i < definitions.size(); i++) { + FunctionDefinition definition = definitions.get(i); + if ("load_file".equalsIgnoreCase(definition.getFunctionName())) { + // load binary data + String fileName = retrieveFileNameFromParameters(definition); + File file = nullSafeFindFileByName(fileName); + preparedStatement.setBinaryStream(i + 1, new FileInputStream(file)); + } else if ("load_clob_file".equalsIgnoreCase(definition.getFunctionName())) { + // load clob data + String fileName = retrieveFileNameFromParameters(definition); + File file = nullSafeFindFileByName(fileName); + preparedStatement.setClob(i + 1, new FileReader(file)); + } else { + throw new NotImplementedException("Unsupport function call " + definition.getFunctionName()); + } + } boolean isResultSet; try { - isResultSet = statement.execute(sql); - } catch (Exception e) { - latch.countDown(); + isResultSet = preparedStatement.execute(); + } catch (SQLException e) { return handleException(e, statement, sqlTuple); } latch.countDown(); return consumeStatement(statement, sqlTuple, isResultSet); } - } - // use ps protocal - String preparedSql = OBJECT_VALUE_PATTERN.matcher(sql).replaceAll("?"); - log.info( - "Load_file call is detected in sql, use ps protocol to rewrite " - + "the original sql, originalSql={}, modifiedSql={}", - sql, preparedSql); - List definitions = retrieveFunctionCalls(sql); - log.info("There is a function call in sql, functions={}", definitions); - try (PreparedStatement preparedStatement = statement.getConnection().prepareStatement(preparedSql); - TraceStage stage = sqlTuple.getSqlWatch().start(SqlExecuteStages.EXECUTE)) { - for (int i = 0; i < definitions.size(); i++) { - FunctionDefinition definition = definitions.get(i); - if ("load_file".equalsIgnoreCase(definition.getFunctionName())) { - // load binary data - String fileName = retrieveFileNameFromParameters(definition); - File file = nullSafeFindFileByName(fileName); - preparedStatement.setBinaryStream(i + 1, new FileInputStream(file)); - } else if ("load_clob_file".equalsIgnoreCase(definition.getFunctionName())) { - // load clob data - String fileName = retrieveFileNameFromParameters(definition); - File file = nullSafeFindFileByName(fileName); - preparedStatement.setClob(i + 1, new FileReader(file)); - } else { - throw new NotImplementedException("Unsupport function call " + definition.getFunctionName()); - } - } - boolean isResult = preparedStatement.execute(); + } catch (Exception e) { + return Collections.singletonList(JdbcGeneralResult.failedResult(sqlTuple, e)); + } finally { latch.countDown(); - return consumeStatement(preparedStatement, sqlTuple, isResult); } } - private List handleException(Exception exception, Statement statement, SqlTuple sqlTuple) { + protected List handleException(Exception exception, Statement statement, SqlTuple sqlTuple) { if (exception instanceof SQLTransientConnectionException && ((SQLTransientConnectionException) exception).getErrorCode() == 1094) { // ERROR 1094 (HY000) : Unknown thread id: %lu when kill a not exists session diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java index 9d36731266..e9357ee503 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java @@ -20,40 +20,26 @@ import java.util.ArrayList; import java.util.Collections; -import java.util.HashMap; import java.util.List; -import java.util.Map; import java.util.Queue; import java.util.concurrent.ConcurrentLinkedQueue; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import java.util.concurrent.Future; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; import java.util.function.BiConsumer; +import java.util.function.Function; import org.springframework.jdbc.core.StatementCallback; import com.oceanbase.odc.common.util.StringUtils; -import com.oceanbase.odc.common.util.TraceStage; -import com.oceanbase.odc.common.util.TraceWatch; import com.oceanbase.odc.common.util.VersionUtils; import com.oceanbase.odc.core.session.ConnectionSession; import com.oceanbase.odc.core.session.ConnectionSessionUtil; -import com.oceanbase.odc.core.sql.execute.SqlExecuteStages; import com.oceanbase.odc.core.sql.execute.model.JdbcGeneralResult; import com.oceanbase.odc.core.sql.execute.model.SqlTuple; import com.oceanbase.odc.core.sql.util.OBUtils; -import com.oceanbase.odc.service.db.browser.DBSchemaAccessors; import com.oceanbase.odc.service.iam.model.User; import com.oceanbase.odc.service.iam.util.SecurityContextUtils; -import com.oceanbase.odc.service.session.interceptor.SqlExecuteInterceptorService; -import com.oceanbase.odc.service.session.model.OdcResultSetMetaData.OdcTable; -import com.oceanbase.tools.dbbrowser.schema.DBSchemaAccessor; import lombok.Getter; -import lombok.NonNull; import lombok.Setter; import lombok.extern.slf4j.Slf4j; @@ -63,58 +49,39 @@ */ @Getter @Slf4j -public class AsyncExecuteContext { +public class AsyncExecuteContext { public static final String SHOW_TABLE_COLUMN_INFO = "SHOW_TABLE_COLUMN_INFO"; + private final Function mapper; private final ConnectionSession session; private final List sqlTuples; - private final Queue results = new ConcurrentLinkedQueue<>(); - private final Map contextMap = new HashMap<>(); - private final SqlExecuteInterceptorService sqlInterceptService; - private final ExecutorService executor = Executors.newSingleThreadExecutor(); + private final Queue results = new ConcurrentLinkedQueue<>(); private final List> traceIdHooks = new ArrayList<>(); private final User user; + private final Long waitTimeMillis; private final List sessionIds; - private final boolean queryRuntimeTrace; @Setter private Future> future; - @Setter private String currentTraceId; - private Future handle; private int count = 0; public AsyncExecuteContext(ConnectionSession session, List sqlTuples, - SqlExecuteInterceptorService sqlInterceptService, User user) { + Function mapper, User user, Long waitTimeMillis) { this.session = session; this.sqlTuples = sqlTuples; - this.sqlInterceptService = sqlInterceptService; + this.mapper = mapper; this.user = user; - this.queryRuntimeTrace = VersionUtils.isGreaterThanOrEqualsTo(ConnectionSessionUtil.getVersion(session), "4.2") - && sqlTuples.size() <= 10; + this.waitTimeMillis = waitTimeMillis; this.sessionIds = initSessionIds(session); } - public boolean await(long timeout, TimeUnit unit) { - if (future == null) { - return false; - } - try { - future.get(timeout, unit); - return true; - } catch (TimeoutException e) { - return false; - } catch (Exception e) { - throw new IllegalStateException(e); - } - } - public int getTotal() { return sqlTuples.size(); } - public List getResults() { - List copiedResults = new ArrayList<>(); + public List getResults() { + List copiedResults = new ArrayList<>(); while (!results.isEmpty()) { copiedResults.add(results.poll()); } @@ -128,38 +95,24 @@ public void setCurrentTraceId(String traceId) { } } - public void startQuery(CountDownLatch latch) { + public void onQueryStart() { count++; - if (handle != null && !handle.isDone()) { - handle.cancel(true); - } - if (queryRuntimeTrace) { - handle = executor.submit(() -> { - try { - if (!latch.await(1100, TimeUnit.MILLISECONDS)) { - String traceId = session.getSyncJdbcExecutor(BACKEND_DS_KEY) - .execute((StatementCallback) stmt -> OBUtils - .queryTraceIdFromASH(stmt, sessionIds, session.getConnectType())); - if (traceId != null) { - setCurrentTraceId(traceId); - } - } - } catch (InterruptedException e) { - log.warn("Failed to get trace id.", e); - } - return null; - }); - } } - public void finishQuery(List results) { - if (handle != null && !handle.isDone()) { - handle.cancel(true); + public void onQueryExecuting() { + String traceId = session.getSyncJdbcExecutor(BACKEND_DS_KEY) + .execute((StatementCallback) stmt -> OBUtils + .queryTraceIdFromASH(stmt, sessionIds, session.getConnectType())); + if (traceId != null) { + setCurrentTraceId(traceId); } + } + + public void onQueryFinish(List results) { SecurityContextUtils.setCurrentUser(user); try { for (JdbcGeneralResult result : results) { - this.results.add(mapResult(result)); + this.results.add(mapper.apply(result)); for (BiConsumer hook : traceIdHooks) { hook.accept(session, result.getTraceId()); } @@ -170,7 +123,8 @@ public void finishQuery(List results) { } private List initSessionIds(ConnectionSession session) { - if (!queryRuntimeTrace) { + if (VersionUtils.isLessThan(ConnectionSessionUtil.getVersion(session), "4.2") + || sqlTuples.size() > 10) { return null; } String proxySessId = ConnectionSessionUtil.getConsoleConnectionProxySessId(session); @@ -180,50 +134,4 @@ private List initSessionIds(ConnectionSession session) { return session.getSyncJdbcExecutor(CONSOLE_DS_KEY).execute((StatementCallback>) stmt -> OBUtils .querySessionIdsByProxySessId(stmt, proxySessId, session.getConnectType())); } - - private SqlExecuteResult mapResult(JdbcGeneralResult jdbcResult) { - SqlExecuteResult res = generateResult(session, jdbcResult, contextMap); - TraceWatch traceWatch = res.getSqlTuple().getSqlWatch(); - try (TraceStage stage = traceWatch.start(SqlExecuteStages.SQL_AFTER_CHECK)) { - sqlInterceptService.afterCompletion(res, session, contextMap); - } catch (Exception e) { - throw new IllegalStateException(e); - } - if (!traceWatch.isClosed()) { - traceWatch.close(); - } - return res; - } - - private SqlExecuteResult generateResult(@NonNull ConnectionSession connectionSession, - @NonNull JdbcGeneralResult generalResult, @NonNull Map cxt) { - SqlExecuteResult result = new SqlExecuteResult(generalResult); - TraceWatch watch = generalResult.getSqlTuple().getSqlWatch(); - OdcTable resultTable = null; - DBSchemaAccessor schemaAccessor = DBSchemaAccessors.create(connectionSession); - try (TraceStage s = watch.start(SqlExecuteStages.INIT_SQL_TYPE)) { - result.initSqlType(connectionSession.getDialectType()); - } catch (Exception e) { - log.warn("Failed to init sql type", e); - } - try (TraceStage s = watch.start(SqlExecuteStages.INIT_EDITABLE_INFO)) { - resultTable = result.initEditableInfo(); - } catch (Exception e) { - log.warn("Failed to init editable info", e); - } - if (Boolean.TRUE.equals(cxt.get(SHOW_TABLE_COLUMN_INFO))) { - try (TraceStage s = watch.start(SqlExecuteStages.INIT_COLUMN_INFO)) { - result.initColumnInfo(connectionSession, resultTable, schemaAccessor); - } catch (Exception e) { - log.warn("Failed to init column comment", e); - } - } - try (TraceStage s = watch.start(SqlExecuteStages.INIT_WARNING_MESSAGE)) { - result.initWarningMessage(connectionSession); - } catch (Exception e) { - log.warn("Failed to init warning message", e); - } - return result; - } - } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java index 78524e3e30..d16d4e9e73 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java @@ -33,7 +33,7 @@ public class AsyncExecuteResultResp { private int count; private SqlExecuteStatus status; - public AsyncExecuteResultResp(SqlExecuteStatus status, AsyncExecuteContext context) { + public AsyncExecuteResultResp(SqlExecuteStatus status, AsyncExecuteContext context) { this.status = status; results = context.getResults(); traceId = context.getCurrentTraceId(); From 9a97a6f4f858593b48855bf03d39109d95df9547 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Mon, 22 Apr 2024 15:12:14 +0800 Subject: [PATCH 04/16] add sql text into async resp --- .../oceanbase/odc/service/session/OdcStatementCallBack.java | 2 +- .../odc/service/session/model/AsyncExecuteContext.java | 5 ++++- .../odc/service/session/model/AsyncExecuteResultResp.java | 2 ++ 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java index f372b3cac6..4687f6a24e 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java @@ -202,7 +202,7 @@ public List doInStatement(Statement statement) throws SQLExce } } if (context != null) { - context.onQueryStart(); + context.onQueryStart(sqlTuple.getExecutedSql()); } try { applyConnectionSettings(statement); diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java index e9357ee503..5285c2d845 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java @@ -64,6 +64,8 @@ public class AsyncExecuteContext { @Setter private Future> future; private String currentTraceId; + @Setter + private String sql; private int count = 0; public AsyncExecuteContext(ConnectionSession session, List sqlTuples, @@ -95,8 +97,9 @@ public void setCurrentTraceId(String traceId) { } } - public void onQueryStart() { + public void onQueryStart(String sql) { count++; + this.sql = sql; } public void onQueryExecuting() { diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java index d16d4e9e73..58f046885f 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java @@ -32,6 +32,7 @@ public class AsyncExecuteResultResp { private int total; private int count; private SqlExecuteStatus status; + private String sql; public AsyncExecuteResultResp(SqlExecuteStatus status, AsyncExecuteContext context) { this.status = status; @@ -39,5 +40,6 @@ public AsyncExecuteResultResp(SqlExecuteStatus status, AsyncExecuteContext Date: Tue, 23 Apr 2024 11:55:10 +0800 Subject: [PATCH 05/16] response to CR comments --- .../oceanbase/odc/core/sql/util/OBUtils.java | 17 +++++++++++++++- .../connection/util/ConnectionInfoUtil.java | 18 ++--------------- .../session/ConnectConsoleService.java | 18 +++++++---------- .../service/session/OdcStatementCallBack.java | 20 ++++++++++--------- .../session/model/AsyncExecuteContext.java | 6 ++++-- .../session/model/AsyncExecuteResultResp.java | 8 +++----- 6 files changed, 43 insertions(+), 44 deletions(-) diff --git a/server/odc-core/src/main/java/com/oceanbase/odc/core/sql/util/OBUtils.java b/server/odc-core/src/main/java/com/oceanbase/odc/core/sql/util/OBUtils.java index 31ef191eab..286024fb23 100644 --- a/server/odc-core/src/main/java/com/oceanbase/odc/core/sql/util/OBUtils.java +++ b/server/odc-core/src/main/java/com/oceanbase/odc/core/sql/util/OBUtils.java @@ -382,6 +382,21 @@ public static String queryDBMSOutput(@NonNull Connection connection, Integer max } } + public static String queryOBProxySessId(@NonNull Statement statement, @NonNull DialectType dialectType, + @NonNull String connectionId) throws SQLException { + String proxySessId = null; + String sql = "select proxy_sessid from " + + (dialectType.isMysql() ? "oceanbase" : "sys") + + ".v$ob_processlist where id = " + + connectionId; + try (ResultSet rs = statement.executeQuery(sql)) { + if (rs.next()) { + proxySessId = rs.getString(1); + } + } + return proxySessId; + } + public static List querySessionIdsByProxySessId(@NonNull Statement statement, @NonNull String proxySessId, ConnectType connectType) throws SQLException { DialectType dialectType = connectType.getDialectType(); @@ -431,7 +446,7 @@ public static String queryPlanIdByTraceId(@NonNull Statement statement, String t SqlBuilder sqlBuilder = getBuilder(connectType) .append("select plan_id from ") .append(dialectType.isMysql() ? "oceanbase" : "sys") - .append(".active_session_history where trace_id=") + .append(".v$active_session_history where trace_id=") .value(traceId); try (ResultSet rs = statement.executeQuery(sqlBuilder.toString())) { if (!rs.next()) { diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java index a2b638d359..2181e81b08 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java @@ -33,6 +33,7 @@ import com.oceanbase.odc.core.shared.constant.DialectType; import com.oceanbase.odc.core.sql.execute.GeneralSyncJdbcExecutor; import com.oceanbase.odc.core.sql.execute.SyncJdbcExecutor; +import com.oceanbase.odc.core.sql.util.OBUtils; import com.oceanbase.odc.plugin.connect.api.InformationExtensionPoint; import com.oceanbase.odc.service.plugin.ConnectionPluginUtil; @@ -66,7 +67,7 @@ public static void initConnectionId(@NonNull Statement statement, @NonNull Conne connectionSession.setAttribute(ConnectionSessionConstants.CONNECTION_ID_KEY, sessionId); if (connectionSession.getDialectType().isOceanbase() && VersionUtils.isGreaterThanOrEqualsTo( ConnectionSessionUtil.getVersion(connectionSession), "4.2")) { - String proxySessId = queryOBProxySessId(statement, connectionSession.getDialectType(), sessionId); + String proxySessId = OBUtils.queryOBProxySessId(statement, connectionSession.getDialectType(), sessionId); connectionSession.setAttribute(ConnectionSessionConstants.OB_PROXY_SESSID_KEY, proxySessId); } } catch (Exception exception) { @@ -79,21 +80,6 @@ public static String queryConnectionId(@NonNull Statement statement, @NonNull Di return ConnectionPluginUtil.getSessionExtension(dialectType).getConnectionId(statement.getConnection()); } - public static String queryOBProxySessId(@NonNull Statement statement, @NonNull DialectType dialectType, - @NonNull String connectionId) throws SQLException { - String proxySessId = null; - String sql = "select proxy_sessid from " - + (dialectType.isMysql() ? "oceanbase" : "sys") - + ".v$ob_processlist where id = " - + connectionId; - try (ResultSet rs = statement.executeQuery(sql)) { - if (rs.next()) { - proxySessId = rs.getString(1); - } - } - return proxySessId; - } - public static void initSessionVersion(@NonNull ConnectionSession connectionSession) { InformationExtensionPoint point = ConnectionPluginUtil.getInformationExtension(connectionSession.getDialectType()); diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java index d1d7a790e6..d6b1a3bff9 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java @@ -15,8 +15,6 @@ */ package com.oceanbase.odc.service.session; -import static com.oceanbase.odc.service.session.model.AsyncExecuteContext.SHOW_TABLE_COLUMN_INFO; - import java.io.IOException; import java.io.InputStream; import java.util.Arrays; @@ -26,8 +24,6 @@ import java.util.Map; import java.util.Objects; import java.util.concurrent.ExecutionException; -import java.util.concurrent.Executor; -import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; @@ -76,7 +72,6 @@ import com.oceanbase.odc.core.sql.execute.cache.table.VirtualElement; import com.oceanbase.odc.core.sql.execute.cache.table.VirtualTable; import com.oceanbase.odc.core.sql.execute.model.JdbcGeneralResult; -import com.oceanbase.odc.core.sql.execute.model.SqlExecuteStatus; import com.oceanbase.odc.core.sql.execute.model.SqlTuple; import com.oceanbase.odc.core.sql.parser.AbstractSyntaxTree; import com.oceanbase.odc.core.sql.parser.AbstractSyntaxTreeFactories; @@ -132,8 +127,9 @@ public class ConnectConsoleService { public static final int DEFAULT_GET_RESULT_TIMEOUT_SECONDS = 3; - private final Executor queryProfileMonitor = - Executors.newFixedThreadPool(5, r -> new Thread(r, "query-profile-monitor-" + r.hashCode())); + public static final String SHOW_TABLE_COLUMN_INFO = "SHOW_TABLE_COLUMN_INFO"; + private static final Long EXECUTING_CONTEXT_WAIT_MILLIS = 1100L; + @Autowired private ConnectSessionService sessionService; @Autowired @@ -387,7 +383,7 @@ public SqlAsyncExecuteResp executeV2(@NotNull String sessionId, throw new IllegalStateException(e); } return res; - }, authenticationFacade.currentUser(), 1100L); + }, authenticationFacade.currentUser(), EXECUTING_CONTEXT_WAIT_MILLIS); OdcStatementCallBack statementCallBack = new OdcStatementCallBack(sqlTuples, connectionSession, request.getAutoCommit(), queryLimit, stopOnError, executeContext); @@ -451,14 +447,14 @@ public AsyncExecuteResultResp getAsyncResultV2(@NotNull String sessionId, String ConnectionSession connectionSession = sessionService.nullSafeGet(sessionId); AsyncExecuteContext executeContext = (AsyncExecuteContext) ConnectionSessionUtil.getExecuteContext(connectionSession, requestId); - if (executeContext.getFuture().isDone()) { + if (executeContext.isFinished()) { ConnectionSessionUtil.removeExecuteContext(connectionSession, requestId); - return new AsyncExecuteResultResp(SqlExecuteStatus.SUCCESS, executeContext); + return new AsyncExecuteResultResp(true, executeContext); } else { if (log.isDebugEnabled()) { log.debug("Get sql execution result timed out, sessionId={}, requestId={}", sessionId, requestId); } - return new AsyncExecuteResultResp(SqlExecuteStatus.RUNNING, executeContext); + return new AsyncExecuteResultResp(false, executeContext); } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java index 3c6ba33b30..73ec59d5c9 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java @@ -186,13 +186,13 @@ public List doInStatement(Statement statement) throws SQLExce return results; } boolean currentAutoCommit = statement.getConnection().getAutoCommit(); + List returnVal = new LinkedList<>(); try { applyStatementSettings(statement); // 对于修改表数据DML，如果是自动提交，为了保证原子性，在执行过程设置为手动，执行完成后再进行reset if (this.autoCommit ^ currentAutoCommit) { statement.getConnection().setAutoCommit(this.autoCommit); } - List returnVal = new LinkedList<>(); Future handle = null; for (SqlTuple sqlTuple : this.sqls) { if (handle != null) { @@ -220,7 +220,7 @@ public List doInStatement(Statement statement) throws SQLExce if (context != null) { handle = executor.submit(() -> { try { - if (!latch.await(1100, TimeUnit.MILLISECONDS)) { + if (!latch.await(context.getWaitTimeMillis(), TimeUnit.MILLISECONDS)) { context.onQueryExecuting(); } } catch (InterruptedException e) { @@ -229,7 +229,7 @@ public List doInStatement(Statement statement) throws SQLExce return null; }); } - executeResults = doExecuteSql(statement, sqlTuple); + executeResults = doExecuteSql(statement, sqlTuple, latch); } } else { executeResults = Collections.singletonList(JdbcGeneralResult.canceledResult(sqlTuple)); @@ -242,13 +242,14 @@ public List doInStatement(Statement statement) throws SQLExce Optional failed = returnVal .stream().filter(r -> r.getStatus() == SqlExecuteStatus.FAILED).findFirst(); if (failed.isPresent()) { - try { - ConnectionSessionUtil.logSocketInfo(statement.getConnection(), "console error"); - } catch (Exception exception) { - log.warn("Failed to execute abnormal replenishment logic", exception); - } + throw failed.get().getThrown(); + } + } catch (Exception e){ + try { + ConnectionSessionUtil.logSocketInfo(statement.getConnection(), "console error"); + } catch (Exception exception) { + log.warn("Failed to execute abnormal replenishment logic", exception); } - return returnVal; } finally { if (this.autoCommit ^ currentAutoCommit) { statement.getConnection().setAutoCommit(currentAutoCommit); @@ -260,6 +261,7 @@ public List doInStatement(Statement statement) throws SQLExce } } } + return returnVal; } private void applyConnectionSettings(Statement statement) throws SQLException { diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java index 5285c2d845..33f8e2ce05 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java @@ -50,8 +50,6 @@ @Getter @Slf4j public class AsyncExecuteContext { - public static final String SHOW_TABLE_COLUMN_INFO = "SHOW_TABLE_COLUMN_INFO"; - private final Function mapper; private final ConnectionSession session; private final List sqlTuples; @@ -78,6 +76,10 @@ public AsyncExecuteContext(ConnectionSession session, List sqlTuples, this.sessionIds = initSessionIds(session); } + public boolean isFinished() { + return future.isDone(); + } + public int getTotal() { return sqlTuples.size(); } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java index 58f046885f..72287014c1 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java @@ -17,8 +17,6 @@ import java.util.List; -import com.oceanbase.odc.core.sql.execute.model.SqlExecuteStatus; - import lombok.Data; /** @@ -31,11 +29,11 @@ public class AsyncExecuteResultResp { private String traceId; private int total; private int count; - private SqlExecuteStatus status; + private boolean finished; private String sql; - public AsyncExecuteResultResp(SqlExecuteStatus status, AsyncExecuteContext context) { - this.status = status; + public AsyncExecuteResultResp(boolean finished, AsyncExecuteContext context) { + this.finished = finished; results = context.getResults(); traceId = context.getCurrentTraceId(); total = context.getTotal(); From caa5728f6c3dc24932479edda7c7f20d3ec86f76 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Tue, 23 Apr 2024 11:57:18 +0800 Subject: [PATCH 06/16] fix typo --- .../oceanbase/odc/service/session/ConnectConsoleService.java | 3 +++ .../oceanbase/odc/service/session/OdcStatementCallBack.java | 4 ++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java index d6b1a3bff9..f56d7b0805 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java @@ -384,6 +384,9 @@ public SqlAsyncExecuteResp executeV2(@NotNull String sessionId, } return res; }, authenticationFacade.currentUser(), EXECUTING_CONTEXT_WAIT_MILLIS); + executeContext.getTraceIdHooks().add((session, traceId) -> { + profileManager.addProfile(session, traceId); + }); OdcStatementCallBack statementCallBack = new OdcStatementCallBack(sqlTuples, connectionSession, request.getAutoCommit(), queryLimit, stopOnError, executeContext); diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java index 73ec59d5c9..ab5334d808 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java @@ -145,7 +145,7 @@ public OdcStatementCallBack(@NonNull List sqls, @NonNull ConnectionSes public OdcStatementCallBack(@NonNull List sqls, @NonNull ConnectionSession connectionSession, Boolean autoCommit, Integer queryLimit, boolean stopWhenError) { - this(sqls, connectionSession, autoCommit, queryLimit, true, null); + this(sqls, connectionSession, autoCommit, queryLimit, stopWhenError, null); } public OdcStatementCallBack(@NonNull List sqls, @NonNull ConnectionSession connectionSession, @@ -471,7 +471,7 @@ private SqlExecTime getTraceIdAndAndSetStage(Statement statement, TraceWatch tra try { StopWatch stopWatch = StopWatch.createStarted(); String version = ConnectionSessionUtil.getVersion(connectionSession); - SqlExecTime executeDetails = new SqlExecTime();; + SqlExecTime executeDetails = new SqlExecTime(); if (useFullLinkTrace && VersionUtils.isGreaterThanOrEqualsTo(version, "4.1") && connectionSession.getDialectType().isOceanbase()) { try { From 180fc84b3ad2959373a32f29aea421b90de78990 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Tue, 23 Apr 2024 12:44:09 +0800 Subject: [PATCH 07/16] format code --- .../odc/service/connection/util/ConnectionInfoUtil.java | 4 ++-- .../oceanbase/odc/service/session/ConnectConsoleService.java | 3 --- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java index 2181e81b08..9764edfecc 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java @@ -15,7 +15,6 @@ */ package com.oceanbase.odc.service.connection.util; -import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Statement; @@ -67,7 +66,8 @@ public static void initConnectionId(@NonNull Statement statement, @NonNull Conne connectionSession.setAttribute(ConnectionSessionConstants.CONNECTION_ID_KEY, sessionId); if (connectionSession.getDialectType().isOceanbase() && VersionUtils.isGreaterThanOrEqualsTo( ConnectionSessionUtil.getVersion(connectionSession), "4.2")) { - String proxySessId = OBUtils.queryOBProxySessId(statement, connectionSession.getDialectType(), sessionId); + String proxySessId = + OBUtils.queryOBProxySessId(statement, connectionSession.getDialectType(), sessionId); connectionSession.setAttribute(ConnectionSessionConstants.OB_PROXY_SESSID_KEY, proxySessId); } } catch (Exception exception) { diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java index f56d7b0805..d6b1a3bff9 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java @@ -384,9 +384,6 @@ public SqlAsyncExecuteResp executeV2(@NotNull String sessionId, } return res; }, authenticationFacade.currentUser(), EXECUTING_CONTEXT_WAIT_MILLIS); - executeContext.getTraceIdHooks().add((session, traceId) -> { - profileManager.addProfile(session, traceId); - }); OdcStatementCallBack statementCallBack = new OdcStatementCallBack(sqlTuples, connectionSession, request.getAutoCommit(), queryLimit, stopOnError, executeContext); From a3179d3ef15d926eff5367158055a66af75bbc83 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Tue, 23 Apr 2024 17:13:55 +0800 Subject: [PATCH 08/16] add SqlExecutionListener to listen execute status --- .../session/ConnectConsoleService.java | 33 +++--- .../service/session/OBExecutionListener.java | 88 +++++++++++++++ .../service/session/SqlExecutionListener.java | 40 +++++++ .../session/model/AsyncExecuteContext.java | 100 ++++-------------- .../session/model/AsyncExecuteResultResp.java | 10 +- 5 files changed, 171 insertions(+), 100 deletions(-) create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/session/OBExecutionListener.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/session/SqlExecutionListener.java diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java index d6b1a3bff9..0e0f7bbdc1 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java @@ -373,17 +373,7 @@ public SqlAsyncExecuteResp executeV2(@NotNull String sessionId, Objects.nonNull(request.getContinueExecutionOnError()) ? request.getContinueExecutionOnError() : userConfigFacade.isContinueExecutionOnError(); boolean stopOnError = !continueExecutionOnError; - AsyncExecuteContext executeContext = - new AsyncExecuteContext<>(connectionSession, sqlTuples, jdbcResult -> { - SqlExecuteResult res = generateResult(connectionSession, jdbcResult, context); - try (TraceStage stage = - res.getSqlTuple().getSqlWatch().start(SqlExecuteStages.SQL_AFTER_CHECK)) { - sqlInterceptService.afterCompletion(res, connectionSession, context); - } catch (Exception e) { - throw new IllegalStateException(e); - } - return res; - }, authenticationFacade.currentUser(), EXECUTING_CONTEXT_WAIT_MILLIS); + AsyncExecuteContext executeContext = new AsyncExecuteContext(sqlTuples, context); OdcStatementCallBack statementCallBack = new OdcStatementCallBack(sqlTuples, connectionSession, request.getAutoCommit(), queryLimit, stopOnError, executeContext); @@ -398,6 +388,9 @@ public SqlAsyncExecuteResp executeV2(@NotNull String sessionId, statementCallBack.setMaxCachedSize(sessionProperties.getResultSetMaxCachedSize()); statementCallBack.setMaxCachedLines(sessionProperties.getResultSetMaxCachedLines()); statementCallBack.setLocale(LocaleContextHolder.getLocale()); + if (connectionSession.getDialectType().isOceanbase() && sqlTuples.size() <= 10) { + statementCallBack.getListeners().add(new OBExecutionListener(connectionSession)); + } Future> futureResult = connectionSession.getAsyncJdbcExecutor( ConnectionSessionConstants.CONSOLE_DS_KEY).execute(statementCallBack); @@ -445,16 +438,26 @@ public List getAsyncResult(@NotNull String sessionId, String r public AsyncExecuteResultResp getAsyncResultV2(@NotNull String sessionId, String requestId) { PreConditions.validArgumentState(Objects.nonNull(requestId), ErrorCodes.SqlRegulationRuleBlocked, null, null); ConnectionSession connectionSession = sessionService.nullSafeGet(sessionId); - AsyncExecuteContext executeContext = + AsyncExecuteContext context = (AsyncExecuteContext) ConnectionSessionUtil.getExecuteContext(connectionSession, requestId); - if (executeContext.isFinished()) { + List resultList = context.getResults(); + List results = resultList.stream().map(jdbcGeneralResult -> { + SqlExecuteResult result = generateResult(connectionSession, jdbcGeneralResult, context.getContextMap()); + try (TraceStage stage = result.getSqlTuple().getSqlWatch().start(SqlExecuteStages.SQL_AFTER_CHECK)) { + sqlInterceptService.afterCompletion(result, connectionSession, context.getContextMap()); + } catch (Exception e) { + throw new IllegalStateException(e); + } + return result; + }).collect(Collectors.toList()); + if (context.isFinished()) { ConnectionSessionUtil.removeExecuteContext(connectionSession, requestId); - return new AsyncExecuteResultResp(true, executeContext); + return new AsyncExecuteResultResp(true, context, results); } else { if (log.isDebugEnabled()) { log.debug("Get sql execution result timed out, sessionId={}, requestId={}", sessionId, requestId); } - return new AsyncExecuteResultResp(false, executeContext); + return new AsyncExecuteResultResp(false, context, results); } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OBExecutionListener.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OBExecutionListener.java new file mode 100644 index 0000000000..ca1c835b16 --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OBExecutionListener.java @@ -0,0 +1,88 @@ +/* + * Copyright (c) 2023 OceanBase. + * + * Licensed 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 com.oceanbase.odc.service.session; + +import static com.oceanbase.odc.core.session.ConnectionSessionConstants.BACKEND_DS_KEY; +import static com.oceanbase.odc.core.session.ConnectionSessionConstants.CONSOLE_DS_KEY; + +import java.util.Collections; +import java.util.List; + +import org.apache.commons.collections4.CollectionUtils; +import org.springframework.jdbc.core.StatementCallback; + +import com.oceanbase.odc.common.util.StringUtils; +import com.oceanbase.odc.common.util.VersionUtils; +import com.oceanbase.odc.core.session.ConnectionSession; +import com.oceanbase.odc.core.session.ConnectionSessionUtil; +import com.oceanbase.odc.core.sql.execute.model.JdbcGeneralResult; +import com.oceanbase.odc.core.sql.execute.model.SqlTuple; +import com.oceanbase.odc.core.sql.util.OBUtils; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; + +/** + * @author: liuyizhuo.lyz + * @date: 2024/4/23 + */ +public class OBExecutionListener implements SqlExecutionListener { + private static final Long DEFAULT_QUERY_TRACE_ID_WAIT_MILLIS = 1100L; + private static final String ENABLE_QUERY_PROFILE_VERSION = "4.2"; + + private final ConnectionSession session; + private final List sessionIds; + + public OBExecutionListener(ConnectionSession session) { + this.session = session; + sessionIds = getSessionIds(); + } + + @Override + public void onExecutionStart(SqlTuple sqlTuple, AsyncExecuteContext context) {} + + @Override + public void onExecutionEnd(SqlTuple sqlTuple, List results, AsyncExecuteContext context) {} + + @Override + public void onExecutionCanceled(SqlTuple sqlTuple, List results, AsyncExecuteContext context) {} + + public void onExecutionStartAfter(SqlTuple sqlTuple, AsyncExecuteContext context) { + if (CollectionUtils.isEmpty(sessionIds)) { + return; + } + String traceId = session.getSyncJdbcExecutor(BACKEND_DS_KEY).execute((StatementCallback) stmt -> OBUtils + .queryTraceIdFromASH(stmt, sessionIds, session.getConnectType())); + if (traceId != null) { + context.setCurrentExecutingSqlTraceId(traceId); + } + } + + @Override + public Long getOnExecutionStartAfterMillis() { + return DEFAULT_QUERY_TRACE_ID_WAIT_MILLIS; + } + + private List getSessionIds() { + if (VersionUtils.isLessThan(ConnectionSessionUtil.getVersion(session), ENABLE_QUERY_PROFILE_VERSION)) { + return Collections.emptyList(); + } + String proxySessId = ConnectionSessionUtil.getConsoleConnectionProxySessId(session); + if (StringUtils.isEmpty(proxySessId)) { + return Collections.singletonList(ConnectionSessionUtil.getConsoleConnectionId(session)); + } + return session.getSyncJdbcExecutor(CONSOLE_DS_KEY).execute((StatementCallback>) stmt -> OBUtils + .querySessionIdsByProxySessId(stmt, proxySessId, session.getConnectType())); + } +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/SqlExecutionListener.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/SqlExecutionListener.java new file mode 100644 index 0000000000..1c07102214 --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/SqlExecutionListener.java @@ -0,0 +1,40 @@ +/* + * Copyright (c) 2023 OceanBase. + * + * Licensed 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 com.oceanbase.odc.service.session; + +import java.util.List; + +import com.oceanbase.odc.core.sql.execute.model.JdbcGeneralResult; +import com.oceanbase.odc.core.sql.execute.model.SqlTuple; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; + +/** + * @author: liuyizhuo.lyz + * @date: 2024/4/23 + */ +public interface SqlExecutionListener { + + void onExecutionStart(SqlTuple sqlTuple, AsyncExecuteContext context); + + void onExecutionEnd(SqlTuple sqlTuple, List results, AsyncExecuteContext context); + + void onExecutionCanceled(SqlTuple sqlTuple, List results, AsyncExecuteContext context); + + void onExecutionStartAfter(SqlTuple sqlTuple, AsyncExecuteContext context); + + Long getOnExecutionStartAfterMillis(); + +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java index 33f8e2ce05..8102935dad 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java @@ -15,29 +15,15 @@ */ package com.oceanbase.odc.service.session.model; -import static com.oceanbase.odc.core.session.ConnectionSessionConstants.BACKEND_DS_KEY; -import static com.oceanbase.odc.core.session.ConnectionSessionConstants.CONSOLE_DS_KEY; - import java.util.ArrayList; -import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.Queue; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.Future; -import java.util.function.BiConsumer; -import java.util.function.Function; - -import org.springframework.jdbc.core.StatementCallback; -import com.oceanbase.odc.common.util.StringUtils; -import com.oceanbase.odc.common.util.VersionUtils; -import com.oceanbase.odc.core.session.ConnectionSession; -import com.oceanbase.odc.core.session.ConnectionSessionUtil; import com.oceanbase.odc.core.sql.execute.model.JdbcGeneralResult; import com.oceanbase.odc.core.sql.execute.model.SqlTuple; -import com.oceanbase.odc.core.sql.util.OBUtils; -import com.oceanbase.odc.service.iam.model.User; -import com.oceanbase.odc.service.iam.util.SecurityContextUtils; import lombok.Getter; import lombok.Setter; @@ -48,95 +34,49 @@ * @date 2024/4/15 */ @Getter +@Setter @Slf4j -public class AsyncExecuteContext { - private final Function mapper; - private final ConnectionSession session; +public class AsyncExecuteContext { private final List sqlTuples; - private final Queue results = new ConcurrentLinkedQueue<>(); - private final List> traceIdHooks = new ArrayList<>(); - private final User user; - private final Long waitTimeMillis; - private final List sessionIds; + private final Queue results = new ConcurrentLinkedQueue<>(); + private final Map contextMap; - @Setter private Future> future; - private String currentTraceId; - @Setter - private String sql; - private int count = 0; + private String currentExecutingSqlTraceId; + private String currentExecutingSql; + private int totalSqlExecutedCount = 0; - public AsyncExecuteContext(ConnectionSession session, List sqlTuples, - Function mapper, User user, Long waitTimeMillis) { - this.session = session; + public AsyncExecuteContext(List sqlTuples, Map contextMap) { this.sqlTuples = sqlTuples; - this.mapper = mapper; - this.user = user; - this.waitTimeMillis = waitTimeMillis; - this.sessionIds = initSessionIds(session); + this.contextMap = contextMap; } public boolean isFinished() { return future.isDone(); } + public void addCount() { + totalSqlExecutedCount++; + } + public int getTotal() { return sqlTuples.size(); } - public List getResults() { - List copiedResults = new ArrayList<>(); + public List getResults() { + List copiedResults = new ArrayList<>(); while (!results.isEmpty()) { copiedResults.add(results.poll()); } return copiedResults; } - public void setCurrentTraceId(String traceId) { - currentTraceId = traceId; - for (BiConsumer hook : traceIdHooks) { - hook.accept(session, traceId); - } - } - - public void onQueryStart(String sql) { - count++; - this.sql = sql; + public void addResult(JdbcGeneralResult result) { + this.results.add(result); } - public void onQueryExecuting() { - String traceId = session.getSyncJdbcExecutor(BACKEND_DS_KEY) - .execute((StatementCallback) stmt -> OBUtils - .queryTraceIdFromASH(stmt, sessionIds, session.getConnectType())); - if (traceId != null) { - setCurrentTraceId(traceId); - } - } - - public void onQueryFinish(List results) { - SecurityContextUtils.setCurrentUser(user); - try { - for (JdbcGeneralResult result : results) { - this.results.add(mapper.apply(result)); - for (BiConsumer hook : traceIdHooks) { - hook.accept(session, result.getTraceId()); - } - } - } finally { - SecurityContextUtils.clear(); - } + public void addResults(List results) { + this.results.addAll(results); } - private List initSessionIds(ConnectionSession session) { - if (VersionUtils.isLessThan(ConnectionSessionUtil.getVersion(session), "4.2") - || sqlTuples.size() > 10) { - return null; - } - String proxySessId = ConnectionSessionUtil.getConsoleConnectionProxySessId(session); - if (StringUtils.isEmpty(proxySessId)) { - return Collections.singletonList(ConnectionSessionUtil.getConsoleConnectionId(session)); - } - return session.getSyncJdbcExecutor(CONSOLE_DS_KEY).execute((StatementCallback>) stmt -> OBUtils - .querySessionIdsByProxySessId(stmt, proxySessId, session.getConnectType())); - } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java index 72287014c1..eb6edac6e9 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java @@ -32,12 +32,12 @@ public class AsyncExecuteResultResp { private boolean finished; private String sql; - public AsyncExecuteResultResp(boolean finished, AsyncExecuteContext context) { + public AsyncExecuteResultResp(boolean finished, AsyncExecuteContext context, List results) { this.finished = finished; - results = context.getResults(); - traceId = context.getCurrentTraceId(); + this.results = results; + traceId = context.getCurrentExecutingSqlTraceId(); total = context.getTotal(); - count = context.getCount(); - sql = context.getSql(); + count = context.getTotalSqlExecutedCount(); + sql = context.getCurrentExecutingSql(); } } From d93beb4e66c95ac5cfd3c1089444121cb34ba8b4 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Tue, 23 Apr 2024 17:16:57 +0800 Subject: [PATCH 09/16] add SqlExecutionListener to listen execute status --- .../service/session/OdcStatementCallBack.java | 59 +++++++++++-------- 1 file changed, 35 insertions(+), 24 deletions(-) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java index ab5334d808..4d9d8cca0f 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java @@ -30,6 +30,7 @@ import java.sql.Statement; import java.util.ArrayList; import java.util.Collections; +import java.util.Comparator; import java.util.LinkedList; import java.util.List; import java.util.Locale; @@ -83,7 +84,6 @@ import com.oceanbase.odc.core.sql.util.OBUtils; import com.oceanbase.odc.service.plugin.ConnectionPluginUtil; import com.oceanbase.odc.service.session.model.AsyncExecuteContext; -import com.oceanbase.odc.service.session.model.SqlExecuteResult; import lombok.Getter; import lombok.NonNull; @@ -120,7 +120,8 @@ public class OdcStatementCallBack implements StatementCallback context; + private final AsyncExecuteContext context; + private final List listeners = new ArrayList<>(); private final ExecutorService executor = Executors.newSingleThreadExecutor(); @Setter private boolean useFullLinkTrace = false; @@ -149,8 +150,7 @@ public OdcStatementCallBack(@NonNull List sqls, @NonNull ConnectionSes } public OdcStatementCallBack(@NonNull List sqls, @NonNull ConnectionSession connectionSession, - Boolean autoCommit, Integer queryLimit, boolean stopWhenError, - AsyncExecuteContext context) { + Boolean autoCommit, Integer queryLimit, boolean stopWhenError, AsyncExecuteContext context) { this.sqls = sqls; this.autoCommit = autoCommit == null ? connectionSession.getDefaultAutoCommit() : autoCommit; this.connectType = connectionSession.getConnectType(); @@ -175,15 +175,16 @@ public List doInStatement(Statement statement) throws SQLExce } if (ConnectionSessionUtil.isConsoleSessionReset(connectionSession)) { ConnectionSessionUtil.setConsoleSessionResetFlag(connectionSession, false); - List results = this.sqls.stream().map(sqlTuple -> { + return this.sqls.stream().map(sqlTuple -> { JdbcGeneralResult result = JdbcGeneralResult.canceledResult(sqlTuple); result.setConnectionReset(true); + if (context != null) { + context.addResult(result); + } + listeners.forEach( + listener -> listener.onExecutionCanceled(sqlTuple, Collections.singletonList(result), context)); return result; }).collect(Collectors.toList()); - if (context != null) { - context.onQueryFinish(results); - } - return results; } boolean currentAutoCommit = statement.getConnection().getAutoCommit(); List returnVal = new LinkedList<>(); @@ -203,8 +204,10 @@ public List doInStatement(Statement statement) throws SQLExce } } if (context != null) { - context.onQueryStart(sqlTuple.getExecutedSql()); + context.setCurrentExecutingSql(sqlTuple.getExecutedSql()); + context.addCount(); } + listeners.forEach(listener -> listener.onExecutionStart(sqlTuple, context)); try { applyConnectionSettings(statement); } catch (Exception e) { @@ -215,36 +218,44 @@ public List doInStatement(Statement statement) throws SQLExce if (Thread.currentThread().isInterrupted() || ConnectionSessionUtil.isConsoleSessionKillQuery(connectionSession)) { executeResults = Collections.singletonList(JdbcGeneralResult.canceledResult(sqlTuple)); + listeners.forEach(listener -> listener.onExecutionCanceled(sqlTuple, executeResults, context)); } else { CountDownLatch latch = new CountDownLatch(1); - if (context != null) { - handle = executor.submit(() -> { - try { - if (!latch.await(context.getWaitTimeMillis(), TimeUnit.MILLISECONDS)) { - context.onQueryExecuting(); - } - } catch (InterruptedException e) { - log.warn("Failed to call back.", e); + handle = executor.submit(() -> { + long startTs = System.currentTimeMillis(); + List sortedListeners = listeners.stream() + .filter(listener -> listener.getOnExecutionStartAfterMillis() != null + && listener.getOnExecutionStartAfterMillis() > 0) + .sorted(Comparator + .comparingLong(SqlExecutionListener::getOnExecutionStartAfterMillis)) + .collect(Collectors.toList()); + for (SqlExecutionListener listener : sortedListeners) { + long waitTs = System.currentTimeMillis() - startTs; + Long expectedTs = listener.getOnExecutionStartAfterMillis(); + if (waitTs > expectedTs || !latch.await(expectedTs - waitTs, TimeUnit.MILLISECONDS)) { + listener.onExecutionStartAfter(sqlTuple, context); } - return null; - }); - } + } + return null; + }); executeResults = doExecuteSql(statement, sqlTuple, latch); + listeners.forEach(listener -> listener.onExecutionEnd(sqlTuple, executeResults, context)); } } else { executeResults = Collections.singletonList(JdbcGeneralResult.canceledResult(sqlTuple)); + listeners.forEach(listener -> listener.onExecutionCanceled(sqlTuple, executeResults, context)); } - returnVal.addAll(executeResults); if (context != null) { - context.onQueryFinish(executeResults); + context.addResults(executeResults); } + returnVal.addAll(executeResults); } Optional failed = returnVal .stream().filter(r -> r.getStatus() == SqlExecuteStatus.FAILED).findFirst(); if (failed.isPresent()) { throw failed.get().getThrown(); } - } catch (Exception e){ + } catch (Exception e) { try { ConnectionSessionUtil.logSocketInfo(statement.getConnection(), "console error"); } catch (Exception exception) { From 7df93978fc71ad66df44e0884c796e0b8a81b778 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Tue, 23 Apr 2024 17:18:09 +0800 Subject: [PATCH 10/16] set trace id to null when starting --- .../com/oceanbase/odc/service/session/OdcStatementCallBack.java | 1 + 1 file changed, 1 insertion(+) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java index 4d9d8cca0f..86fe217807 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java @@ -206,6 +206,7 @@ public List doInStatement(Statement statement) throws SQLExce if (context != null) { context.setCurrentExecutingSql(sqlTuple.getExecutedSql()); context.addCount(); + context.setCurrentExecutingSqlTraceId(null); } listeners.forEach(listener -> listener.onExecutionStart(sqlTuple, context)); try { From 769a879f19b29a7126f7e83c116961006d7410ce Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Tue, 23 Apr 2024 17:34:02 +0800 Subject: [PATCH 11/16] response to CR --- .../oceanbase/odc/service/session/OdcStatementCallBack.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java index 86fe217807..e6ef73b9b4 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java @@ -373,7 +373,7 @@ protected List doExecuteSql(Statement statement, SqlTuple sql boolean isResultSet; try { isResultSet = statement.execute(sql); - } catch (SQLException e) { + } catch (Exception e) { return handleException(e, statement, sqlTuple); } latch.countDown(); @@ -409,7 +409,7 @@ protected List doExecuteSql(Statement statement, SqlTuple sql boolean isResultSet; try { isResultSet = preparedStatement.execute(); - } catch (SQLException e) { + } catch (Exception e) { return handleException(e, statement, sqlTuple); } latch.countDown(); From 6997eec39751fd9b49cd9ac2754236907abb7634 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Tue, 23 Apr 2024 17:51:39 +0800 Subject: [PATCH 12/16] fix an error --- .../oceanbase/odc/service/session/OdcStatementCallBack.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java index e6ef73b9b4..ecec407fea 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java @@ -233,8 +233,10 @@ public List doInStatement(Statement statement) throws SQLExce for (SqlExecutionListener listener : sortedListeners) { long waitTs = System.currentTimeMillis() - startTs; Long expectedTs = listener.getOnExecutionStartAfterMillis(); - if (waitTs > expectedTs || !latch.await(expectedTs - waitTs, TimeUnit.MILLISECONDS)) { + if (!latch.await(expectedTs - waitTs, TimeUnit.MILLISECONDS)) { listener.onExecutionStartAfter(sqlTuple, context); + } else { + break; } } return null; From ce6eced419d465356813693ae5909777824feaf6 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Wed, 24 Apr 2024 14:27:57 +0800 Subject: [PATCH 13/16] response to CR comments --- .../session/NlsFormatInterceptorTest.java | 25 ++-- .../datasecurity/DataMaskingInterceptor.java | 6 +- .../integration/ExternalSqlInterceptor.java | 6 +- .../session/ConnectConsoleService.java | 50 ++++---- .../service/session/OBExecutionListener.java | 2 +- .../service/session/OdcStatementCallBack.java | 110 ++++++++++++------ .../service/session/SqlExecutionListener.java | 2 +- .../BaseTimeConsumingInterceptor.java | 10 +- .../DatabasePermissionInterceptor.java | 5 +- .../interceptor/NlsFormatInterceptor.java | 6 +- .../interceptor/SqlCheckInterceptor.java | 19 +-- .../interceptor/SqlConsoleInterceptor.java | 14 ++- .../interceptor/SqlExecuteInterceptor.java | 7 +- .../SqlExecuteInterceptorService.java | 5 +- .../SwitchDatabaseInterceptor.java | 4 +- .../session/model/AsyncExecuteContext.java | 25 ++-- .../session/model/AsyncExecuteResultResp.java | 8 +- 17 files changed, 179 insertions(+), 125 deletions(-) diff --git a/server/integration-test/src/test/java/com/oceanbase/odc/service/session/NlsFormatInterceptorTest.java b/server/integration-test/src/test/java/com/oceanbase/odc/service/session/NlsFormatInterceptorTest.java index 13b3ff6fb9..5bfdf17fcf 100644 --- a/server/integration-test/src/test/java/com/oceanbase/odc/service/session/NlsFormatInterceptorTest.java +++ b/server/integration-test/src/test/java/com/oceanbase/odc/service/session/NlsFormatInterceptorTest.java @@ -27,6 +27,7 @@ import com.oceanbase.odc.core.sql.execute.model.SqlExecuteStatus; import com.oceanbase.odc.core.sql.execute.model.SqlTuple; import com.oceanbase.odc.service.session.interceptor.NlsFormatInterceptor; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; import com.oceanbase.odc.service.session.model.SqlExecuteResult; /** @@ -43,7 +44,7 @@ public void afterCompletion_mysql_notingSet() throws Exception { ConnectionSession session = getConnectionSession(ConnectType.OB_MYSQL); NlsFormatInterceptor interceptor = new NlsFormatInterceptor(); SqlExecuteResult r = getResponse("set session nls_date_format='DD-MON-RR'", SqlExecuteStatus.SUCCESS); - interceptor.afterCompletion(r, session, new HashMap<>()); + interceptor.afterCompletion(r, session, getContext()); Assert.assertNull(ConnectionSessionUtil.getNlsDateFormat(session)); } @@ -52,7 +53,7 @@ public void afterCompletion_failedSqlResult_notingSet() throws Exception { ConnectionSession session = getConnectionSession(ConnectType.OB_ORACLE); NlsFormatInterceptor interceptor = new NlsFormatInterceptor(); SqlExecuteResult r = getResponse("set session nls_date_format='DD-MON-RR'", SqlExecuteStatus.FAILED); - interceptor.afterCompletion(r, session, new HashMap<>()); + interceptor.afterCompletion(r, session, getContext()); Assert.assertNull(ConnectionSessionUtil.getNlsDateFormat(session)); } @@ -65,7 +66,7 @@ public void afterCompletion_multiSqls_notingSet() throws Exception { + "begin\n" + "dbms_output.put_line('aaaa');\n" + "end;", SqlExecuteStatus.SUCCESS); - interceptor.afterCompletion(r, session, new HashMap<>()); + interceptor.afterCompletion(r, session, getContext()); Assert.assertNull(ConnectionSessionUtil.getNlsDateFormat(session)); } @@ -74,7 +75,7 @@ public void afterCompletion_noSetVarExists_notingSet() throws Exception { ConnectionSession session = getConnectionSession(ConnectType.OB_ORACLE); NlsFormatInterceptor interceptor = new NlsFormatInterceptor(); SqlExecuteResult r = getResponse("-- comment\nselect 123 from dual;", SqlExecuteStatus.SUCCESS); - interceptor.afterCompletion(r, session, new HashMap<>()); + interceptor.afterCompletion(r, session, getContext()); Assert.assertNull(ConnectionSessionUtil.getNlsDateFormat(session)); } @@ -85,7 +86,7 @@ public void afterCompletion_commentWithSetVar_setSucceed() throws Exception { String expect = "DD-MON-RR"; SqlExecuteResult r = getResponse("-- comment\nset session nls_date_format='" + expect + "';", SqlExecuteStatus.SUCCESS); - interceptor.afterCompletion(r, session, new HashMap<>()); + interceptor.afterCompletion(r, session, getContext()); Assert.assertEquals(expect, ConnectionSessionUtil.getNlsDateFormat(session)); } @@ -96,7 +97,7 @@ public void afterCompletion_multiCommentsWithSetVar_setSucceed() throws Exceptio String expect = "DD-MON-RR"; SqlExecuteResult r = getResponse("/*asdasdasd*/ set session nls_date_format='" + expect + "';", SqlExecuteStatus.SUCCESS); - interceptor.afterCompletion(r, session, new HashMap<>()); + interceptor.afterCompletion(r, session, getContext()); Assert.assertEquals(expect, ConnectionSessionUtil.getNlsDateFormat(session)); } @@ -107,7 +108,7 @@ public void afterCompletion_nlsTimestampFormat_setSucceed() throws Exception { String expect = "DD-MON-RR"; SqlExecuteResult r = getResponse("/*asdasdasd*/ set session nls_timestamp_format='" + expect + "';", SqlExecuteStatus.SUCCESS); - interceptor.afterCompletion(r, session, new HashMap<>()); + interceptor.afterCompletion(r, session, getContext()); Assert.assertEquals(expect, ConnectionSessionUtil.getNlsTimestampFormat(session)); } @@ -118,7 +119,7 @@ public void afterCompletion_nlsTimestampTZFormat_setSucceed() throws Exception { String expect = "DD-MON-RR"; SqlExecuteResult r = getResponse("/*asdasdasd*/ set session nls_timestamp_tz_format='" + expect + "';", SqlExecuteStatus.SUCCESS); - interceptor.afterCompletion(r, session, new HashMap<>()); + interceptor.afterCompletion(r, session, getContext()); Assert.assertEquals(expect, ConnectionSessionUtil.getNlsTimestampTZFormat(session)); } @@ -129,7 +130,7 @@ public void afterCompletion_setGlobal_nothingSet() throws Exception { String expect = "DD-MON-RR"; SqlExecuteResult r = getResponse("/*asdasdasd*/ set global nls_timestamp_tz_format='" + expect + "';", SqlExecuteStatus.SUCCESS); - interceptor.afterCompletion(r, session, new HashMap<>()); + interceptor.afterCompletion(r, session, getContext()); Assert.assertNull(ConnectionSessionUtil.getNlsTimestampTZFormat(session)); } @@ -140,7 +141,7 @@ public void afterCompletion_alterSession_setSucceed() throws Exception { String expect = "DD-MON-RR"; SqlExecuteResult r = getResponse("/*asdsd*/ alter session \n\t\r set \"nls_date_format\"='" + expect + "';", SqlExecuteStatus.SUCCESS); - interceptor.afterCompletion(r, session, new HashMap<>()); + interceptor.afterCompletion(r, session, getContext()); Assert.assertEquals(expect, ConnectionSessionUtil.getNlsDateFormat(session)); } @@ -159,4 +160,8 @@ private ConnectionSession getConnectionSession(ConnectType type) { return session; } + private AsyncExecuteContext getContext() { + return new AsyncExecuteContext(null, new HashMap<>()); + } + } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/DataMaskingInterceptor.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/DataMaskingInterceptor.java index 2cc642cf6f..6bc34cea53 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/DataMaskingInterceptor.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/DataMaskingInterceptor.java @@ -20,7 +20,6 @@ import java.util.Comparator; import java.util.HashSet; import java.util.List; -import java.util.Map; import java.util.Objects; import java.util.Set; import java.util.stream.Collectors; @@ -40,6 +39,7 @@ import com.oceanbase.odc.service.datasecurity.util.DataMaskingUtil; import com.oceanbase.odc.service.db.browser.DBSchemaAccessors; import com.oceanbase.odc.service.session.interceptor.BaseTimeConsumingInterceptor; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; import com.oceanbase.odc.service.session.model.DBResultSetMetaData; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteReq; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteResp; @@ -64,14 +64,14 @@ public class DataMaskingInterceptor extends BaseTimeConsumingInterceptor { @Override public boolean preHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyncExecuteResp response, - @NonNull ConnectionSession session, @NonNull Map context) { + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) { return true; } @Override @SuppressWarnings("all") public void doAfterCompletion(@NonNull SqlExecuteResult response, @NonNull ConnectionSession session, - @NonNull Map context) throws Exception { + @NonNull AsyncExecuteContext context) throws Exception { // TODO: May intercept sensitive column operation (WHERE / ORDER BY / HAVING) if (!maskingService.isMaskingEnabled()) { return; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/integration/ExternalSqlInterceptor.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/integration/ExternalSqlInterceptor.java index 37976a3c62..84602fd2c0 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/integration/ExternalSqlInterceptor.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/integration/ExternalSqlInterceptor.java @@ -16,7 +16,6 @@ package com.oceanbase.odc.service.integration; import java.util.List; -import java.util.Map; import java.util.Map.Entry; import java.util.Objects; import java.util.Optional; @@ -48,6 +47,7 @@ import com.oceanbase.odc.service.regulation.ruleset.model.Rule.RuleViolation; import com.oceanbase.odc.service.regulation.ruleset.model.SqlConsoleRules; import com.oceanbase.odc.service.session.interceptor.BaseTimeConsumingInterceptor; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteReq; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteResp; import com.oceanbase.odc.service.session.model.SqlExecuteResult; @@ -81,7 +81,7 @@ public class ExternalSqlInterceptor extends BaseTimeConsumingInterceptor { @Override public boolean doPreHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyncExecuteResp response, - @NonNull ConnectionSession session, @NonNull Map context) { + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) { Long ruleSetId = ConnectionSessionUtil.getRuleSetId(session); if (Objects.isNull(ruleSetId) || isIndividualTeam()) { return true; @@ -147,7 +147,7 @@ public boolean doPreHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyn @Override public void afterCompletion(@NonNull SqlExecuteResult response, - @NonNull ConnectionSession session, @NonNull Map context) {} + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) {} @Override diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java index 0e0f7bbdc1..0140483656 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java @@ -128,7 +128,6 @@ public class ConnectConsoleService { public static final int DEFAULT_GET_RESULT_TIMEOUT_SECONDS = 3; public static final String SHOW_TABLE_COLUMN_INFO = "SHOW_TABLE_COLUMN_INFO"; - private static final Long EXECUTING_CONTEXT_WAIT_MILLIS = 1100L; @Autowired private ConnectSessionService sessionService; @@ -270,11 +269,12 @@ public SqlAsyncExecuteResp execute(@NotNull String sessionId, context.put(SHOW_TABLE_COLUMN_INFO, request.getShowTableColumnInfo()); context.put(SqlCheckInterceptor.NEED_SQL_CHECK_KEY, needSqlRuleCheck); context.put(SqlConsoleInterceptor.NEED_SQL_CONSOLE_CHECK, needSqlRuleCheck); + AsyncExecuteContext executeContext = new AsyncExecuteContext(sqlTuples, context); List stages = sqlTuples.stream() .map(s -> s.getSqlWatch().start(SqlExecuteStages.SQL_PRE_CHECK)) .collect(Collectors.toList()); try { - if (!sqlInterceptService.preHandle(request, response, connectionSession, context)) { + if (!sqlInterceptService.preHandle(request, response, connectionSession, executeContext)) { return response; } } finally { @@ -352,11 +352,12 @@ public SqlAsyncExecuteResp executeV2(@NotNull String sessionId, context.put(SHOW_TABLE_COLUMN_INFO, request.getShowTableColumnInfo()); context.put(SqlCheckInterceptor.NEED_SQL_CHECK_KEY, needSqlRuleCheck); context.put(SqlConsoleInterceptor.NEED_SQL_CONSOLE_CHECK, needSqlRuleCheck); + AsyncExecuteContext executeContext = new AsyncExecuteContext(sqlTuples, context); List stages = sqlTuples.stream() .map(s -> s.getSqlWatch().start(SqlExecuteStages.SQL_PRE_CHECK)) .collect(Collectors.toList()); try { - if (!sqlInterceptService.preHandle(request, response, connectionSession, context)) { + if (!sqlInterceptService.preHandle(request, response, connectionSession, executeContext)) { return response; } } finally { @@ -373,7 +374,6 @@ public SqlAsyncExecuteResp executeV2(@NotNull String sessionId, Objects.nonNull(request.getContinueExecutionOnError()) ? request.getContinueExecutionOnError() : userConfigFacade.isContinueExecutionOnError(); boolean stopOnError = !continueExecutionOnError; - AsyncExecuteContext executeContext = new AsyncExecuteContext(sqlTuples, context); OdcStatementCallBack statementCallBack = new OdcStatementCallBack(sqlTuples, connectionSession, request.getAutoCommit(), queryLimit, stopOnError, executeContext); @@ -415,10 +415,10 @@ public List getAsyncResult(@NotNull String sessionId, String r Map context = ConnectionSessionUtil.getFutureJdbcContext(connectionSession, requestId); ConnectionSessionUtil.removeFutureJdbc(connectionSession, requestId); return resultList.stream().map(jdbcGeneralResult -> { - Map cxt = context == null ? new HashMap<>() : context; - SqlExecuteResult result = generateResult(connectionSession, jdbcGeneralResult, cxt); + Map ctx = context == null ? new HashMap<>() : context; + SqlExecuteResult result = generateResult(connectionSession, jdbcGeneralResult, ctx); try (TraceStage stage = result.getSqlTuple().getSqlWatch().start(SqlExecuteStages.SQL_AFTER_CHECK)) { - sqlInterceptService.afterCompletion(result, connectionSession, cxt); + sqlInterceptService.afterCompletion(result, connectionSession, new AsyncExecuteContext(null, ctx)); } catch (Exception e) { throw new IllegalStateException(e); } @@ -440,24 +440,26 @@ public AsyncExecuteResultResp getAsyncResultV2(@NotNull String sessionId, String ConnectionSession connectionSession = sessionService.nullSafeGet(sessionId); AsyncExecuteContext context = (AsyncExecuteContext) ConnectionSessionUtil.getExecuteContext(connectionSession, requestId); - List resultList = context.getResults(); - List results = resultList.stream().map(jdbcGeneralResult -> { - SqlExecuteResult result = generateResult(connectionSession, jdbcGeneralResult, context.getContextMap()); - try (TraceStage stage = result.getSqlTuple().getSqlWatch().start(SqlExecuteStages.SQL_AFTER_CHECK)) { - sqlInterceptService.afterCompletion(result, connectionSession, context.getContextMap()); - } catch (Exception e) { - throw new IllegalStateException(e); - } - return result; - }).collect(Collectors.toList()); - if (context.isFinished()) { - ConnectionSessionUtil.removeExecuteContext(connectionSession, requestId); - return new AsyncExecuteResultResp(true, context, results); - } else { - if (log.isDebugEnabled()) { - log.debug("Get sql execution result timed out, sessionId={}, requestId={}", sessionId, requestId); + boolean shouldRemoveContext = context.isFinished(); + try { + List resultList = context.getFinishedSqlExecutionResults(); + List results = resultList.stream().map(jdbcGeneralResult -> { + SqlExecuteResult result = generateResult(connectionSession, jdbcGeneralResult, context.getContextMap()); + try (TraceStage stage = result.getSqlTuple().getSqlWatch().start(SqlExecuteStages.SQL_AFTER_CHECK)) { + sqlInterceptService.afterCompletion(result, connectionSession, context); + } catch (Exception e) { + throw new IllegalStateException(e); + } + return result; + }).collect(Collectors.toList()); + return new AsyncExecuteResultResp(context, results); + } catch (Exception e) { + shouldRemoveContext = true; + throw e; + } finally { + if (shouldRemoveContext) { + ConnectionSessionUtil.removeExecuteContext(connectionSession, requestId); } - return new AsyncExecuteResultResp(false, context, results); } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OBExecutionListener.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OBExecutionListener.java index ca1c835b16..9c2792ad1f 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OBExecutionListener.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OBExecutionListener.java @@ -56,7 +56,7 @@ public void onExecutionStart(SqlTuple sqlTuple, AsyncExecuteContext context) {} public void onExecutionEnd(SqlTuple sqlTuple, List results, AsyncExecuteContext context) {} @Override - public void onExecutionCanceled(SqlTuple sqlTuple, List results, AsyncExecuteContext context) {} + public void onExecutionCancelled(SqlTuple sqlTuple, List results, AsyncExecuteContext context) {} public void onExecutionStartAfter(SqlTuple sqlTuple, AsyncExecuteContext context) { if (CollectionUtils.isEmpty(sessionIds)) { diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java index ecec407fea..7c8d148088 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/OdcStatementCallBack.java @@ -178,11 +178,7 @@ public List doInStatement(Statement statement) throws SQLExce return this.sqls.stream().map(sqlTuple -> { JdbcGeneralResult result = JdbcGeneralResult.canceledResult(sqlTuple); result.setConnectionReset(true); - if (context != null) { - context.addResult(result); - } - listeners.forEach( - listener -> listener.onExecutionCanceled(sqlTuple, Collections.singletonList(result), context)); + onExecutionCancelled(sqlTuple, Collections.singletonList(result)); return result; }).collect(Collectors.toList()); } @@ -203,12 +199,7 @@ public List doInStatement(Statement statement) throws SQLExce // eat exception } } - if (context != null) { - context.setCurrentExecutingSql(sqlTuple.getExecutedSql()); - context.addCount(); - context.setCurrentExecutingSqlTraceId(null); - } - listeners.forEach(listener -> listener.onExecutionStart(sqlTuple, context)); + onExecutionStart(sqlTuple); try { applyConnectionSettings(statement); } catch (Exception e) { @@ -219,37 +210,16 @@ public List doInStatement(Statement statement) throws SQLExce if (Thread.currentThread().isInterrupted() || ConnectionSessionUtil.isConsoleSessionKillQuery(connectionSession)) { executeResults = Collections.singletonList(JdbcGeneralResult.canceledResult(sqlTuple)); - listeners.forEach(listener -> listener.onExecutionCanceled(sqlTuple, executeResults, context)); + onExecutionCancelled(sqlTuple, executeResults); } else { CountDownLatch latch = new CountDownLatch(1); - handle = executor.submit(() -> { - long startTs = System.currentTimeMillis(); - List sortedListeners = listeners.stream() - .filter(listener -> listener.getOnExecutionStartAfterMillis() != null - && listener.getOnExecutionStartAfterMillis() > 0) - .sorted(Comparator - .comparingLong(SqlExecutionListener::getOnExecutionStartAfterMillis)) - .collect(Collectors.toList()); - for (SqlExecutionListener listener : sortedListeners) { - long waitTs = System.currentTimeMillis() - startTs; - Long expectedTs = listener.getOnExecutionStartAfterMillis(); - if (!latch.await(expectedTs - waitTs, TimeUnit.MILLISECONDS)) { - listener.onExecutionStartAfter(sqlTuple, context); - } else { - break; - } - } - return null; - }); + handle = executor.submit(() -> onExecutionStartAfterMillis(sqlTuple, latch)); executeResults = doExecuteSql(statement, sqlTuple, latch); - listeners.forEach(listener -> listener.onExecutionEnd(sqlTuple, executeResults, context)); + onExecutionEnd(sqlTuple, executeResults); } } else { executeResults = Collections.singletonList(JdbcGeneralResult.canceledResult(sqlTuple)); - listeners.forEach(listener -> listener.onExecutionCanceled(sqlTuple, executeResults, context)); - } - if (context != null) { - context.addResults(executeResults); + onExecutionCancelled(sqlTuple, executeResults); } returnVal.addAll(executeResults); } @@ -274,6 +244,7 @@ public List doInStatement(Statement statement) throws SQLExce log.info("Clear dbms_output cache, dbmsInfo={}", dbmsInfo); } } + executor.shutdownNow(); } return returnVal; } @@ -596,6 +567,73 @@ private void rollback(Connection connection) { } } + private void onExecutionStart(SqlTuple sqlTuple) { + if (context != null) { + context.setCurrentExecutingSql(sqlTuple.getExecutedSql()); + context.incrementTotalExecutedSqlCount(); + context.setCurrentExecutingSqlTraceId(null); + } + listeners.forEach(listener -> { + try { + listener.onExecutionStart(sqlTuple, context); + } catch (Exception e) { + log.warn("An error occurred in listener {}.", listener.getClass(), e); + } + }); + } + + private void onExecutionCancelled(SqlTuple sqlTuple, List results) { + if (context != null) { + context.addSqlExecutionResults(results); + } + listeners.forEach(listener -> { + try { + listener.onExecutionCancelled(sqlTuple, results, context); + } catch (Exception e) { + log.warn("An error occurred in listener {}.", listener.getClass(), e); + } + }); + } + + private void onExecutionEnd(SqlTuple sqlTuple, List results) { + if (context != null) { + context.addSqlExecutionResults(results); + } + listeners.forEach(listener -> { + try { + listener.onExecutionEnd(sqlTuple, results, context); + } catch (Exception e) { + log.warn("An error occurred in listener {}.", listener.getClass(), e); + } + }); + } + + private Void onExecutionStartAfterMillis(SqlTuple sqlTuple, CountDownLatch latch) { + long startTs = System.currentTimeMillis(); + List sortedListeners = listeners.stream() + .filter(listener -> listener.getOnExecutionStartAfterMillis() != null + && listener.getOnExecutionStartAfterMillis() > 0) + .sorted(Comparator + .comparingLong(SqlExecutionListener::getOnExecutionStartAfterMillis)) + .collect(Collectors.toList()); + for (SqlExecutionListener listener : sortedListeners) { + long waitTs = System.currentTimeMillis() - startTs; + Long expectedTs = listener.getOnExecutionStartAfterMillis(); + try { + if (!latch.await(expectedTs - waitTs, TimeUnit.MILLISECONDS)) { + listener.onExecutionStartAfter(sqlTuple, context); + } else { + break; + } + } catch (InterruptedException e) { + return null; + } catch (Exception e) { + log.warn("An error occurred in listener {}.", listener.getClass(), e); + } + } + return null; + } + @Getter @ToString static class FunctionDefinition { diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/SqlExecutionListener.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/SqlExecutionListener.java index 1c07102214..1143412a21 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/SqlExecutionListener.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/SqlExecutionListener.java @@ -31,7 +31,7 @@ public interface SqlExecutionListener { void onExecutionEnd(SqlTuple sqlTuple, List results, AsyncExecuteContext context); - void onExecutionCanceled(SqlTuple sqlTuple, List results, AsyncExecuteContext context); + void onExecutionCancelled(SqlTuple sqlTuple, List results, AsyncExecuteContext context); void onExecutionStartAfter(SqlTuple sqlTuple, AsyncExecuteContext context); diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/BaseTimeConsumingInterceptor.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/BaseTimeConsumingInterceptor.java index 5ba5bd2bf2..77475f26a4 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/BaseTimeConsumingInterceptor.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/BaseTimeConsumingInterceptor.java @@ -17,11 +17,11 @@ package com.oceanbase.odc.service.session.interceptor; import java.util.List; -import java.util.Map; import java.util.stream.Collectors; import com.oceanbase.odc.common.util.TraceStage; import com.oceanbase.odc.core.session.ConnectionSession; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteReq; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteResp; import com.oceanbase.odc.service.session.model.SqlExecuteResult; @@ -32,7 +32,7 @@ public abstract class BaseTimeConsumingInterceptor implements SqlExecuteIntercep @Override public boolean preHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyncExecuteResp response, - @NonNull ConnectionSession session, @NonNull Map context) throws Exception { + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) throws Exception { List stageList = response.getSqls().stream() .map(v -> v.getSqlTuple().getSqlWatch().start(getExecuteStageName())) .collect(Collectors.toList()); @@ -51,19 +51,19 @@ public boolean preHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyncE @Override public void afterCompletion(@NonNull SqlExecuteResult response, - @NonNull ConnectionSession session, @NonNull Map context) throws Exception { + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) throws Exception { try (TraceStage stage = response.getSqlTuple().getSqlWatch().start(getExecuteStageName())) { doAfterCompletion(response, session, context); } } protected boolean doPreHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyncExecuteResp response, - @NonNull ConnectionSession session, @NonNull Map context) throws Exception { + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) throws Exception { return true; } protected void doAfterCompletion(@NonNull SqlExecuteResult response, - @NonNull ConnectionSession session, @NonNull Map context) throws Exception {} + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) throws Exception {} protected abstract String getExecuteStageName(); diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/DatabasePermissionInterceptor.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/DatabasePermissionInterceptor.java index 2a3fa79b3f..ba1696c133 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/DatabasePermissionInterceptor.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/DatabasePermissionInterceptor.java @@ -37,6 +37,7 @@ import com.oceanbase.odc.service.iam.auth.AuthenticationFacade; import com.oceanbase.odc.service.permission.database.model.DatabasePermissionType; import com.oceanbase.odc.service.permission.database.model.UnauthorizedDatabase; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteReq; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteResp; import com.oceanbase.odc.service.session.model.SqlExecuteResult; @@ -69,7 +70,7 @@ public int getOrder() { @Override public boolean doPreHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyncExecuteResp response, - @NonNull ConnectionSession session, @NonNull Map context) throws Exception { + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) throws Exception { if (authenticationFacade.currentUser().getOrganizationType() == OrganizationType.INDIVIDUAL) { return true; } @@ -100,7 +101,7 @@ public boolean doPreHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyn @Override public void afterCompletion(@NonNull SqlExecuteResult response, @NonNull ConnectionSession session, - @NonNull Map context) {} + @NonNull AsyncExecuteContext context) {} @Override protected String getExecuteStageName() { diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/NlsFormatInterceptor.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/NlsFormatInterceptor.java index 1a2e79d350..53e74ef523 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/NlsFormatInterceptor.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/NlsFormatInterceptor.java @@ -17,7 +17,6 @@ import java.util.Collections; import java.util.List; -import java.util.Map; import java.util.Objects; import java.util.stream.Collectors; @@ -36,6 +35,7 @@ import com.oceanbase.odc.core.sql.parser.AbstractSyntaxTreeFactories; import com.oceanbase.odc.core.sql.split.OffsetString; import com.oceanbase.odc.core.sql.split.SqlCommentProcessor; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteReq; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteResp; import com.oceanbase.odc.service.session.model.SqlExecuteResult; @@ -67,13 +67,13 @@ public class NlsFormatInterceptor extends BaseTimeConsumingInterceptor { @Override public boolean preHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyncExecuteResp response, - @NonNull ConnectionSession session, @NonNull Map context) { + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) { return true; } @Override public void doAfterCompletion(@NonNull SqlExecuteResult response, - @NonNull ConnectionSession session, @NonNull Map context) { + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) { DialectType dialect = session.getDialectType(); if (response.getStatus() != SqlExecuteStatus.SUCCESS || dialect != DialectType.OB_ORACLE) { return; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlCheckInterceptor.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlCheckInterceptor.java index e0d1403db4..19b8585156 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlCheckInterceptor.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlCheckInterceptor.java @@ -37,6 +37,7 @@ import com.oceanbase.odc.service.iam.auth.AuthenticationFacade; import com.oceanbase.odc.service.regulation.ruleset.RuleService; import com.oceanbase.odc.service.regulation.ruleset.model.Rule; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteReq; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteResp; import com.oceanbase.odc.service.session.model.SqlExecuteResult; @@ -74,11 +75,12 @@ public class SqlCheckInterceptor extends BaseTimeConsumingInterceptor { @Override public boolean doPreHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyncExecuteResp response, - @NonNull ConnectionSession session, @NonNull Map context) { - boolean sqlCheckIntercepted = handle(request, response, session, context); - context.put(SQL_CHECK_INTERCEPTED, sqlCheckIntercepted); - if (Objects.nonNull(context.get(SqlConsoleInterceptor.SQL_CONSOLE_INTERCEPTED))) { - return sqlCheckIntercepted && (Boolean) context.get(SqlConsoleInterceptor.SQL_CONSOLE_INTERCEPTED); + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) { + Map ctx = context.getContextMap(); + boolean sqlCheckIntercepted = handle(request, response, session, ctx); + ctx.put(SQL_CHECK_INTERCEPTED, sqlCheckIntercepted); + if (Objects.nonNull(ctx.get(SqlConsoleInterceptor.SQL_CONSOLE_INTERCEPTED))) { + return sqlCheckIntercepted && (Boolean) ctx.get(SqlConsoleInterceptor.SQL_CONSOLE_INTERCEPTED); } else { return true; } @@ -132,11 +134,12 @@ protected String getExecuteStageName() { @Override @SuppressWarnings("all") public void afterCompletion(@NonNull SqlExecuteResult response, @NonNull ConnectionSession session, - @NonNull Map context) throws Exception { - if (!context.containsKey(SQL_CHECK_RESULT_KEY)) { + @NonNull AsyncExecuteContext context) throws Exception { + Map ctx = context.getContextMap(); + if (!ctx.containsKey(SQL_CHECK_RESULT_KEY)) { return; } - Map> map = (Map>) context.get(SQL_CHECK_RESULT_KEY); + Map> map = (Map>) ctx.get(SQL_CHECK_RESULT_KEY); List results = map.get(response.getSqlTuple().getOffset()); if (CollectionUtils.isEmpty(results)) { return; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlConsoleInterceptor.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlConsoleInterceptor.java index 7610a3f1ff..35de0eff18 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlConsoleInterceptor.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlConsoleInterceptor.java @@ -42,6 +42,7 @@ import com.oceanbase.odc.service.regulation.ruleset.model.Rule; import com.oceanbase.odc.service.regulation.ruleset.model.Rule.RuleViolation; import com.oceanbase.odc.service.regulation.ruleset.model.SqlConsoleRules; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteReq; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteResp; import com.oceanbase.odc.service.session.model.SqlExecuteResult; @@ -74,11 +75,12 @@ public class SqlConsoleInterceptor extends BaseTimeConsumingInterceptor { @Override public boolean doPreHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyncExecuteResp response, - @NonNull ConnectionSession session, @NonNull Map context) { - boolean sqlConsoleIntercepted = handle(request, response, session, context); - context.put(SQL_CONSOLE_INTERCEPTED, sqlConsoleIntercepted); - if (Objects.nonNull(context.get(SqlCheckInterceptor.SQL_CHECK_INTERCEPTED))) { - return sqlConsoleIntercepted && (Boolean) context.get(SqlCheckInterceptor.SQL_CHECK_INTERCEPTED); + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) { + Map ctx = context.getContextMap(); + boolean sqlConsoleIntercepted = handle(request, response, session, ctx); + ctx.put(SQL_CONSOLE_INTERCEPTED, sqlConsoleIntercepted); + if (Objects.nonNull(ctx.get(SqlCheckInterceptor.SQL_CHECK_INTERCEPTED))) { + return sqlConsoleIntercepted && (Boolean) ctx.get(SqlCheckInterceptor.SQL_CHECK_INTERCEPTED); } else { return true; } @@ -199,7 +201,7 @@ protected String getExecuteStageName() { @Override public void doAfterCompletion(@NonNull SqlExecuteResult response, @NonNull ConnectionSession session, - @NonNull Map context) { + @NonNull AsyncExecuteContext context) { if (response.getStatus() != SqlExecuteStatus.SUCCESS) { return; } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlExecuteInterceptor.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlExecuteInterceptor.java index 5fb77ecaa9..225e350614 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlExecuteInterceptor.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlExecuteInterceptor.java @@ -15,11 +15,10 @@ */ package com.oceanbase.odc.service.session.interceptor; -import java.util.Map; - import org.springframework.core.Ordered; import com.oceanbase.odc.core.session.ConnectionSession; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteReq; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteResp; import com.oceanbase.odc.service.session.model.SqlExecuteResult; @@ -44,11 +43,11 @@ public interface SqlExecuteInterceptor extends Ordered { * @return whether to execute this sql */ default boolean preHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyncExecuteResp response, - @NonNull ConnectionSession session, @NonNull Map context) throws Exception { + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) throws Exception { return true; } default void afterCompletion(@NonNull SqlExecuteResult response, - @NonNull ConnectionSession session, @NonNull Map context) throws Exception {} + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) throws Exception {} } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlExecuteInterceptorService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlExecuteInterceptorService.java index 76ad156cc9..662ed3bb5b 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlExecuteInterceptorService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SqlExecuteInterceptorService.java @@ -29,6 +29,7 @@ import com.oceanbase.odc.core.authority.util.SkipAuthorize; import com.oceanbase.odc.core.session.ConnectionSession; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteReq; import com.oceanbase.odc.service.session.model.SqlAsyncExecuteResp; import com.oceanbase.odc.service.session.model.SqlExecuteResult; @@ -61,7 +62,7 @@ public void init() { } public boolean preHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyncExecuteResp response, - @NonNull ConnectionSession session, @NonNull Map context) throws Exception { + @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) throws Exception { for (SqlExecuteInterceptor interceptor : interceptors) { if (interceptor.preHandle(request, response, session, context)) { continue; @@ -72,7 +73,7 @@ public boolean preHandle(@NonNull SqlAsyncExecuteReq request, @NonNull SqlAsyncE } public void afterCompletion(@NonNull SqlExecuteResult response, @NonNull ConnectionSession session, - @NonNull Map context) throws Exception { + @NonNull AsyncExecuteContext context) throws Exception { for (SqlExecuteInterceptor interceptor : interceptors) { interceptor.afterCompletion(response, session, context); } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SwitchDatabaseInterceptor.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SwitchDatabaseInterceptor.java index 73a9b31fa7..9d41f99e8e 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SwitchDatabaseInterceptor.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/interceptor/SwitchDatabaseInterceptor.java @@ -16,7 +16,6 @@ package com.oceanbase.odc.service.session.interceptor; import java.util.Collections; -import java.util.Map; import java.util.Objects; import java.util.Optional; @@ -25,6 +24,7 @@ import com.oceanbase.odc.core.session.ConnectionSession; import com.oceanbase.odc.core.session.ConnectionSessionUtil; import com.oceanbase.odc.core.sql.execute.model.SqlExecuteStatus; +import com.oceanbase.odc.service.session.model.AsyncExecuteContext; import com.oceanbase.odc.service.session.model.SqlExecuteResult; import com.oceanbase.odc.service.session.util.SchemaExtractor; @@ -41,7 +41,7 @@ public class SwitchDatabaseInterceptor implements SqlExecuteInterceptor { @Override public void afterCompletion(@NonNull SqlExecuteResult response, @NonNull ConnectionSession session, - @NonNull Map context) throws Exception { + @NonNull AsyncExecuteContext context) throws Exception { if (response.getStatus() != SqlExecuteStatus.SUCCESS) { return; } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java index 8102935dad..69e5b9d770 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java @@ -44,7 +44,7 @@ public class AsyncExecuteContext { private Future> future; private String currentExecutingSqlTraceId; private String currentExecutingSql; - private int totalSqlExecutedCount = 0; + private int totalExecutedSqlCount = 0; public AsyncExecuteContext(List sqlTuples, Map contextMap) { this.sqlTuples = sqlTuples; @@ -52,18 +52,25 @@ public AsyncExecuteContext(List sqlTuples, Map context } public boolean isFinished() { - return future.isDone(); + return future != null && future.isDone(); } - public void addCount() { - totalSqlExecutedCount++; + public boolean isCancelled() { + return future != null && future.isCancelled(); } - public int getTotal() { + public void incrementTotalExecutedSqlCount() { + totalExecutedSqlCount++; + } + + public int getToBeExecutedSqlCount() { return sqlTuples.size(); } - public List getResults() { + /** + * only return the incremental results + */ + public List getFinishedSqlExecutionResults() { List copiedResults = new ArrayList<>(); while (!results.isEmpty()) { copiedResults.add(results.poll()); @@ -71,11 +78,7 @@ public List getResults() { return copiedResults; } - public void addResult(JdbcGeneralResult result) { - this.results.add(result); - } - - public void addResults(List results) { + public void addSqlExecutionResults(List results) { this.results.addAll(results); } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java index eb6edac6e9..0eb747b342 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java @@ -32,12 +32,12 @@ public class AsyncExecuteResultResp { private boolean finished; private String sql; - public AsyncExecuteResultResp(boolean finished, AsyncExecuteContext context, List results) { - this.finished = finished; + public AsyncExecuteResultResp(AsyncExecuteContext context, List results) { + this.finished = context.isFinished(); this.results = results; traceId = context.getCurrentExecutingSqlTraceId(); - total = context.getTotal(); - count = context.getTotalSqlExecutedCount(); + total = context.getToBeExecutedSqlCount(); + count = context.getTotalExecutedSqlCount(); sql = context.getCurrentExecutingSql(); } } From 1ab297d369d4cc323399980900032855fe01afc1 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Thu, 25 Apr 2024 16:18:26 +0800 Subject: [PATCH 14/16] response to CR comments --- .../web/controller/v2/ConnectSessionController.java | 12 ++++++------ .../odc/service/session/ConnectConsoleService.java | 6 +++--- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/ConnectSessionController.java b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/ConnectSessionController.java index 265d7b44b1..e633fb6b00 100644 --- a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/ConnectSessionController.java +++ b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/ConnectSessionController.java @@ -121,11 +121,11 @@ public SuccessResponse asyncSqlExecute(@PathVariable String return Responses.success(consoleService.execute(SidUtils.getSessionId(sessionId), req)); } - @RequestMapping(value = {"/sessions/{sessionId}/sqls/asyncExecuteV2"}, method = RequestMethod.POST) + @RequestMapping(value = {"/sessions/{sessionId}/sqls/streamExecute"}, method = RequestMethod.POST) @StatefulRoute(stateName = StateName.DB_SESSION, stateIdExpression = "#sessionId") - public SuccessResponse asyncSqlExecuteV2(@PathVariable String sessionId, + public SuccessResponse streamExecute(@PathVariable String sessionId, @RequestBody SqlAsyncExecuteReq req) throws Exception { - return Responses.success(consoleService.executeV2(SidUtils.getSessionId(sessionId), req, true)); + return Responses.success(consoleService.streamExecute(SidUtils.getSessionId(sessionId), req, true)); } /** @@ -142,11 +142,11 @@ public SuccessResponse> getAsyncSqlExecute(@PathVariable return Responses.success(consoleService.getAsyncResult(SidUtils.getSessionId(sessionId), requestId, null)); } - @RequestMapping(value = "/sessions/{sessionId}/sqls/getResultV2", method = RequestMethod.GET) + @RequestMapping(value = "/sessions/{sessionId}/sqls/getMoreResults", method = RequestMethod.GET) @StatefulRoute(stateName = StateName.DB_SESSION, stateIdExpression = "#sessionId") - public SuccessResponse getAsyncSqlExecuteV2(@PathVariable String sessionId, + public SuccessResponse getMoreResults(@PathVariable String sessionId, @RequestParam String requestId) { - return Responses.success(consoleService.getAsyncResultV2(SidUtils.getSessionId(sessionId), requestId)); + return Responses.success(consoleService.getMoreResults(SidUtils.getSessionId(sessionId), requestId)); } /** diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java index 0140483656..99d8243b6c 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java @@ -313,7 +313,7 @@ public SqlAsyncExecuteResp execute(@NotNull String sessionId, return response; } - public SqlAsyncExecuteResp executeV2(@NotNull String sessionId, + public SqlAsyncExecuteResp streamExecute(@NotNull String sessionId, @NotNull @Valid SqlAsyncExecuteReq request, boolean needSqlRuleCheck) throws Exception { ConnectionSession connectionSession = sessionService.nullSafeGet(sessionId, true); @@ -435,14 +435,14 @@ public List getAsyncResult(@NotNull String sessionId, String r } } - public AsyncExecuteResultResp getAsyncResultV2(@NotNull String sessionId, String requestId) { + public AsyncExecuteResultResp getMoreResults(@NotNull String sessionId, String requestId) { PreConditions.validArgumentState(Objects.nonNull(requestId), ErrorCodes.SqlRegulationRuleBlocked, null, null); ConnectionSession connectionSession = sessionService.nullSafeGet(sessionId); AsyncExecuteContext context = (AsyncExecuteContext) ConnectionSessionUtil.getExecuteContext(connectionSession, requestId); boolean shouldRemoveContext = context.isFinished(); try { - List resultList = context.getFinishedSqlExecutionResults(); + List resultList = context.getMoreSqlExecutionResults(); List results = resultList.stream().map(jdbcGeneralResult -> { SqlExecuteResult result = generateResult(connectionSession, jdbcGeneralResult, context.getContextMap()); try (TraceStage stage = result.getSqlTuple().getSqlWatch().start(SqlExecuteStages.SQL_AFTER_CHECK)) { From 4eb6f8dadc8adfee3fc53d78520f04cc608b5c46 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Thu, 25 Apr 2024 16:55:21 +0800 Subject: [PATCH 15/16] fix wrong method name --- .../odc/service/session/model/AsyncExecuteContext.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java index 69e5b9d770..11027990fe 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteContext.java @@ -70,7 +70,7 @@ public int getToBeExecutedSqlCount() { /** * only return the incremental results */ - public List getFinishedSqlExecutionResults() { + public List getMoreSqlExecutionResults() { List copiedResults = new ArrayList<>(); while (!results.isEmpty()) { copiedResults.add(results.poll()); From 6e1b48cf5ffbcbef4eff37cb56b21c56bc81af52 Mon Sep 17 00:00:00 2001 From: LuckyPickleZZ Date: Fri, 26 Apr 2024 13:51:29 +0800 Subject: [PATCH 16/16] response to CR comments --- .../oceanbase/odc/service/session/ConnectConsoleService.java | 2 +- .../odc/service/session/model/AsyncExecuteResultResp.java | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java index 99d8243b6c..e55505d00b 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/ConnectConsoleService.java @@ -452,7 +452,7 @@ public AsyncExecuteResultResp getMoreResults(@NotNull String sessionId, String r } return result; }).collect(Collectors.toList()); - return new AsyncExecuteResultResp(context, results); + return new AsyncExecuteResultResp(shouldRemoveContext, context, results); } catch (Exception e) { shouldRemoveContext = true; throw e; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java index 0eb747b342..ae94614894 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/session/model/AsyncExecuteResultResp.java @@ -32,8 +32,8 @@ public class AsyncExecuteResultResp { private boolean finished; private String sql; - public AsyncExecuteResultResp(AsyncExecuteContext context, List results) { - this.finished = context.isFinished(); + public AsyncExecuteResultResp(boolean finished, AsyncExecuteContext context, List results) { + this.finished = finished; this.results = results; traceId = context.getCurrentExecutingSqlTraceId(); total = context.getToBeExecutedSqlCount();