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 af24f9a83c7dfe..a11e247960d303 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 @@ -264,6 +264,8 @@ public enum TableFrom { private boolean privChecked; + private final Map cteProducerPlansForPrivCheck = Maps.newHashMap(); + // if greater than 0 means the duration has used private long materializedViewRewriteDuration = 0L; @@ -1020,6 +1022,14 @@ public void setPrivChecked(boolean privChecked) { this.privChecked = privChecked; } + public void addCteProducerPlanForPrivCheck(CTEId cteId, Plan producerPlan) { + cteProducerPlansForPrivCheck.put(cteId, producerPlan); + } + + public Plan getCteProducerPlanForPrivCheck(CTEId cteId) { + return cteProducerPlansForPrivCheck.get(cteId); + } + public List getUsedBackendsDistributing() { return usedBackendsDistributing; } 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..07bce890d88c3f 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.getCteProducerPlanForPrivCheck(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/main/java/org/apache/doris/nereids/rules/rewrite/RewriteCteChildren.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/RewriteCteChildren.java index 0360c4aa0e7ab2..f996c8fee3ccfc 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/RewriteCteChildren.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/RewriteCteChildren.java @@ -80,6 +80,8 @@ public Plan visit(Plan plan, CascadesContext context) { @Override public Plan visitLogicalCTEAnchor(LogicalCTEAnchor cteAnchor, CascadesContext cascadesContext) { + cascadesContext.getStatementContext().addCteProducerPlanForPrivCheck( + cteAnchor.getCteId(), cteAnchor.child(0)); LogicalPlan outer; if (cascadesContext.getStatementContext().getRewrittenCteConsumer().containsKey(cteAnchor.getCteId())) { outer = cascadesContext.getStatementContext().getRewrittenCteConsumer().get(cteAnchor.getCteId()); 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) -> {