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

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);
     }
@@ -522,6 +568,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,63 @@ public static String queryDBMSOutput(@NonNull Connection connection, Integer max
         }
     }
 
+    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(".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

@@ -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<SqlAsyncExecuteResp> 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<SqlAsyncExecuteResp> 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<List<SqlExecuteResult>> 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<AsyncExecuteResultResp> getAsyncSqlExecuteV2(@PathVariable String sessionId,
+            @RequestParam String requestId) {
+        return Responses.success(consoleService.getAsyncResultV2(SidUtils.getSessionId(sessionId), requestId));
+    }
+
     /**
      * 对 sql 脚本的内容进行静态检查
      *

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 = 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);
         }
@@ -71,6 +79,21 @@ 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());