diff --git a/fe/fe-core/src/main/java/org/apache/doris/common/cache/NereidsSqlCacheManager.java b/fe/fe-core/src/main/java/org/apache/doris/common/cache/NereidsSqlCacheManager.java index 73d6b7ebb25ac3..66f72c4192c4e6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/common/cache/NereidsSqlCacheManager.java +++ b/fe/fe-core/src/main/java/org/apache/doris/common/cache/NereidsSqlCacheManager.java @@ -147,6 +147,7 @@ public void tryAddFeSqlCache(ConnectContext connectContext, String sql) { } SqlCacheContext sqlCacheContext = sqlCacheContextOpt.get(); + sqlCacheContext.setQueryId(connectContext.queryId()); String key = sqlCacheContext.getCacheKeyType() == CacheKeyType.SQL ? generateCacheKey(connectContext, normalizeSql(sql)) : generateCacheKey(connectContext, DebugUtil.printId(sqlCacheContext.getOrComputeCacheKeyMd5())); @@ -174,6 +175,7 @@ public void tryAddBeCache(ConnectContext connectContext, String sql, CacheAnalyz return; } SqlCacheContext sqlCacheContext = sqlCacheContextOpt.get(); + sqlCacheContext.setQueryId(connectContext.queryId()); String key = sqlCacheContext.getCacheKeyType() == CacheKeyType.SQL ? generateCacheKey(connectContext, normalizeSql(sql)) : generateCacheKey(connectContext, DebugUtil.printId(sqlCacheContext.getOrComputeCacheKeyMd5())); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/SqlCacheContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/SqlCacheContext.java index 0794d0aca0b2cd..458cdae3b739f5 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/SqlCacheContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/SqlCacheContext.java @@ -60,7 +60,7 @@ public class SqlCacheContext { private static final Logger LOG = LogManager.getLogger(SqlCacheContext.class); private final UserIdentity userIdentity; - private final TUniqueId queryId; + private volatile TUniqueId queryId; // if contains udf/udaf/tableValuesFunction we can not process it and skip use sql cache private volatile boolean cannotProcessExpression; private volatile String originSql; @@ -99,9 +99,8 @@ public class SqlCacheContext { private volatile CacheKeyType cacheKeyType = CacheKeyType.SQL; - public SqlCacheContext(UserIdentity userIdentity, TUniqueId queryId) { + public SqlCacheContext(UserIdentity userIdentity) { this.userIdentity = Objects.requireNonNull(userIdentity, "userIdentity cannot be null"); - this.queryId = Objects.requireNonNull(queryId, "queryId cannot be null"); } public String getPhysicalPlan() { @@ -437,6 +436,10 @@ public void setCacheKeyType(CacheKeyType cacheKeyType) { this.cacheKeyType = cacheKeyType; } + public void setQueryId(TUniqueId queryId) { + this.queryId = queryId; + } + /** FullTableName */ @lombok.Data @lombok.AllArgsConstructor diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java index a378087252f304..41a9944edabd65 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/StatementContext.java @@ -258,10 +258,10 @@ public StatementContext(ConnectContext connectContext, OriginStatement originSta this.originStatement = originStatement; exprIdGenerator = ExprId.createGenerator(initialId); if (connectContext != null && connectContext.getSessionVariable() != null - && connectContext.queryId() != null && CacheAnalyzer.canUseSqlCache(connectContext.getSessionVariable())) { + // cannot set the queryId here because the queryId for the current query is set in the subsequent steps. this.sqlCacheContext = new SqlCacheContext( - connectContext.getCurrentUserIdentity(), connectContext.queryId()); + connectContext.getCurrentUserIdentity()); if (originStatement != null) { this.sqlCacheContext.setOriginSql(originStatement.originStmt); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java index b06a9594b6b83c..05de6977371267 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java @@ -939,6 +939,13 @@ public void setQueryId(TUniqueId queryId) { } } + public void resetQueryId() { + if (this.queryId != null) { + this.lastQueryId = this.queryId.deepCopy(); + } + this.queryId = null; + } + public void setTraceId(String traceId) { this.traceId = traceId; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectProcessor.java index 98214597252c50..8ad600f465e8be 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectProcessor.java @@ -232,6 +232,9 @@ protected void auditAfterExec(String origStmt, StatementBase parsedStmt, // only throw an exception when there is a problem interacting with the requesting client protected void handleQuery(String originStmt) throws ConnectionException { + // Before executing the query, the queryId should be set to empty. + // Otherwise, if SQL parsing fails, the audit log will record the queryId from the previous query. + ctx.resetQueryId(); if (Config.isCloudMode()) { if (!ctx.getCurrentUserIdentity().isRootUser() && ((CloudSystemInfoService) Env.getCurrentSystemInfo()).getInstanceStatus() diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java index f16bc58a4fcd48..573479a4ec5960 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectContextTest.java @@ -32,6 +32,7 @@ import java.nio.channels.SocketChannel; import java.util.List; import java.util.Optional; +import java.util.UUID; public class ConnectContextTest { @Mocked @@ -273,4 +274,27 @@ public void testInsertQueryTimeoutS() { result = context.getInsertTimeoutS(); Assert.assertEquals(propertyValue, result); } + + @Test + public void testResetQueryId() { + ConnectContext context = new ConnectContext(); + Assert.assertNull(context.queryId); + Assert.assertNull(context.lastQueryId); + + UUID uuid = UUID.randomUUID(); + TUniqueId queryId = new TUniqueId(uuid.getMostSignificantBits(), uuid.getLeastSignificantBits()); + context.setQueryId(queryId); + Assert.assertEquals(queryId, context.queryId); + Assert.assertNull(context.lastQueryId); + + context.resetQueryId(); + Assert.assertNull(context.queryId); + Assert.assertEquals(queryId, context.lastQueryId); + + UUID uuid2 = UUID.randomUUID(); + TUniqueId queryId2 = new TUniqueId(uuid2.getMostSignificantBits(), uuid2.getLeastSignificantBits()); + context.setQueryId(queryId2); + Assert.assertEquals(queryId2, context.queryId); + Assert.assertEquals(queryId, context.lastQueryId); + } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/SqlCacheTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/SqlCacheTest.java index afe95a49bde6c2..f67d037fdf64fb 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/SqlCacheTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/SqlCacheTest.java @@ -37,11 +37,11 @@ public void testCacheKey() { queryId.setLo(uuid.getLeastSignificantBits()); UserIdentity admin = new UserIdentity("admin", "127.0.0.1"); - SqlCacheContext cacheContext = new SqlCacheContext(admin, queryId); + SqlCacheContext cacheContext = new SqlCacheContext(admin); cacheContext.setOriginSql("SELECT * FROM tbl"); PUniqueId key1 = cacheContext.doComputeCacheKeyMd5(ImmutableSet.of()); - SqlCacheContext cacheContext2 = new SqlCacheContext(admin, queryId); + SqlCacheContext cacheContext2 = new SqlCacheContext(admin); cacheContext2.setOriginSql( "-- Same query with comments and extra spaces\n" + "/* Comment */ SELECT * FROM tbl " @@ -49,7 +49,7 @@ public void testCacheKey() { PUniqueId key2 = cacheContext2.doComputeCacheKeyMd5(ImmutableSet.of()); Assertions.assertEquals(key1, key2); - SqlCacheContext cacheContext3 = new SqlCacheContext(admin, queryId); + SqlCacheContext cacheContext3 = new SqlCacheContext(admin); cacheContext3.setOriginSql( "-- Same query with comments and extra spaces\n" + "/* Comment */ SELeCT * FROM tbl "