refactor(sql-execute): refactor SQL async execute api into streaming return

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<>());
+ }
+
 }

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 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

server/odc-core/src/main/java/com/oceanbase/odc/core/session/ConnectionSessionUtil.java

@@ -200,6 +200,52 @@ public static <T> void removeFutureJdbc(@NonNull ConnectionSession connectionSes
  id2FutureResult.remove(requestId);
  }
 
+ @SuppressWarnings("all")
+ public static <T> String setExecuteContext(@NonNull ConnectionSession connectionSession,
+ @NonNull Object context) {
+ Object value = connectionSession.getAttribute(ConnectionSessionConstants.ASYNC_EXECUTE_CONTEXT_KEY);
+ Map<String, Object> id2ExecuteContext;
+ if (value instanceof Map) {
+ id2ExecuteContext = (Map<String, Object>) 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 <T> Object getExecuteContext(@NonNull ConnectionSession connectionSession,
+ @NonNull String requestId) {
+ Object value = connectionSession.getAttribute(ConnectionSessionConstants.ASYNC_EXECUTE_CONTEXT_KEY);
+ Map<String, Object> id2ExecuteContext;
+ if (value instanceof Map) {
+ id2ExecuteContext = (Map<String, Object>) 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 <T> void removeExecuteContext(@NonNull ConnectionSession connectionSession,
+ @NonNull String requestId) {
+ Object value = connectionSession.getAttribute(ConnectionSessionConstants.ASYNC_EXECUTE_CONTEXT_KEY);
+ Map<String, Object> id2ExecuteContext;
+ if (!(value instanceof Map)) {
+ throw new NullPointerException("Result not found by session id " + connectionSession.getId());
+ } else {
+ id2ExecuteContext = (Map<String, Object>) value;
+ }
+ id2ExecuteContext.remove(requestId);
+ }
+
  public static void setUserId(@NonNull ConnectionSession connectionSession, @NonNull Long userId) {
  connectionSession.setAttribute(ConnectionSessionConstants.USER_ID_KEY, userId);
  }
@@ -535,6 +581,11 @@ public static String getConsoleConnectionId(@NonNull ConnectionSession connectio
  return (String) connectionSession.getAttribute(ConnectionSessionConstants.CONNECTION_ID_KEY);
  }
 
+ public static String getConsoleConnectionProxySessId(@NonNull ConnectionSession connectionSession) {
+ Object proxySessId = connectionSession.getAttribute(ConnectionSessionConstants.OB_PROXY_SESSID_KEY);
+ return proxySessId == null ? null : (String) proxySessId;
+ }
+
  public static String getVersion(@NonNull ConnectionSession connectionSession) {
  return (String) connectionSession.getAttribute(ConnectionSessionConstants.OB_VERSION);
  }

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,78 @@ 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<String> 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<String> 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<String> 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(".v$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);
+ }
+ }
+
 }

server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/ConnectSessionController.java

@@ -56,6 +56,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;
@@ -120,6 +121,13 @@ public SuccessResponse<SqlAsyncExecuteResp> asyncSqlExecute(@PathVariable String
  return Responses.success(consoleService.execute(SidUtils.getSessionId(sessionId), req));
  }
 
+ @RequestMapping(value = {"/sessions/{sessionId}/sqls/streamExecute"}, method = RequestMethod.POST)
+ @StatefulRoute(stateName = StateName.DB_SESSION, stateIdExpression = "#sessionId")
+ public SuccessResponse<SqlAsyncExecuteResp> streamExecute(@PathVariable String sessionId,
+ @RequestBody SqlAsyncExecuteReq req) throws Exception {
+ return Responses.success(consoleService.streamExecute(SidUtils.getSessionId(sessionId), req, true));
+ }
+
  /**
  * 获取异步执行sql的结果 Todo 这里的sqlIds后续需要改成一个string类型的requestId，异步api请求需要有超时机制
  *
@@ -134,6 +142,13 @@ public SuccessResponse<List<SqlExecuteResult>> getAsyncSqlExecute(@PathVariable
  return Responses.success(consoleService.getAsyncResult(SidUtils.getSessionId(sessionId), requestId, null));
  }
 
+ @RequestMapping(value = "/sessions/{sessionId}/sqls/getMoreResults", method = RequestMethod.GET)
+ @StatefulRoute(stateName = StateName.DB_SESSION, stateIdExpression = "#sessionId")
+ public SuccessResponse<AsyncExecuteResultResp> getMoreResults(@PathVariable String sessionId,
+ @RequestParam String requestId) {
+ return Responses.success(consoleService.getMoreResults(SidUtils.getSessionId(sessionId), requestId));
+ }
+
  /**
  * 对 sql 脚本的内容进行静态检查
  *

server/odc-service/src/main/java/com/oceanbase/odc/service/connection/util/ConnectionInfoUtil.java

@@ -23,13 +23,16 @@
 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;
 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;
 
@@ -61,6 +64,12 @@ 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 =
+ OBUtils.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);
  }

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<String, Object> context) {
+ @NonNull ConnectionSession session, @NonNull AsyncExecuteContext context) {
  return true;
  }
 
  @Override
  @SuppressWarnings("all")
  public void doAfterCompletion(@NonNull SqlExecuteResult response, @NonNull ConnectionSession session,
- @NonNull Map<String, Object> context) throws Exception {
+ @NonNull AsyncExecuteContext context) throws Exception {
  // TODO: May intercept sensitive column operation (WHERE / ORDER BY / HAVING)
  if (!maskingService.isMaskingEnabled()) {
  return;