diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/CheckPrivileges.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/CheckPrivileges.java index b65f7245ca7603..bce4862aeac94b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/CheckPrivileges.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/CheckPrivileges.java @@ -25,19 +25,23 @@ import org.apache.doris.nereids.exceptions.AnalysisException; import org.apache.doris.nereids.jobs.JobContext; import org.apache.doris.nereids.rules.analysis.UserAuthentication; +import org.apache.doris.nereids.trees.expressions.CTEId; import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.expressions.SlotReference; import org.apache.doris.nereids.trees.expressions.functions.table.TableValuedFunction; import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalCTEConsumer; import org.apache.doris.nereids.trees.plans.logical.LogicalCatalogRelation; import org.apache.doris.nereids.trees.plans.logical.LogicalRelation; import org.apache.doris.nereids.trees.plans.logical.LogicalTVFRelation; import org.apache.doris.nereids.trees.plans.logical.LogicalView; import org.apache.doris.qe.ConnectContext; +import com.google.common.collect.ImmutableList; import com.google.common.collect.Sets; import org.roaringbitmap.RoaringBitmap; +import java.util.HashSet; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -46,20 +50,28 @@ /** * CheckPrivileges - * This rule should only check once, because after check would set setPrivChecked in statementContext + * + * The privChecked flag is on StatementContext to support view permission passthrough: + * when InlineLogicalView expands a view, the outer CheckPrivileges has already verified + * the view-level permission, so the inner tables should not be re-checked. + * + * For CTEs, since LogicalCTEConsumer is a leaf node and the producer subtree is not + * traversed by ColumnPruning, we override visitLogicalCTEConsumer to explicitly + * traverse the producer plan and check privileges on its tables. */ public class CheckPrivileges extends ColumnPruning { private JobContext jobContext; + private final Set checkedCteIds = new HashSet<>(); @Override public Plan rewriteRoot(Plan plan, JobContext jobContext) { - // Only enter once, if repeated, the permissions of the table in the view will be checked - if (jobContext.getCascadesContext().getStatementContext().isPrivChecked()) { + StatementContext stmtCtx = jobContext.getCascadesContext().getStatementContext(); + if (stmtCtx.isPrivChecked()) { return plan; } this.jobContext = jobContext; super.rewriteRoot(plan, jobContext); - jobContext.getCascadesContext().getStatementContext().setPrivChecked(true); + stmtCtx.setPrivChecked(true); // don't rewrite plan, because Reorder expect no LogicalProject on LogicalJoin return plan; } @@ -88,6 +100,19 @@ public Plan visitLogicalRelation(LogicalRelation relation, PruneContext context) return super.visitLogicalRelation(relation, context); } + @Override + public Plan visitLogicalCTEConsumer(LogicalCTEConsumer consumer, PruneContext context) { + CTEId cteId = consumer.getCteId(); + StatementContext stmtCtx = jobContext.getCascadesContext().getStatementContext(); + Plan producerPlan = stmtCtx.getCteProducerByCteId(cteId); + if (producerPlan != null && checkedCteIds.add(cteId)) { + PruneContext producerContext = new PruneContext( + null, producerPlan.getOutputExprIdBitSet(), ImmutableList.of(), true); + producerPlan.accept(this, producerContext); + } + return super.visitLogicalCTEConsumer(consumer, context); + } + private Set computeUsedColumns(Plan plan, RoaringBitmap requiredSlotIds) { List outputs = plan.getOutput(); Map idToSlot = new LinkedHashMap<>(outputs.size()); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/privileges/TestCheckPrivileges.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/privileges/TestCheckPrivileges.java index 3de42904b25485..e8c1733b1d6af7 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/privileges/TestCheckPrivileges.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/privileges/TestCheckPrivileges.java @@ -165,6 +165,33 @@ public void testPrivilegesAndPolicies() throws Exception { ); } + // test CTE with JOIN privilege checking + // Verifies that column-level privileges are enforced when CTE is + // referenced multiple times via JOIN (CTE won't be inlined due to + // inlineCTEReferencedThreshold). CheckPrivileges.visitLogicalCTEConsumer + // explicitly traverses the CTE producer plan to check privileges. + { + // CTE + JOIN on fully-privileged table should succeed + query("WITH cte AS (SELECT id, name FROM custom_catalog.test_db.test_tbl1) " + + "SELECT a.id FROM cte a LEFT JOIN cte b ON a.id = b.id"); + + // CTE + JOIN accessing restricted column should be denied + Assertions.assertThrows(AnalysisException.class, () -> + query("WITH cte AS (SELECT * FROM custom_catalog.test_db.test_tbl2) " + + "SELECT a.id FROM cte a LEFT JOIN cte b ON a.id = b.id") + ); + + // CTE + JOIN accessing only allowed columns should succeed + query("WITH cte AS (SELECT id FROM custom_catalog.test_db.test_tbl2) " + + "SELECT a.id FROM cte a LEFT JOIN cte b ON a.id = b.id"); + + // CTE + INNER JOIN accessing restricted column should also be denied + Assertions.assertThrows(AnalysisException.class, () -> + query("WITH cte AS (SELECT * FROM custom_catalog.test_db.test_tbl2) " + + "SELECT a.id FROM cte a INNER JOIN cte b ON a.id = b.id") + ); + } + // test row policy with data masking { Function checkId = (NamedExpression ne) -> {