diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertOverwriteTableCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertOverwriteTableCommand.java index 5b11acf5de0e8e..eb058e6f60cf86 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertOverwriteTableCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertOverwriteTableCommand.java @@ -194,7 +194,6 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { ConnectContext.get().getQualifiedUser(), ConnectContext.get().getRemoteIP(), ((OlapTable) targetTable).getQualifiedDbName() + ": " + targetTable.getName()); } - ConnectContext.get().setSkipAuth(true); } partitionNames = ((UnboundTableSink) logicalQuery).getPartitions(); // If not specific partition to overwrite, means it's a command to overwrite the table. @@ -228,6 +227,12 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { isRunning.set(true); long taskId = 0; try { + // OLAP overwrite runs its internal partition replacement with the auth check skipped. + // Set the flag here, inside the try, so the finally below always pairs the reset even if + // an earlier step (e.g. the @branch guard) throws before we get here. + if (physicalTableSink instanceof PhysicalOlapTableSink && targetTable instanceof OlapTable) { + ctx.setSkipAuth(true); + } if (isAutoDetectOverwrite(getLogicalQuery())) { // taskId here is a group id. it contains all replace tasks made and registered in rpc process. taskId = insertOverwriteManager.registerTaskGroup(targetTable); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertOverwriteSkipAuthResetTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertOverwriteSkipAuthResetTest.java new file mode 100644 index 00000000000000..7aab2f1184c6ad --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertOverwriteSkipAuthResetTest.java @@ -0,0 +1,81 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 org.apache.doris.nereids.trees.plans.commands.insert; + +import org.apache.doris.nereids.StatementContext; +import org.apache.doris.nereids.parser.NereidsParser; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.qe.OriginStatement; +import org.apache.doris.qe.StmtExecutor; +import org.apache.doris.utframe.TestWithFeService; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** + * Regression test for the {@code skipAuth} flag lifecycle in + * {@link InsertOverwriteTableCommand#run}. + * + *

For an OLAP target the command runs its internal partition-replacement work with the + * auth check skipped, and must reset that flag before returning. The flag used to be flipped + * on before the surrounding try/finally, so any statement that threw in between (for + * example the {@code @branch}-on-non-iceberg guard) returned without the finally ever running, + * leaving the flag set for the rest of the connection. This test drives exactly that failing + * path and asserts the flag is back to its original value afterwards.

+ */ +public class InsertOverwriteSkipAuthResetTest extends TestWithFeService { + + @Override + protected void runBeforeAll() throws Exception { + createDatabaseAndUse("iot_skipauth"); + createTable("CREATE TABLE t (k INT) DISTRIBUTED BY HASH(k) BUCKETS 1 " + + "PROPERTIES ('replication_num' = '1')"); + } + + @Test + public void skipAuthIsResetAfterFailingBranchOverwrite() throws Exception { + // @branch is only valid for iceberg; against an OLAP table run() plans the sink, marks the + // overwrite as needing the auth-skip, then throws the guard below. The flag must not leak. + String sql = "INSERT OVERWRITE TABLE iot_skipauth.t@BRANCH(anything) SELECT * FROM iot_skipauth.t"; + + LogicalPlan parsed = new NereidsParser().parseSingle(sql); + Assertions.assertTrue(parsed instanceof InsertOverwriteTableCommand, + "an INSERT OVERWRITE ... @branch statement should parse to InsertOverwriteTableCommand"); + InsertOverwriteTableCommand command = (InsertOverwriteTableCommand) parsed; + + StatementContext statementContext = new StatementContext(connectContext, new OriginStatement(sql, 0)); + connectContext.setStatementContext(statementContext); + statementContext.setConnectContext(connectContext); + StmtExecutor executor = new StmtExecutor(connectContext, sql); + + // Baseline: the connection starts with the flag off. + connectContext.setSkipAuth(false); + + Exception thrown = Assertions.assertThrows(Exception.class, + () -> command.run(connectContext, executor)); + // Prove we actually reached the guard that fires AFTER the overwrite is marked as needing the + // auth-skip -- otherwise this test would pass without exercising the leak path at all. + Assertions.assertTrue(thrown.getMessage() != null + && thrown.getMessage().contains("Only support insert overwrite into iceberg table's branch"), + "expected the @branch-on-non-iceberg guard to fire, but got: " + thrown.getMessage()); + + // The flag must have been reset even though run() exited via an exception. + Assertions.assertFalse(connectContext.isSkipAuth(), + "skipAuth must be reset after a failed INSERT OVERWRITE, but it was left set"); + } +}