From 5f315b45f9ab5a8d5c35c717993a5b0858cbca80 Mon Sep 17 00:00:00 2001 From: tiancaiamao Date: Tue, 9 Jun 2026 22:54:26 +0800 Subject: [PATCH 01/13] feat: implement masking policy projection (AT RESULT semantics) Implements the core masking policy projection feature: - Logical plan: buildFinalProjectionWithMasking applies masking as a final projection on top of SELECT/UNION results, preserving original values for WHERE/JOIN/GROUP BY/HAVING/ORDER BY processing (AT RESULT semantics). - PointGet/BatchPointGet fast path: builds MaskingExprs directly in the plan and applies them in the executor after row decoding. - CTE handling: sets isCTE/buildingCTE flags to skip masking inside CTE definitions, preserving OrigTblName/OrigColName for downstream lookup. - Expression cache: caches parsed masking expressions in SessionVars keyed by schema version with recursion guard. - Set operator handling: uses buildingSetOprOperands depth counter to apply masking only at the outermost set-operator result. - Tests: TestMaskingPolicyExprCache, TestMaskingPolicyProjection, TestMaskingPolicyAtResult (HAVING/ORDER BY/UNION/Aggregate coverage). --- pkg/executor/batch_point_get.go | 11 + pkg/executor/builder.go | 1 + pkg/executor/point_get.go | 39 +++ pkg/planner/core/logical_plan_builder.go | 280 +++++++++++++++++- .../core/masking_policy_at_result_test.go | 94 ++++++ pkg/planner/core/masking_policy_expr_cache.go | 101 +++++++ .../core/masking_policy_expr_cache_test.go | 103 +++++++ .../core/masking_policy_projection_test.go | 100 +++++++ .../physicalop/physical_batch_point_get.go | 13 +- pkg/planner/core/planbuilder.go | 4 + pkg/planner/core/point_get_plan.go | 98 ++++++ .../variable/masking_policy_cache.go | 26 ++ pkg/sessionctx/variable/session.go | 5 + 13 files changed, 869 insertions(+), 6 deletions(-) create mode 100644 pkg/planner/core/masking_policy_at_result_test.go create mode 100644 pkg/planner/core/masking_policy_expr_cache.go create mode 100644 pkg/planner/core/masking_policy_expr_cache_test.go create mode 100644 pkg/planner/core/masking_policy_projection_test.go create mode 100644 pkg/sessionctx/variable/masking_policy_cache.go diff --git a/pkg/executor/batch_point_get.go b/pkg/executor/batch_point_get.go index 83109b77391a3..9c8ed4f4f70ce 100644 --- a/pkg/executor/batch_point_get.go +++ b/pkg/executor/batch_point_get.go @@ -24,6 +24,7 @@ import ( "github.com/pingcap/failpoint" "github.com/pingcap/tidb/pkg/executor/internal/exec" + "github.com/pingcap/tidb/pkg/expression" "github.com/pingcap/tidb/pkg/kv" "github.com/pingcap/tidb/pkg/meta/model" "github.com/pingcap/tidb/pkg/parser/ast" @@ -80,6 +81,9 @@ type BatchPointGetExec struct { // virtualColumnRetFieldTypes records the RetFieldTypes of virtual columns. virtualColumnRetFieldTypes []*types.FieldType + // MaskingExprs stores the masking expressions for columns that have masking policies. + MaskingExprs []expression.Expression + snapshot kv.Snapshot stats *runtimeStatsWithSnapshot } @@ -246,6 +250,13 @@ func (e *BatchPointGetExec) Next(ctx context.Context, req *chunk.Chunk) error { if err != nil { return err } + + if e.MaskingExprs != nil && req.NumRows() > 0 { + if err := applyMaskingExprs(sctx, schema, e.MaskingExprs, req); err != nil { + return err + } + } + return nil } diff --git a/pkg/executor/builder.go b/pkg/executor/builder.go index a1655fb5765ef..fbddb15e56f1e 100644 --- a/pkg/executor/builder.go +++ b/pkg/executor/builder.go @@ -5761,6 +5761,7 @@ func (b *executorBuilder) buildBatchPointGet(plan *physicalop.BatchPointGetPlan) handles: handles, idxVals: plan.IndexValues, partitionNames: plan.PartitionNames, + MaskingExprs: plan.MaskingExprs, } e.snapshot, err = b.getSnapshot() diff --git a/pkg/executor/point_get.go b/pkg/executor/point_get.go index 588c503591900..e49831e0134fe 100644 --- a/pkg/executor/point_get.go +++ b/pkg/executor/point_get.go @@ -95,6 +95,7 @@ func (b *executorBuilder) buildPointGet(p *physicalop.PointGetPlan) exec.Executo e.SetInitCap(1) e.SetMaxChunkSize(1) e.Init(p) + e.MaskingExprs = p.MaskingExprs snapshotTS, err := b.getSnapshotTS() if err != nil { @@ -145,6 +146,9 @@ type PointGetExecutor struct { // virtualColumnRetFieldTypes records the RetFieldTypes of virtual columns. virtualColumnRetFieldTypes []*types.FieldType + // MaskingExprs stores the masking expressions for columns that have masking policies. + MaskingExprs []expression.Expression + stats *runtimeStatsWithSnapshot } @@ -448,6 +452,41 @@ func (e *PointGetExecutor) Next(ctx context.Context, req *chunk.Chunk) error { if err != nil { return err } + + if e.MaskingExprs != nil && req.NumRows() > 0 { + if err := applyMaskingExprs(sctx, schema, e.MaskingExprs, req); err != nil { + return err + } + } + + return nil +} + +// applyMaskingExprs evaluates masking expressions for each row in the chunk and +// replaces the original column values with masked values. +func applyMaskingExprs(sctx sessionctx.Context, schema *expression.Schema, maskingExprs []expression.Expression, req *chunk.Chunk) error { + fieldTypes := make([]*types.FieldType, len(schema.Columns)) + for i, col := range schema.Columns { + fieldTypes[i] = col.RetType + } + maskedChunk := chunk.NewChunkWithCapacity(fieldTypes, req.NumRows()) + evalCtx := sctx.GetExprCtx().GetEvalCtx() + for rowIdx := 0; rowIdx < req.NumRows(); rowIdx++ { + row := req.GetRow(rowIdx) + for colIdx, expr := range maskingExprs { + if expr != nil { + result, err := expr.Eval(evalCtx, row) + if err != nil { + return err + } + maskedChunk.AppendDatum(colIdx, &result) + } else { + datum := row.GetDatum(colIdx, schema.Columns[colIdx].RetType) + maskedChunk.AppendDatum(colIdx, &datum) + } + } + } + req.SwapColumns(maskedChunk) return nil } diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go index 2ab42c23158f2..1276118b5f37a 100644 --- a/pkg/planner/core/logical_plan_builder.go +++ b/pkg/planner/core/logical_plan_builder.go @@ -1531,6 +1531,203 @@ func buildExpandFieldName(ctx expression.EvalContext, expr expression.Expression return newName } +// buildMaskingReplaceExprs builds a list of replacement expressions for masking. +// For each column in the schema, if the column has a masking policy, the replacement +// expression is the masking function applied to the column. Otherwise it's the column itself. +func (b *PlanBuilder) buildMaskingReplaceExprs(ctx context.Context, p base.LogicalPlan) ([]expression.Expression, error) { + if b.is == nil || p == nil { + return nil, nil + } + sv := b.ctx.GetSessionVars() + if sv != nil && sv.InRestrictedSQL { + return nil, nil + } + if len(b.is.AllMaskingPolicies()) == 0 { + return nil, nil + } + cols := p.Schema().Columns + names := p.OutputNames() + if len(cols) == 0 || len(cols) != len(names) { + return nil, nil + } + replaceExprs := make([]expression.Expression, len(cols)) + for i, col := range cols { + replaceExprs[i] = col + } + schemaVersion := b.is.SchemaMetaVersion() + hasMask := false + for i, col := range cols { + policy, tblInfo, colInfo := b.findMaskingPolicy(ctx, names[i], col) + if policy == nil { + continue + } + if policy.Status != model.MaskingPolicyStatusEnable { + continue + } + expr, placeholder, err := getMaskingPolicyExpr(b.ctx.GetExprCtx(), sv, schemaVersion, policy, tblInfo, colInfo) + if err != nil { + return nil, err + } + if placeholder == nil { + continue + } + masked := expression.ColumnSubstitute(b.ctx.GetExprCtx(), expr, expression.NewSchema(placeholder), []expression.Expression{col}) + replaceExprs[i] = masked + hasMask = true + } + if !hasMask { + return nil, nil + } + return replaceExprs, nil +} + +func (b *PlanBuilder) findMaskingPolicy(_ context.Context, name *types.FieldName, col *expression.Column) (*model.MaskingPolicyInfo, *model.TableInfo, *model.ColumnInfo) { + if b.is == nil || name == nil { + return nil, nil, nil + } + tblName := name.OrigTblName + colName := name.OrigColName + if tblName.L == "" { + tblName = name.TblName + } + if colName.L == "" { + colName = name.ColName + } + dbName := name.DBName + if dbName.L == "" { + dbName = ast.NewCIStr(b.ctx.GetSessionVars().CurrentDB) + } + tbl, err := b.is.TableByName(context.Background(), dbName, tblName) + if err != nil { + return nil, nil, nil + } + tblInfo := tbl.Meta() + if tblInfo == nil { + return nil, nil, nil + } + colInfo := tblInfo.FindPublicColumnByName(colName.L) + if colInfo == nil { + return nil, nil, nil + } + policy, ok := b.is.MaskingPolicyByTableColumn(tblInfo.ID, colInfo.ID) + if !ok || policy == nil { + return nil, nil, nil + } + return policy, tblInfo, colInfo +} + +// buildFinalProjectionWithMasking builds a final projection that applies masking policies +// to the result. This implements the "AT RESULT" semantics where masking is applied +// only after all relational operations (HAVING, ORDER BY, set operators) have been +// computed using original values. +// +// The originalFields parameter contains the original SELECT field expressions before +// they were materialized. This allows us to correctly apply masking to wrapper expressions +// like CONCAT(c, ”) or CAST(c AS CHAR) by processing the original expression tree. +func (b *PlanBuilder) buildFinalProjectionWithMasking(ctx context.Context, p base.LogicalPlan, oldLen int, originalFields []*ast.SelectField) (base.LogicalPlan, error) { + if b.is == nil || p == nil { + return p, nil + } + sv := b.ctx.GetSessionVars() + if sv != nil && sv.InRestrictedSQL { + return p, nil + } + if len(b.is.AllMaskingPolicies()) == 0 { + return p, nil + } + // If we don't have original fields, fall back to the simple column-based + // masking path (used by set operators). + if len(originalFields) == 0 || len(originalFields) < oldLen { + return b.buildFinalProjectionWithMaskingSimple(ctx, p, oldLen) + } + for i := 0; i < oldLen; i++ { + field := originalFields[i] + if field == nil || field.Auxiliary || field.WildCard != nil || field.Expr == nil { + return b.buildFinalProjectionWithMaskingSimple(ctx, p, oldLen) + } + } + + // Prefer rebuilding from the current projection expressions. + // This keeps wrapper semantics correct, e.g. CONCAT(c, '-') should become + // CONCAT(mask_expr(c), '-') instead of mask_expr(CONCAT(c, '-')). + proj, ok := p.(*logicalop.LogicalProjection) + if !ok || len(proj.Children()) == 0 || len(proj.Exprs) < oldLen { + return b.buildFinalProjectionWithMaskingSimple(ctx, p, oldLen) + } + child := proj.Children()[0] + childMaskExprs, err := b.buildMaskingReplaceExprs(ctx, child) + if err != nil { + return nil, err + } + if childMaskExprs == nil { + return p, nil + } + + finalProj := logicalop.LogicalProjection{Exprs: make([]expression.Expression, 0, oldLen)}.Init(b.ctx, b.getSelectOffset()) + schema := expression.NewSchema(make([]*expression.Column, 0, oldLen)...) + newNames := make([]*types.FieldName, 0, oldLen) + for i := 0; i < oldLen; i++ { + expr := expression.ColumnSubstitute(b.ctx.GetExprCtx(), proj.Exprs[i], child.Schema(), childMaskExprs) + finalProj.Exprs = append(finalProj.Exprs, expr) + + baseCol := p.Schema().Columns[i] + newCol := &expression.Column{ + UniqueID: b.ctx.GetSessionVars().AllocPlanColumnID(), + RetType: expr.GetType(b.ctx.GetExprCtx().GetEvalCtx()).Clone(), + } + newCol.ID = baseCol.ID + newCol.SetCoercibility(baseCol.Coercibility()) + newCol.SetRepertoire(baseCol.Repertoire()) + schema.Append(newCol) + newNames = append(newNames, p.OutputNames()[i]) + } + finalProj.SetSchema(schema) + finalProj.SetOutputNames(newNames) + finalProj.SetChildren(child) + return finalProj, nil +} + +// buildFinalProjectionWithMaskingSimple is a fallback that applies masking based on column references +// when originalFields are not available. This is used for cases like set operators where we don't +// have the original field expressions. +func (b *PlanBuilder) buildFinalProjectionWithMaskingSimple(ctx context.Context, p base.LogicalPlan, oldLen int) (base.LogicalPlan, error) { + maskExprs, err := b.buildMaskingReplaceExprs(ctx, p) + if err != nil { + return nil, err + } + if maskExprs == nil { + return p, nil + } + + actualLen := p.Schema().Len() + if oldLen < 0 || oldLen > actualLen { + oldLen = actualLen + } + + proj := logicalop.LogicalProjection{Exprs: make([]expression.Expression, 0, oldLen)}.Init(b.ctx, b.getSelectOffset()) + schema := expression.NewSchema(make([]*expression.Column, 0, oldLen)...) + newNames := make([]*types.FieldName, 0, oldLen) + for i := 0; i < oldLen; i++ { + col := p.Schema().Columns[i] + expr := expression.ColumnSubstitute(b.ctx.GetExprCtx(), col, p.Schema(), maskExprs) + proj.Exprs = append(proj.Exprs, expr) + + newCol := &expression.Column{ + UniqueID: b.ctx.GetSessionVars().AllocPlanColumnID(), + RetType: expr.GetType(b.ctx.GetExprCtx().GetEvalCtx()).Clone(), + } + newCol.ID = col.ID + newCol.SetCoercibility(col.Coercibility()) + newCol.SetRepertoire(col.Repertoire()) + schema.Append(newCol) + newNames = append(newNames, p.OutputNames()[i]) + } + proj.SetSchema(schema) + proj.SetOutputNames(newNames) + proj.SetChildren(p) + return proj, nil +} + // buildProjectionField builds the field object according to SelectField in projection. func (b *PlanBuilder) buildProjectionField(ctx context.Context, p base.LogicalPlan, field *ast.SelectField, expr expression.Expression) (*expression.Column, *types.FieldName, error) { var origTblName, tblName, origColName, colName, dbName ast.CIStr @@ -1765,7 +1962,7 @@ func (b *PlanBuilder) implicitProjectGroupingSetCols(projSchema *expression.Sche // buildProjection returns a Projection plan and non-aux columns length. func (b *PlanBuilder) buildProjection(ctx context.Context, p base.LogicalPlan, fields []*ast.SelectField, mapper map[*ast.AggregateFuncExpr]int, - windowMapper map[*ast.WindowFuncExpr]int, considerWindow bool, expandGenerateColumn bool) (base.LogicalPlan, []expression.Expression, int, error) { + windowMapper map[*ast.WindowFuncExpr]int, considerWindow bool, expandGenerateColumn bool, applyMasking bool) (base.LogicalPlan, []expression.Expression, int, error) { err := b.preprocessUserVarTypes(ctx, p, fields, mapper) if err != nil { return nil, nil, 0, err @@ -1776,6 +1973,20 @@ func (b *PlanBuilder) buildProjection(ctx context.Context, p base.LogicalPlan, f schema := expression.NewSchema(make([]*expression.Column, 0, len(fields))...) oldLen := 0 newNames := make([]*types.FieldName, 0, len(fields)) + var cachedMaskPlan base.LogicalPlan + var cachedMaskExprs []expression.Expression + getMaskExprs := func(p base.LogicalPlan) ([]expression.Expression, error) { + if cachedMaskPlan == p { + return cachedMaskExprs, nil + } + exprs, err := b.buildMaskingReplaceExprs(ctx, p) + if err != nil { + return nil, err + } + cachedMaskPlan = p + cachedMaskExprs = exprs + return exprs, nil + } for i, field := range fields { if !field.Auxiliary { oldLen++ @@ -1823,6 +2034,15 @@ func (b *PlanBuilder) buildProjection(ctx context.Context, p base.LogicalPlan, f } p = np + if applyMasking && !field.Auxiliary { + maskExprs, err := getMaskExprs(p) + if err != nil { + return nil, nil, 0, err + } + if maskExprs != nil { + newExpr = expression.ColumnSubstitute(b.ctx.GetExprCtx(), newExpr, p.Schema(), maskExprs) + } + } proj.Exprs = append(proj.Exprs, newExpr) col, name, err := b.buildProjectionField(ctx, p, field, newExpr) @@ -2119,6 +2339,10 @@ func (b *PlanBuilder) buildSetOpr(ctx context.Context, setOpr *ast.SetOprStmt) ( } } + // Track nesting depth so masking is applied only on the outermost set-operator result. + b.buildingSetOprOperands++ + defer func() { b.buildingSetOprOperands-- }() + // Because INTERSECT has higher precedence than UNION and EXCEPT. We build it first. selectPlans := make([]base.LogicalPlan, 0, len(setOpr.SelectList.Selects)) afterSetOprs := make([]*ast.SetOprType, 0, len(setOpr.SelectList.Selects)) @@ -2193,8 +2417,19 @@ func (b *PlanBuilder) buildSetOpr(ctx context.Context, setOpr *ast.SetOprStmt) ( } proj.SetOutputNames(setOprPlan.OutputNames()[:oldLen]) proj.SetSchema(schema) - return b.tryToBuildSequence(currentLayerCTEs, proj), nil + setOprPlan = proj + } + + if b.buildingSetOprOperands == 1 && !b.isCTE && !b.buildingCTE { + // Apply masking at the final result stage (AT RESULT semantics). + // This ensures set operators (UNION/INTERSECT/EXCEPT) use original values. + // For CTEs, we skip masking here because CTE definitions should preserve original values. + setOprPlan, err = b.buildFinalProjectionWithMasking(ctx, setOprPlan, oldLen, nil) + if err != nil { + return nil, err + } } + return b.tryToBuildSequence(currentLayerCTEs, setOprPlan), nil } @@ -4522,7 +4757,7 @@ func (b *PlanBuilder) buildSelect(ctx context.Context, sel *ast.SelectStmt) (p b var oldLen int // According to https://dev.mysql.com/doc/refman/8.0/en/window-functions-usage.html, // we can only process window functions after having clause, so `considerWindow` is false now. - p, projExprs, oldLen, err = b.buildProjection(ctx, p, sel.Fields.Fields, totalMap, nil, false, sel.OrderBy != nil) + p, projExprs, oldLen, err = b.buildProjection(ctx, p, sel.Fields.Fields, totalMap, nil, false, sel.OrderBy != nil, false) if err != nil { return nil, err } @@ -4560,7 +4795,7 @@ func (b *PlanBuilder) buildSelect(ctx context.Context, sel *ast.SelectStmt) (p b // In such case plan `p` is not changed, so we don't have to build another projection. if hasWindowFuncField { // Now we build the window function fields. - p, projExprs, oldLen, err = b.buildProjection(ctx, p, sel.Fields.Fields, windowAggMap, windowMapper, true, false) + p, projExprs, oldLen, err = b.buildProjection(ctx, p, sel.Fields.Fields, windowAggMap, windowMapper, true, false, false) if err != nil { return nil, err } @@ -4612,7 +4847,19 @@ func (b *PlanBuilder) buildSelect(ctx context.Context, sel *ast.SelectStmt) (p b } proj.SetOutputNames(p.OutputNames()[:oldLen]) proj.SetSchema(schema) - return b.tryToBuildSequence(currentLayerCTEs, proj), nil + p = proj + } + + // Apply masking at the final result stage (AT RESULT semantics). + // This ensures HAVING, ORDER BY, set operators, etc. all used original values. + // For CTEs, we skip masking here because: + // 1. CTE definitions should preserve original values for correct filtering/joining + // 2. Masking is applied when CTE results are materialized to the final output + if b.buildingSetOprOperands == 0 && !b.isCTE && !b.buildingCTE { + p, err = b.buildFinalProjectionWithMasking(ctx, p, oldLen, originalFields) + if err != nil { + return nil, err + } } return b.tryToBuildSequence(currentLayerCTEs, p), nil @@ -4898,6 +5145,17 @@ func (b *PlanBuilder) computeCTEInlineFlag(cte *cteInfo) { } func (b *PlanBuilder) buildDataSourceFromCTEMerge(ctx context.Context, cte *ast.CommonTableExpression) (base.LogicalPlan, error) { + // Set b.isCTE to true to ensure masking is skipped during CTE definition building. + // This preserves original values in CTE for correct WHERE/HAVING/GROUP BY behavior. + oldIsCTE := b.isCTE + oldBuildingCTE := b.buildingCTE + b.isCTE = true + b.buildingCTE = true + defer func() { + b.isCTE = oldIsCTE + b.buildingCTE = oldBuildingCTE + }() + p, err := b.buildResultSetNode(ctx, cte.Query.Query, true) if err != nil { return nil, err @@ -4905,6 +5163,14 @@ func (b *PlanBuilder) buildDataSourceFromCTEMerge(ctx context.Context, cte *ast. b.handleHelper.popMap() outPutNames := p.OutputNames() for _, name := range outPutNames { + // Preserve OrigTblName and OrigColName for masking policy lookup. + // If they are not set, copy from TblName and ColName before overwriting. + if name.OrigTblName.L == "" { + name.OrigTblName = name.TblName + } + if name.OrigColName.L == "" { + name.OrigColName = name.ColName + } name.TblName = cte.Name name.DBName = ast.NewCIStr(b.ctx.GetSessionVars().CurrentDB) } @@ -4914,6 +5180,10 @@ func (b *PlanBuilder) buildDataSourceFromCTEMerge(ctx context.Context, cte *ast. return nil, errors.New("CTE columns length is not consistent") } for i, n := range cte.ColNameList { + // Preserve original column name before overwriting. + if outPutNames[i].OrigColName.L == "" { + outPutNames[i].OrigColName = outPutNames[i].ColName + } outPutNames[i].ColName = n } } diff --git a/pkg/planner/core/masking_policy_at_result_test.go b/pkg/planner/core/masking_policy_at_result_test.go new file mode 100644 index 0000000000000..ac8a6bc8bd93c --- /dev/null +++ b/pkg/planner/core/masking_policy_at_result_test.go @@ -0,0 +1,94 @@ +// Copyright 2026 PingCAP, Inc. +// +// Licensed 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 core_test + +import ( + "testing" + + "github.com/pingcap/tidb/pkg/testkit" + "github.com/stretchr/testify/require" +) + +// TestMaskingPolicyAtResult tests that masking follows AT RESULT semantics: +// HAVING, ORDER BY, and set operators use original values, not masked values. +func TestMaskingPolicyAtResult(t *testing.T) { + store, _ := testkit.CreateMockStoreAndDomain(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test") + + // Test HAVING uses original values + t.Run("HAVING uses original values", func(t *testing.T) { + tk.MustExec("drop table if exists t_having") + tk.MustExec("create table t_having(id int, val varchar(100))") + tk.MustExec("insert into t_having values (1, 'apple'), (2, 'banana'), (3, 'cherry')") + tk.MustExec("create masking policy p_having on t_having(val) as concat('[', val, ']') enable") + + // HAVING should compare original values, not masked values + result := tk.MustQuery("select val, count(*) from t_having group by val having val > 'apple'") + result.Check(testkit.Rows("[banana] 1", "[cherry] 1")) + + // This should return no rows because no value > 'zzzz' + rows := tk.MustQuery("select val from t_having group by val having val > 'zzzz'").Rows() + require.Len(t, rows, 0) + }) + + // Test ORDER BY uses original values + t.Run("ORDER BY uses original values", func(t *testing.T) { + tk.MustExec("drop table if exists t_orderby") + tk.MustExec("create table t_orderby(id int, val varchar(100))") + tk.MustExec("insert into t_orderby values (1, 'apple'), (2, 'banana'), (3, 'cherry')") + tk.MustExec("create masking policy p_orderby on t_orderby(val) as concat('[', val, ']') enable") + + // ORDER BY should sort by original values + result := tk.MustQuery("select id, val from t_orderby order by val") + result.Check(testkit.Rows("1 [apple]", "2 [banana]", "3 [cherry]")) + + // Reverse order + result = tk.MustQuery("select id, val from t_orderby order by val desc") + result.Check(testkit.Rows("3 [cherry]", "2 [banana]", "1 [apple]")) + }) + + // Test UNION DISTINCT uses original values for deduplication + t.Run("UNION DISTINCT uses original values", func(t *testing.T) { + tk.MustExec("drop table if exists t_union") + tk.MustExec("create table t_union(id int, val varchar(100))") + tk.MustExec("insert into t_union values (1, 'value1'), (2, 'value2')") + tk.MustExec("create masking policy p_union on t_union(val) as concat('[', val, ']') enable") + + // Two different original values should NOT be deduplicated + // even though they mask to the same pattern + result := tk.MustQuery("select val from t_union where id = 1 union distinct select val from t_union where id = 2") + rows := result.Rows() + require.Len(t, rows, 2) + }) + + // Test aggregate with HAVING + t.Run("Aggregate with HAVING uses original values", func(t *testing.T) { + tk.MustExec("drop table if exists t_agg") + tk.MustExec("create table t_agg(category varchar(50), amount int)") + tk.MustExec("insert into t_agg values ('A', 100), ('B', 200), ('C', 300)") + tk.MustExec("create masking policy p_agg on t_agg(category) as concat('_', category) enable") + + // HAVING should compare original category values + result := tk.MustQuery("select category, sum(amount) from t_agg group by category having category > 'A'") + rows := result.Rows() + require.Len(t, rows, 2) // B and C + + // All categories should be masked in output + require.Contains(t, []string{"_B", "_C"}, rows[0][0]) + require.Contains(t, []string{"_B", "_C"}, rows[1][0]) + require.NotEqual(t, rows[0][0], rows[1][0]) + }) +} diff --git a/pkg/planner/core/masking_policy_expr_cache.go b/pkg/planner/core/masking_policy_expr_cache.go new file mode 100644 index 0000000000000..64e79a1548767 --- /dev/null +++ b/pkg/planner/core/masking_policy_expr_cache.go @@ -0,0 +1,101 @@ +// Copyright 2026 PingCAP, Inc. +// +// Licensed 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 core + +import ( + "github.com/pingcap/errors" + "github.com/pingcap/tidb/pkg/expression" + "github.com/pingcap/tidb/pkg/meta/model" + "github.com/pingcap/tidb/pkg/sessionctx/variable" +) + +type maskingPolicyExprCacheEntry struct { + expr expression.Expression + column *expression.Column +} + +type maskingPolicyExprCache struct { + schemaVersion int64 + entries map[int64]*maskingPolicyExprCacheEntry + building map[int64]struct{} +} + +func getMaskingPolicyExprCache(sv *variable.SessionVars, schemaVersion int64) *maskingPolicyExprCache { + cacheObj, cacheVer := sv.GetMaskingPolicyExprCache() + if cacheObj != nil { + if cache, ok := cacheObj.(*maskingPolicyExprCache); ok && cacheVer == schemaVersion { + return cache + } + } + cache := &maskingPolicyExprCache{ + schemaVersion: schemaVersion, + entries: make(map[int64]*maskingPolicyExprCacheEntry), + building: make(map[int64]struct{}), + } + sv.SetMaskingPolicyExprCache(cache, schemaVersion) + return cache +} + +func getMaskingPolicyExpr( + ctx expression.BuildContext, + sv *variable.SessionVars, + schemaVersion int64, + policy *model.MaskingPolicyInfo, + tblInfo *model.TableInfo, + colInfo *model.ColumnInfo, +) (expression.Expression, *expression.Column, error) { + cache := getMaskingPolicyExprCache(sv, schemaVersion) + if entry, ok := cache.entries[policy.ID]; ok { + return entry.expr.Clone(), entry.column, nil + } + if _, ok := cache.building[policy.ID]; ok { + return nil, nil, errors.New("masking policy expression recursion detected") + } + cache.building[policy.ID] = struct{}{} + expr, placeholder, err := buildMaskingPolicyExpr(ctx, policy, tblInfo, colInfo) + delete(cache.building, policy.ID) + if err != nil { + return nil, nil, err + } + cache.entries[policy.ID] = &maskingPolicyExprCacheEntry{ + expr: expr, + column: placeholder, + } + return expr.Clone(), placeholder, nil +} + +func buildMaskingPolicyExpr( + ctx expression.BuildContext, + policy *model.MaskingPolicyInfo, + tblInfo *model.TableInfo, + colInfo *model.ColumnInfo, +) (expression.Expression, *expression.Column, error) { + if policy == nil || tblInfo == nil || colInfo == nil { + return nil, nil, errors.New("masking policy expression requires policy/table/column info") + } + cols, names, err := expression.ColumnInfos2ColumnsAndNames(ctx, policy.DBName, tblInfo.Name, []*model.ColumnInfo{colInfo}, tblInfo) + if err != nil { + return nil, nil, err + } + schema := expression.NewSchema(cols...) + expr, err := expression.ParseSimpleExpr(ctx, policy.Expression, expression.WithInputSchemaAndNames(schema, names, tblInfo)) + if err != nil { + return nil, nil, err + } + if len(cols) != 1 { + return nil, nil, errors.New("masking policy expression expects one input column") + } + return expr, cols[0], nil +} \ No newline at end of file diff --git a/pkg/planner/core/masking_policy_expr_cache_test.go b/pkg/planner/core/masking_policy_expr_cache_test.go new file mode 100644 index 0000000000000..500dc5a34402f --- /dev/null +++ b/pkg/planner/core/masking_policy_expr_cache_test.go @@ -0,0 +1,103 @@ +// Copyright 2026 PingCAP, Inc. +// +// Licensed 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 core + +import ( + "testing" + + "github.com/pingcap/tidb/pkg/expression/exprstatic" + "github.com/pingcap/tidb/pkg/meta/model" + "github.com/pingcap/tidb/pkg/parser/ast" + "github.com/pingcap/tidb/pkg/parser/mysql" + "github.com/pingcap/tidb/pkg/sessionctx/variable" + "github.com/pingcap/tidb/pkg/types" + "github.com/stretchr/testify/require" +) + +func TestMaskingPolicyExprCache(t *testing.T) { + sv := variable.NewSessionVars(nil) + ctx := exprstatic.NewExprContext() + tblInfo := &model.TableInfo{Name: ast.NewCIStr("t")} + colInfo := &model.ColumnInfo{ + ID: 1, + Name: ast.NewCIStr("c"), + Offset: 0, + FieldType: *types.NewFieldType(mysql.TypeString), + } + tblInfo.Columns = []*model.ColumnInfo{colInfo} + + policy := &model.MaskingPolicyInfo{ + ID: 1, + Name: ast.NewCIStr("p"), + DBName: ast.NewCIStr("test"), + TableName: tblInfo.Name, + TableID: 100, + ColumnName: colInfo.Name, + ColumnID: colInfo.ID, + Expression: "c", + } + + expr1, col1, err := getMaskingPolicyExpr(ctx, sv, 1, policy, tblInfo, colInfo) + require.NoError(t, err) + require.NotNil(t, expr1) + require.NotNil(t, col1) + + cache := getMaskingPolicyExprCache(sv, 1) + entry := cache.entries[policy.ID] + require.NotNil(t, entry) + + // Changing the policy expression while on the same schema version + // should still return the cached entry. + policy.Expression = "concat(c, 'x')" + _, _, err = getMaskingPolicyExpr(ctx, sv, 1, policy, tblInfo, colInfo) + require.NoError(t, err) + require.Same(t, entry, cache.entries[policy.ID]) + + // Bumping the schema version should create a fresh cache. + _, _, err = getMaskingPolicyExpr(ctx, sv, 2, policy, tblInfo, colInfo) + require.NoError(t, err) + newCache := getMaskingPolicyExprCache(sv, 2) + require.NotSame(t, cache, newCache) + require.Len(t, newCache.entries, 1) +} + +func TestMaskingPolicyExprCacheRecursion(t *testing.T) { + sv := variable.NewSessionVars(nil) + ctx := exprstatic.NewExprContext() + tblInfo := &model.TableInfo{Name: ast.NewCIStr("t")} + colInfo := &model.ColumnInfo{ + ID: 1, + Name: ast.NewCIStr("c"), + Offset: 0, + FieldType: *types.NewFieldType(mysql.TypeString), + } + tblInfo.Columns = []*model.ColumnInfo{colInfo} + + policy := &model.MaskingPolicyInfo{ + ID: 1, + Name: ast.NewCIStr("p"), + DBName: ast.NewCIStr("test"), + TableName: tblInfo.Name, + TableID: 100, + ColumnName: colInfo.Name, + ColumnID: colInfo.ID, + Expression: "c", + } + + cache := getMaskingPolicyExprCache(sv, 1) + cache.building[policy.ID] = struct{}{} + _, _, err := getMaskingPolicyExpr(ctx, sv, 1, policy, tblInfo, colInfo) + require.Error(t, err) +} \ No newline at end of file diff --git a/pkg/planner/core/masking_policy_projection_test.go b/pkg/planner/core/masking_policy_projection_test.go new file mode 100644 index 0000000000000..d67f1e9f8d278 --- /dev/null +++ b/pkg/planner/core/masking_policy_projection_test.go @@ -0,0 +1,100 @@ +// Copyright 2026 PingCAP, Inc. +// +// Licensed 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 core_test + +import ( + "testing" + + "github.com/pingcap/tidb/pkg/testkit" + "github.com/stretchr/testify/require" +) + +func TestMaskingPolicyProjection(t *testing.T) { + store, _ := testkit.CreateMockStoreAndDomain(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test") + tk.MustExec("drop table if exists t") + tk.MustExec("create table t(id int primary key, c varchar(10))") + tk.MustExec("insert into t values (1, 'a'), (2, 'b')") + tk.MustExec("create masking policy p on t(c) as concat(c, 'x') enable") + + tk.MustQuery("select c from t order by id").Check(testkit.Rows("ax", "bx")) + tk.MustQuery("select concat(c, '-') from t where c = 'a'").Check(testkit.Rows("ax-")) + + // Predicate should still use original value. + rows := tk.MustQuery("select c from t where c = 'ax'").Rows() + require.Len(t, rows, 0) +} + +func TestMaskingPolicyPointGet(t *testing.T) { + store, _ := testkit.CreateMockStoreAndDomain(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test") + tk.MustExec("drop table if exists t_pointget") + tk.MustExec("create table t_pointget(id int primary key, c varchar(10))") + tk.MustExec("insert into t_pointget values (1, 'secret'), (2, 'hidden')") + tk.MustExec("create masking policy p_pointget on t_pointget(c) as concat('***', c) enable") + + // Test Point_Get with primary key lookup + tk.MustQuery("select c from t_pointget where id = 1").Check(testkit.Rows("***secret")) + tk.MustQuery("select c from t_pointget where id = 2").Check(testkit.Rows("***hidden")) + + // Test Point_Get with multiple columns + tk.MustQuery("select id, c from t_pointget where id = 1").Check(testkit.Rows("1 ***secret")) + + // Test Point_Get with expression using masked column + tk.MustQuery("select concat(c, '-') from t_pointget where id = 1").Check(testkit.Rows("***secret-")) + + // Test Point_Get with unique index lookup + tk.MustExec("create unique index idx_c on t_pointget(c)") + tk.MustQuery("select c from t_pointget where c = 'secret'").Check(testkit.Rows("***secret")) + + // Predicate should still use original value + rows := tk.MustQuery("select c from t_pointget where c = '***secret'").Rows() + require.Len(t, rows, 0) +} + +func TestMaskingPolicyBatchPointGet(t *testing.T) { + store, _ := testkit.CreateMockStoreAndDomain(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test") + tk.MustExec("drop table if exists t_batch_pointget") + tk.MustExec("create table t_batch_pointget(id int primary key, c varchar(10))") + tk.MustExec("insert into t_batch_pointget values (1, 'secret'), (2, 'hidden'), (3, 'masked')") + tk.MustExec("create masking policy p_batch on t_batch_pointget(c) as concat('***', c) enable") + + // Test BatchPointGet with WHERE pk IN (...) + tk.MustQuery("select c from t_batch_pointget where id in (1, 2)").Check(testkit.Rows("***secret", "***hidden")) + + // Test BatchPointGet with multiple values + tk.MustQuery("select c from t_batch_pointget where id in (1, 2, 3)").Check(testkit.Rows("***secret", "***hidden", "***masked")) + + // Test BatchPointGet with multiple columns + tk.MustQuery("select id, c from t_batch_pointget where id in (1, 2)").Check(testkit.Rows("1 ***secret", "2 ***hidden")) + + // Test BatchPointGet with expression using masked column + tk.MustQuery("select concat(c, '-') from t_batch_pointget where id in (1)").Check(testkit.Rows("***secret-")) + + // Test BatchPointGet with unique index IN clause + tk.MustExec("create unique index idx_c on t_batch_pointget(c)") + tk.MustQuery("select c from t_batch_pointget where c in ('secret', 'hidden')").Check(testkit.Rows("***secret", "***hidden")) + + // Predicate should still use original value + rows := tk.MustQuery("select c from t_batch_pointget where c in ('***secret')").Rows() + require.Len(t, rows, 0) + + // Test single value in IN (might use PointGet or BatchPointGet) + tk.MustQuery("select c from t_batch_pointget where id in (1)").Check(testkit.Rows("***secret")) +} diff --git a/pkg/planner/core/operator/physicalop/physical_batch_point_get.go b/pkg/planner/core/operator/physicalop/physical_batch_point_get.go index 938de3de6c4b6..b7d4fe4144590 100644 --- a/pkg/planner/core/operator/physicalop/physical_batch_point_get.go +++ b/pkg/planner/core/operator/physicalop/physical_batch_point_get.go @@ -84,6 +84,11 @@ type PointGetPlan struct { LockWaitTime int64 Columns []*model.ColumnInfo `plan-cache-clone:"shallow"` + // MaskingExprs stores the masking expressions for columns that have masking policies. + // If not nil, each element corresponds to a column in the schema. + // The executor should evaluate these expressions instead of returning raw column values. + MaskingExprs []expression.Expression `plan-cache-clone:"shallow"` + // required by cost model cost float64 PlanCostInit bool @@ -520,7 +525,13 @@ type BatchPointGetPlan struct { Lock bool LockWaitTime int64 Columns []*model.ColumnInfo `plan-cache-clone:"shallow"` - cost float64 + + // MaskingExprs stores the masking expressions for columns that have masking policies. + // If not nil, each element corresponds to a column in the schema. + // The executor should evaluate these expressions instead of returning raw column values. + MaskingExprs []expression.Expression `plan-cache-clone:"shallow"` + + cost float64 // required by cost model PlanCostInit bool diff --git a/pkg/planner/core/planbuilder.go b/pkg/planner/core/planbuilder.go index c4faac5fbbcb8..9cfffbcab671e 100644 --- a/pkg/planner/core/planbuilder.go +++ b/pkg/planner/core/planbuilder.go @@ -285,6 +285,10 @@ type PlanBuilder struct { allocIDForCTEStorage int buildingRecursivePartForCTE bool buildingCTE bool + // buildingSetOprOperands is the nesting depth while building set-operator operands. + // When > 0, final masking projection should be skipped for inner SELECT/SET plans and + // applied only once on the completed set-operator result. + buildingSetOprOperands int // Check whether the current building query is a CTE isCTE bool // CTE table name in lower case, it can be nil diff --git a/pkg/planner/core/point_get_plan.go b/pkg/planner/core/point_get_plan.go index de6a4548f8722..413d44606f91a 100644 --- a/pkg/planner/core/point_get_plan.go +++ b/pkg/planner/core/point_get_plan.go @@ -506,6 +506,15 @@ func tryWhereIn2BatchPointGet(ctx base.PlanContext, selStmt *ast.SelectStmt, res } p.DBName = dbName + // Build masking expressions for columns with masking policies. + if is := ctx.GetInfoSchema(); is != nil && len(names) > 0 { + maskExprs, err := buildMaskingExprsForPointGet(context.Background(), ctx, is.(infoschema.InfoSchema), schema, names, tbl) + if err != nil { + return nil + } + p.MaskingExprs = maskExprs + } + return p } @@ -763,6 +772,20 @@ func newPointGetPlan(ctx base.PlanContext, dbName string, schema *expression.Sch p.Plan.SetStats(&property.StatsInfo{RowCount: 1}) ctx.GetSessionVars().StmtCtx.Tables = []stmtctx.TableEntry{{DB: dbName, Table: tbl.Name.L}} + + // Build masking expressions for columns with masking policies. + // Only apply masking for top-level result-producing SELECT statements. + // We detect synthetic SELECTs for DML by checking if output names are empty. + if is := ctx.GetInfoSchema(); is != nil && len(names) > 0 { + maskExprs, err := buildMaskingExprsForPointGet(context.Background(), ctx, is.(infoschema.InfoSchema), schema, names, tbl) + if err != nil { + // Fail-closed: do not use PointGet fast path when masking cannot be applied. + // The query will fall back to the normal planner path. + return nil + } + p.MaskingExprs = maskExprs + } + return p } @@ -1462,3 +1485,78 @@ func getHashOrKeyPartitionColumnName(ctx base.PlanContext, tbl *model.TableInfo) } return &col.Name.Name } + +// buildMaskingExprsForPointGet builds masking expressions for PointGetPlan/BatchPointGetPlan. +// This is similar to buildMaskingReplaceExprs in logical_plan_builder.go +// but adapted for the fast path. +func buildMaskingExprsForPointGet( + ctx context.Context, + sctx base.PlanContext, + is infoschema.InfoSchema, + schema *expression.Schema, + names []*types.FieldName, + tbl *model.TableInfo, +) ([]expression.Expression, error) { + if is == nil || schema == nil || len(names) == 0 { + return nil, nil + } + sv := sctx.GetSessionVars() + if sv != nil && sv.InRestrictedSQL { + return nil, nil + } + if len(is.AllMaskingPolicies()) == 0 { + return nil, nil + } + cols := schema.Columns + if len(cols) == 0 || len(cols) != len(names) { + return nil, nil + } + replaceExprs := make([]expression.Expression, len(cols)) + for i, col := range cols { + replaceExprs[i] = col + } + schemaVersion := is.SchemaMetaVersion() + hasMask := false + for i, col := range cols { + policy, colInfo := findMaskingPolicyForColumn(is, tbl, names[i]) + if policy == nil || colInfo == nil { + continue + } + if policy.Status != model.MaskingPolicyStatusEnable { + continue + } + expr, placeholder, err := getMaskingPolicyExpr(sctx.GetExprCtx(), sv, schemaVersion, policy, tbl, colInfo) + if err != nil { + return nil, err + } + if placeholder == nil { + continue + } + masked := expression.ColumnSubstitute(sctx.GetExprCtx(), expr, expression.NewSchema(placeholder), []expression.Expression{col}) + replaceExprs[i] = masked + hasMask = true + } + if !hasMask { + return nil, nil + } + return replaceExprs, nil +} + +func findMaskingPolicyForColumn(is infoschema.InfoSchema, tbl *model.TableInfo, name *types.FieldName) (*model.MaskingPolicyInfo, *model.ColumnInfo) { + if tbl == nil || name == nil { + return nil, nil + } + colName := name.OrigColName + if colName.L == "" { + colName = name.ColName + } + colInfo := tbl.FindPublicColumnByName(colName.L) + if colInfo == nil { + return nil, nil + } + policy, ok := is.MaskingPolicyByTableColumn(tbl.ID, colInfo.ID) + if !ok || policy == nil { + return nil, nil + } + return policy, colInfo +} diff --git a/pkg/sessionctx/variable/masking_policy_cache.go b/pkg/sessionctx/variable/masking_policy_cache.go new file mode 100644 index 0000000000000..bd7916b4fa715 --- /dev/null +++ b/pkg/sessionctx/variable/masking_policy_cache.go @@ -0,0 +1,26 @@ +// Copyright 2026 PingCAP, Inc. +// +// Licensed 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 variable + +// GetMaskingPolicyExprCache returns the cached object and the schema version it is built on. +func (s *SessionVars) GetMaskingPolicyExprCache() (any, int64) { + return s.maskingPolicyExprCache, s.maskingPolicyExprCacheSchemaVersion +} + +// SetMaskingPolicyExprCache stores the cache object along with its schema version. +func (s *SessionVars) SetMaskingPolicyExprCache(cache any, schemaVersion int64) { + s.maskingPolicyExprCache = cache + s.maskingPolicyExprCacheSchemaVersion = schemaVersion +} diff --git a/pkg/sessionctx/variable/session.go b/pkg/sessionctx/variable/session.go index 96bcb5ba3abc6..220e7a024f14d 100644 --- a/pkg/sessionctx/variable/session.go +++ b/pkg/sessionctx/variable/session.go @@ -840,6 +840,11 @@ type SessionVars struct { // The Cached Plan for this execution, it should be *plannercore.PlanCacheValue. PlanCacheValue any + // maskingPolicyExprCache caches compiled masking policy expressions per schema version. + maskingPolicyExprCache any + // maskingPolicyExprCacheSchemaVersion is the schema version for maskingPolicyExprCache. + maskingPolicyExprCacheSchemaVersion int64 + // ActiveRoles stores active roles for current user ActiveRoles []*auth.RoleIdentity From a609ceff9917225004e167af94abfb222b0ecf59 Mon Sep 17 00:00:00 2001 From: tiancaiamao Date: Tue, 9 Jun 2026 23:25:51 +0800 Subject: [PATCH 02/13] fix: address CodeRabbit review findings 1. DML masking in PointGet fast path: moved masking out of newPointGetPlan into tryPointGetPlan via maybeApplyPointGetMasking helper, which checks len(selStmt.Fields.Fields) > 0 to skip DML synthetic SELECTs. Similarly fixed tryWhereIn2BatchPointGet. 2. MaskingExprs in Recreated/Init: moved assignment from buildPointGet into Init(p) so Recreated path also refreshes the field. 3. findMaskingPolicy CurrentDB fallback: when DBName is empty, return nil instead of guessing via CurrentDB (avoids wrong-table masking). 4. buildProjection4Union OrigTblName/OrigColName: preserve all FieldName fields from first child, not just ColName. 5. adjustCTEPlanOutputName OrigColName: don't overwrite OrigColName with CTE alias when it was already set by the CTE merge path. 6. applyMaskingExprs arity check: add defensive length validation. --- pkg/executor/point_get.go | 5 ++- pkg/planner/core/logical_plan_builder.go | 19 +++++++++-- pkg/planner/core/point_get_plan.go | 42 +++++++++++++++--------- 3 files changed, 47 insertions(+), 19 deletions(-) diff --git a/pkg/executor/point_get.go b/pkg/executor/point_get.go index e49831e0134fe..c22f9108df424 100644 --- a/pkg/executor/point_get.go +++ b/pkg/executor/point_get.go @@ -95,7 +95,6 @@ func (b *executorBuilder) buildPointGet(p *physicalop.PointGetPlan) exec.Executo e.SetInitCap(1) e.SetMaxChunkSize(1) e.Init(p) - e.MaskingExprs = p.MaskingExprs snapshotTS, err := b.getSnapshotTS() if err != nil { @@ -218,6 +217,7 @@ func (e *PointGetExecutor) Init(p *physicalop.PointGetPlan) { e.rowDecoder = decoder e.partitionDefIdx = p.PartitionIdx e.columns = p.Columns + e.MaskingExprs = p.MaskingExprs e.buildVirtualColumnInfo() sessVars := e.Ctx().GetSessionVars() @@ -465,6 +465,9 @@ func (e *PointGetExecutor) Next(ctx context.Context, req *chunk.Chunk) error { // applyMaskingExprs evaluates masking expressions for each row in the chunk and // replaces the original column values with masked values. func applyMaskingExprs(sctx sessionctx.Context, schema *expression.Schema, maskingExprs []expression.Expression, req *chunk.Chunk) error { + if len(maskingExprs) != len(schema.Columns) { + return errors.Errorf("masking expr count mismatch: got %d, schema has %d columns", len(maskingExprs), len(schema.Columns)) + } fieldTypes := make([]*types.FieldType, len(schema.Columns)) for i, col := range schema.Columns { fieldTypes[i] = col.RetType diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go index 1276118b5f37a..fa3cf67cc60f4 100644 --- a/pkg/planner/core/logical_plan_builder.go +++ b/pkg/planner/core/logical_plan_builder.go @@ -1595,7 +1595,9 @@ func (b *PlanBuilder) findMaskingPolicy(_ context.Context, name *types.FieldName } dbName := name.DBName if dbName.L == "" { - dbName = ast.NewCIStr(b.ctx.GetSessionVars().CurrentDB) + // Cannot determine source database (e.g., derived table aliasing clears DBName). + // Skip masking rather than risk applying the wrong policy via CurrentDB fallback. + return nil, nil, nil } tbl, err := b.is.TableByName(context.Background(), dbName, tblName) if err != nil { @@ -2291,7 +2293,13 @@ func (b *PlanBuilder) buildProjection4Union(_ context.Context, u *logicalop.Logi resultTp.SetCharset(collation.Charset) resultTp.SetCollate(collation.Collation) b.setUnionFlen(resultTp, tmpExprs) - names = append(names, &types.FieldName{ColName: u.Children()[0].OutputNames()[i].ColName}) + firstChildNames := u.Children()[0].OutputNames() + names = append(names, &types.FieldName{ + DBName: firstChildNames[i].DBName, + OrigTblName: firstChildNames[i].OrigTblName, + OrigColName: firstChildNames[i].OrigColName, + ColName: firstChildNames[i].ColName, + }) unionCols = append(unionCols, &expression.Column{ RetType: resultTp, UniqueID: b.ctx.GetSessionVars().AllocPlanColumnID(), @@ -8213,7 +8221,12 @@ func (b *PlanBuilder) adjustCTEPlanOutputName(p base.LogicalPlan, def *ast.Commo } for i, n := range def.ColNameList { clonedNames[i].ColName = n - clonedNames[i].OrigColName = n + // Preserve OrigColName for masking policy lookup — it points to the + // real source column name, which is needed for policy resolution. + // Only set OrigColName when it wasn't already set by the CTE merge path. + if clonedNames[i].OrigColName.L == "" { + clonedNames[i].OrigColName = n + } } } p.SetOutputNames(clonedNames) diff --git a/pkg/planner/core/point_get_plan.go b/pkg/planner/core/point_get_plan.go index 413d44606f91a..2115ba31132b9 100644 --- a/pkg/planner/core/point_get_plan.go +++ b/pkg/planner/core/point_get_plan.go @@ -507,7 +507,9 @@ func tryWhereIn2BatchPointGet(ctx base.PlanContext, selStmt *ast.SelectStmt, res p.DBName = dbName // Build masking expressions for columns with masking policies. - if is := ctx.GetInfoSchema(); is != nil && len(names) > 0 { + // Only apply masking for real SELECT queries (not DML synthetic SELECTs + // which have empty Fields). + if is := ctx.GetInfoSchema(); is != nil && len(selStmt.Fields.Fields) > 0 { maskExprs, err := buildMaskingExprsForPointGet(context.Background(), ctx, is.(infoschema.InfoSchema), schema, names, tbl) if err != nil { return nil @@ -572,6 +574,7 @@ func tryPointGetPlan(ctx base.PlanContext, selStmt *ast.SelectStmt, resolveCtx * handlePair, fieldType := findPKHandle(tbl, pairs) dbName := getLowerDB(tblName.Schema, ctx.GetSessionVars()) + isSelect := len(selStmt.Fields.Fields) > 0 if handlePair.value.Kind() != types.KindNull && len(pairs) == 1 && indexIsAvailableByHints( ctx.GetSessionVars().CurrentDB, @@ -584,7 +587,7 @@ func tryPointGetPlan(ctx base.PlanContext, selStmt *ast.SelectStmt, resolveCtx * if isTableDual { p := newPointGetPlan(ctx, dbName, schema, tbl, names) p.IsTableDual = true - return p + return maybeApplyPointGetMasking(ctx, p, isSelect, schema, names, tbl) } p := newPointGetPlan(ctx, dbName, schema, tbl, names) @@ -594,12 +597,13 @@ func tryPointGetPlan(ctx base.PlanContext, selStmt *ast.SelectStmt, resolveCtx * p.HandleConstant = handlePair.con p.HandleColOffset = pkColOffset p.PartitionNames = tblName.PartitionNames - return p + return maybeApplyPointGetMasking(ctx, p, isSelect, schema, names, tbl) } else if handlePair.value.Kind() != types.KindNull { return nil } - return checkTblIndexForPointPlan(ctx, tnW, schema, tblAlias.L, selStmt.TableHints, names, pairs, isTableDual, check) + p := checkTblIndexForPointPlan(ctx, tnW, schema, tblAlias.L, selStmt.TableHints, names, pairs, isTableDual, check) + return maybeApplyPointGetMasking(ctx, p, isSelect, schema, names, tbl) } func checkTblIndexForPointPlan(ctx base.PlanContext, tblName *resolve.TableNameW, schema *expression.Schema, @@ -775,17 +779,6 @@ func newPointGetPlan(ctx base.PlanContext, dbName string, schema *expression.Sch // Build masking expressions for columns with masking policies. // Only apply masking for top-level result-producing SELECT statements. - // We detect synthetic SELECTs for DML by checking if output names are empty. - if is := ctx.GetInfoSchema(); is != nil && len(names) > 0 { - maskExprs, err := buildMaskingExprsForPointGet(context.Background(), ctx, is.(infoschema.InfoSchema), schema, names, tbl) - if err != nil { - // Fail-closed: do not use PointGet fast path when masking cannot be applied. - // The query will fall back to the normal planner path. - return nil - } - p.MaskingExprs = maskExprs - } - return p } @@ -1542,6 +1535,25 @@ func buildMaskingExprsForPointGet( return replaceExprs, nil } +// maybeApplyPointGetMasking builds and assigns masking expressions to a PointGetPlan. +// It only applies masking for real SELECT queries (isSelect=true), not for DML +// synthetic SELECTs used by UPDATE/DELETE fast paths. +func maybeApplyPointGetMasking(ctx base.PlanContext, p *physicalop.PointGetPlan, isSelect bool, schema *expression.Schema, names []*types.FieldName, tbl *model.TableInfo) *physicalop.PointGetPlan { + if p == nil || !isSelect { + return p + } + if is := ctx.GetInfoSchema(); is != nil && len(names) > 0 { + maskExprs, err := buildMaskingExprsForPointGet(context.Background(), ctx, is.(infoschema.InfoSchema), schema, names, tbl) + if err != nil { + // Fail-closed: do not use PointGet fast path when masking cannot be applied. + // The query will fall back to the normal planner path. + return nil + } + p.MaskingExprs = maskExprs + } + return p +} + func findMaskingPolicyForColumn(is infoschema.InfoSchema, tbl *model.TableInfo, name *types.FieldName) (*model.MaskingPolicyInfo, *model.ColumnInfo) { if tbl == nil || name == nil { return nil, nil From 2ca83898a6c206df7637b9583fc8c1c4ddb8c8f4 Mon Sep 17 00:00:00 2001 From: tiancaiamao Date: Tue, 9 Jun 2026 23:39:06 +0800 Subject: [PATCH 03/13] make fmt & make bazel_prepare --- pkg/planner/core/BUILD.bazel | 4 ++++ pkg/planner/core/masking_policy_expr_cache.go | 2 +- pkg/planner/core/masking_policy_expr_cache_test.go | 2 +- pkg/sessionctx/variable/BUILD.bazel | 1 + 4 files changed, 7 insertions(+), 2 deletions(-) diff --git a/pkg/planner/core/BUILD.bazel b/pkg/planner/core/BUILD.bazel index 679019205ae26..68dfe36e003b9 100644 --- a/pkg/planner/core/BUILD.bazel +++ b/pkg/planner/core/BUILD.bazel @@ -26,6 +26,7 @@ go_library( "initialize.go", "logical_initialize.go", "logical_plan_builder.go", + "masking_policy_expr_cache.go", "memtable_infoschema_extractor.go", "memtable_predicate_extractor.go", "optimizer.go", @@ -225,6 +226,9 @@ go_test( "lateral_join_test.go", "logical_plans_test.go", "main_test.go", + "masking_policy_at_result_test.go", + "masking_policy_expr_cache_test.go", + "masking_policy_projection_test.go", "optimizer_test.go", "physical_plan_test.go", "plan_cache_instance_test.go", diff --git a/pkg/planner/core/masking_policy_expr_cache.go b/pkg/planner/core/masking_policy_expr_cache.go index 64e79a1548767..0a927133ae9d4 100644 --- a/pkg/planner/core/masking_policy_expr_cache.go +++ b/pkg/planner/core/masking_policy_expr_cache.go @@ -98,4 +98,4 @@ func buildMaskingPolicyExpr( return nil, nil, errors.New("masking policy expression expects one input column") } return expr, cols[0], nil -} \ No newline at end of file +} diff --git a/pkg/planner/core/masking_policy_expr_cache_test.go b/pkg/planner/core/masking_policy_expr_cache_test.go index 500dc5a34402f..60ce475260824 100644 --- a/pkg/planner/core/masking_policy_expr_cache_test.go +++ b/pkg/planner/core/masking_policy_expr_cache_test.go @@ -100,4 +100,4 @@ func TestMaskingPolicyExprCacheRecursion(t *testing.T) { cache.building[policy.ID] = struct{}{} _, _, err := getMaskingPolicyExpr(ctx, sv, 1, policy, tblInfo, colInfo) require.Error(t, err) -} \ No newline at end of file +} diff --git a/pkg/sessionctx/variable/BUILD.bazel b/pkg/sessionctx/variable/BUILD.bazel index 0612c7b74a366..851ddb01f0416 100644 --- a/pkg/sessionctx/variable/BUILD.bazel +++ b/pkg/sessionctx/variable/BUILD.bazel @@ -4,6 +4,7 @@ go_library( name = "variable", srcs = [ "error.go", + "masking_policy_cache.go", "mock_globalaccessor.go", "noop.go", "removed.go", From 9790e1b0c579d0ddf010d032c2178626f28d552c Mon Sep 17 00:00:00 2001 From: tiancaiamao Date: Wed, 10 Jun 2026 09:42:16 +0800 Subject: [PATCH 04/13] make lint check happy --- pkg/executor/point_get.go | 2 +- pkg/planner/core/logical_plan_builder.go | 9 ++++++--- pkg/planner/core/point_get_plan.go | 2 +- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/pkg/executor/point_get.go b/pkg/executor/point_get.go index c22f9108df424..69d5da37e6820 100644 --- a/pkg/executor/point_get.go +++ b/pkg/executor/point_get.go @@ -474,7 +474,7 @@ func applyMaskingExprs(sctx sessionctx.Context, schema *expression.Schema, maski } maskedChunk := chunk.NewChunkWithCapacity(fieldTypes, req.NumRows()) evalCtx := sctx.GetExprCtx().GetEvalCtx() - for rowIdx := 0; rowIdx < req.NumRows(); rowIdx++ { + for rowIdx := range req.NumRows() { row := req.GetRow(rowIdx) for colIdx, expr := range maskingExprs { if expr != nil { diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go index fa3cf67cc60f4..c3e8720c977fa 100644 --- a/pkg/planner/core/logical_plan_builder.go +++ b/pkg/planner/core/logical_plan_builder.go @@ -1611,6 +1611,9 @@ func (b *PlanBuilder) findMaskingPolicy(_ context.Context, name *types.FieldName if colInfo == nil { return nil, nil, nil } + if colInfo.ID != col.ID { + return nil, nil, nil + } policy, ok := b.is.MaskingPolicyByTableColumn(tblInfo.ID, colInfo.ID) if !ok || policy == nil { return nil, nil, nil @@ -1642,7 +1645,7 @@ func (b *PlanBuilder) buildFinalProjectionWithMasking(ctx context.Context, p bas if len(originalFields) == 0 || len(originalFields) < oldLen { return b.buildFinalProjectionWithMaskingSimple(ctx, p, oldLen) } - for i := 0; i < oldLen; i++ { + for i := range oldLen { field := originalFields[i] if field == nil || field.Auxiliary || field.WildCard != nil || field.Expr == nil { return b.buildFinalProjectionWithMaskingSimple(ctx, p, oldLen) @@ -1668,7 +1671,7 @@ func (b *PlanBuilder) buildFinalProjectionWithMasking(ctx context.Context, p bas finalProj := logicalop.LogicalProjection{Exprs: make([]expression.Expression, 0, oldLen)}.Init(b.ctx, b.getSelectOffset()) schema := expression.NewSchema(make([]*expression.Column, 0, oldLen)...) newNames := make([]*types.FieldName, 0, oldLen) - for i := 0; i < oldLen; i++ { + for i := range oldLen { expr := expression.ColumnSubstitute(b.ctx.GetExprCtx(), proj.Exprs[i], child.Schema(), childMaskExprs) finalProj.Exprs = append(finalProj.Exprs, expr) @@ -1709,7 +1712,7 @@ func (b *PlanBuilder) buildFinalProjectionWithMaskingSimple(ctx context.Context, proj := logicalop.LogicalProjection{Exprs: make([]expression.Expression, 0, oldLen)}.Init(b.ctx, b.getSelectOffset()) schema := expression.NewSchema(make([]*expression.Column, 0, oldLen)...) newNames := make([]*types.FieldName, 0, oldLen) - for i := 0; i < oldLen; i++ { + for i := range oldLen { col := p.Schema().Columns[i] expr := expression.ColumnSubstitute(b.ctx.GetExprCtx(), col, p.Schema(), maskExprs) proj.Exprs = append(proj.Exprs, expr) diff --git a/pkg/planner/core/point_get_plan.go b/pkg/planner/core/point_get_plan.go index 2115ba31132b9..f4ba7cc86e423 100644 --- a/pkg/planner/core/point_get_plan.go +++ b/pkg/planner/core/point_get_plan.go @@ -1483,7 +1483,7 @@ func getHashOrKeyPartitionColumnName(ctx base.PlanContext, tbl *model.TableInfo) // This is similar to buildMaskingReplaceExprs in logical_plan_builder.go // but adapted for the fast path. func buildMaskingExprsForPointGet( - ctx context.Context, + _ context.Context, sctx base.PlanContext, is infoschema.InfoSchema, schema *expression.Schema, From 19669cf2ab5c4b5f6be59c810ea17f6e8129c9e8 Mon Sep 17 00:00:00 2001 From: tiancaiamao Date: Wed, 10 Jun 2026 16:07:35 +0800 Subject: [PATCH 05/13] fix: CTE panic caused by incorrectly setting buildingCTE flag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The buildingCTE flag has a specific meaning for recursive CTE tracking — it requires b.outerCTEs to be non-empty. Setting it in buildDataSourceFromCTEMerge (for non-recursive CTEs) caused panics when buildDistinct/buildAgg accessed b.outerCTEs[-1]. Fix: only set b.isCTE=true in buildDataSourceFromCTEMerge, not b.buildingCTE. The isCTE flag is sufficient for masking guard checks. Also fixes test ordering flakiness in TestMaskingPolicyAtResult by adding ORDER BY clause to non-deterministic GROUP BY query. --- pkg/planner/core/logical_plan_builder.go | 5 ++--- pkg/planner/core/masking_policy_at_result_test.go | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go index c3e8720c977fa..b07daf41af915 100644 --- a/pkg/planner/core/logical_plan_builder.go +++ b/pkg/planner/core/logical_plan_builder.go @@ -5158,13 +5158,12 @@ func (b *PlanBuilder) computeCTEInlineFlag(cte *cteInfo) { func (b *PlanBuilder) buildDataSourceFromCTEMerge(ctx context.Context, cte *ast.CommonTableExpression) (base.LogicalPlan, error) { // Set b.isCTE to true to ensure masking is skipped during CTE definition building. // This preserves original values in CTE for correct WHERE/HAVING/GROUP BY behavior. + // Note: we do NOT set b.buildingCTE here — that flag has a specific meaning for + // recursive CTE tracking and requires b.outerCTEs to be non-empty. oldIsCTE := b.isCTE - oldBuildingCTE := b.buildingCTE b.isCTE = true - b.buildingCTE = true defer func() { b.isCTE = oldIsCTE - b.buildingCTE = oldBuildingCTE }() p, err := b.buildResultSetNode(ctx, cte.Query.Query, true) diff --git a/pkg/planner/core/masking_policy_at_result_test.go b/pkg/planner/core/masking_policy_at_result_test.go index ac8a6bc8bd93c..2381116305c9e 100644 --- a/pkg/planner/core/masking_policy_at_result_test.go +++ b/pkg/planner/core/masking_policy_at_result_test.go @@ -36,7 +36,7 @@ func TestMaskingPolicyAtResult(t *testing.T) { tk.MustExec("create masking policy p_having on t_having(val) as concat('[', val, ']') enable") // HAVING should compare original values, not masked values - result := tk.MustQuery("select val, count(*) from t_having group by val having val > 'apple'") + result := tk.MustQuery("select val, count(*) from t_having group by val having val > 'apple' order by val") result.Check(testkit.Rows("[banana] 1", "[cherry] 1")) // This should return no rows because no value > 'zzzz' From 6c1ac43d454905aa4d069dcd55a13bdc8f0e7791 Mon Sep 17 00:00:00 2001 From: tiancaiamao Date: Wed, 10 Jun 2026 17:25:44 +0800 Subject: [PATCH 06/13] fix: DDL rename column transition state and CurrentDB fallback 1. buildMaskingReplaceExprs: when getMaskingPolicyExpr fails (e.g., DDL rename column transition where expression still references old column name), log warning and skip instead of failing the query. 2. findMaskingPolicy: restored CurrentDB fallback for empty DBName to match reference implementation behavior. --- pkg/planner/core/logical_plan_builder.go | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go index b07daf41af915..96c14ff8c9d69 100644 --- a/pkg/planner/core/logical_plan_builder.go +++ b/pkg/planner/core/logical_plan_builder.go @@ -1565,6 +1565,15 @@ func (b *PlanBuilder) buildMaskingReplaceExprs(ctx context.Context, p base.Logic continue } expr, placeholder, err := getMaskingPolicyExpr(b.ctx.GetExprCtx(), sv, schemaVersion, policy, tblInfo, colInfo) + if err != nil { + // Log and skip this column — this can happen during DDL transitions + // (e.g., column rename where the expression still references the old name). + logutil.BgLogger().Warn("failed to build masking policy expression, skipping", + zap.String("policy", policy.Name.L), + zap.Error(err), + ) + continue + } if err != nil { return nil, err } @@ -1595,8 +1604,9 @@ func (b *PlanBuilder) findMaskingPolicy(_ context.Context, name *types.FieldName } dbName := name.DBName if dbName.L == "" { - // Cannot determine source database (e.g., derived table aliasing clears DBName). - // Skip masking rather than risk applying the wrong policy via CurrentDB fallback. + dbName = ast.NewCIStr(b.ctx.GetSessionVars().CurrentDB) + } + if dbName.L == "" { return nil, nil, nil } tbl, err := b.is.TableByName(context.Background(), dbName, tblName) From 33ae209ce8b5585e9da2e8b71248118a18798b85 Mon Sep 17 00:00:00 2001 From: tiancaiamao Date: Wed, 10 Jun 2026 18:06:04 +0800 Subject: [PATCH 07/13] fix: remove dead err check causing nogo impossible condition lint error --- pkg/planner/core/logical_plan_builder.go | 3 --- 1 file changed, 3 deletions(-) diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go index 96c14ff8c9d69..630a983d1274b 100644 --- a/pkg/planner/core/logical_plan_builder.go +++ b/pkg/planner/core/logical_plan_builder.go @@ -1574,9 +1574,6 @@ func (b *PlanBuilder) buildMaskingReplaceExprs(ctx context.Context, p base.Logic ) continue } - if err != nil { - return nil, err - } if placeholder == nil { continue } From 334098fe3bfe9dfd6b94aa920a365541f32e3d76 Mon Sep 17 00:00:00 2001 From: tiancaiamao Date: Wed, 10 Jun 2026 20:41:05 +0800 Subject: [PATCH 08/13] fix(ddl): add sys-table fallback in getMaskingPolicyByNameForDDL After TRUNCATE TABLE, infoSchema may have stale masking policy data because loadMaskingPoliciesIfNeeded uses a snapshot TS that predates the DDL worker's UPDATE to mysql.tidb_masking_policy. Add a fallback: when infoSchema lookup fails, query the system table directly via ExecRestrictedSQL with snapshotTS=0 (latest data). Renames getMaskingPolicyByNameForTable to getMaskingPolicyByNameForDDL to match the reference implementation. --- pkg/ddl/executor.go | 46 +++++++++++++++++++++++++++++++++++++-------- 1 file changed, 38 insertions(+), 8 deletions(-) diff --git a/pkg/ddl/executor.go b/pkg/ddl/executor.go index e3cd2c0327e6d..51530fdc2b717 100644 --- a/pkg/ddl/executor.go +++ b/pkg/ddl/executor.go @@ -6540,15 +6540,45 @@ func (e *executor) AddMaskingPolicy(ctx sessionctx.Context, ident ast.Ident, spe return e.createMaskingPolicyWithInfo(ctx, tbl.Meta().ID, tbl.Meta().Columns, policyInfo, OnExistError) } -func (e *executor) getMaskingPolicyByNameForTable(tableID int64, cols []*model.ColumnInfo, policyName ast.CIStr) (*model.MaskingPolicyInfo, bool) { +// getMaskingPolicyByNameForDDL finds a masking policy by name. It first checks the +// infoSchema cache, then falls back to querying mysql.tidb_masking_policy directly. +// The fallback is necessary because DDL operations like TRUNCATE TABLE update the +// system table outside the infoSchema's snapshot window. +func (e *executor) getMaskingPolicyByNameForDDL( + ctx sessionctx.Context, + tableID int64, + columns []*model.ColumnInfo, + policyName ast.CIStr, +) (*model.MaskingPolicyInfo, bool) { + // Fast path: check infoSchema cache. is := e.infoCache.GetLatest() - for _, col := range cols { + for _, col := range columns { policy, ok := is.MaskingPolicyByTableColumn(tableID, col.ID) if ok && policy != nil && policy.Name.L == policyName.L { return policy, true } } - return nil, false + // Slow path: query system table directly. This handles cases where the + // infoSchema's delayed loading has stale data after DDL operations. + rows, _, err := ctx.GetRestrictedSQLExecutor().ExecRestrictedSQL( + kv.WithInternalSourceType(context.Background(), kv.InternalTxnDDL), + nil, + `SELECT policy_id, policy_name, db_name, table_name, table_id, column_name, column_id, expression, status, masking_type, restrict_on, created_at, updated_at, created_by +FROM mysql.tidb_masking_policy +WHERE table_id = %? AND LOWER(policy_name) = %? +ORDER BY policy_id +LIMIT 1`, + tableID, + policyName.L, + ) + if err != nil || len(rows) == 0 { + return nil, false + } + policy, err := maskingPolicyFromSysTableRow(rows[0]) + if err != nil { + return nil, false + } + return policy, true } func (e *executor) AlterTableMaskingPolicyState(ctx sessionctx.Context, ident ast.Ident, spec *ast.AlterTableSpec, enabled bool) error { @@ -6557,7 +6587,7 @@ func (e *executor) AlterTableMaskingPolicyState(ctx sessionctx.Context, ident as return errors.Trace(err) } policyName := spec.MaskingPolicyName - policy, ok := e.getMaskingPolicyByNameForTable(tbl.Meta().ID, tbl.Meta().Columns, policyName) + policy, ok := e.getMaskingPolicyByNameForDDL(ctx, tbl.Meta().ID, tbl.Meta().Columns, policyName) if !ok { return errors.Errorf("masking policy %s doesn't exist", policyName.O) } @@ -6603,7 +6633,7 @@ func (e *executor) DropMaskingPolicy(ctx sessionctx.Context, ident ast.Ident, sp return errors.Trace(err) } policyName := spec.MaskingPolicyName - policy, ok := e.getMaskingPolicyByNameForTable(tbl.Meta().ID, tbl.Meta().Columns, policyName) + policy, ok := e.getMaskingPolicyByNameForDDL(ctx, tbl.Meta().ID, tbl.Meta().Columns, policyName) if !ok { return errors.Errorf("masking policy %s doesn't exist", policyName.O) } @@ -6642,7 +6672,7 @@ func (e *executor) ModifyMaskingPolicyExpression(ctx sessionctx.Context, ident a return errors.Trace(err) } policyName := spec.MaskingPolicyName - policy, ok := e.getMaskingPolicyByNameForTable(tbl.Meta().ID, tbl.Meta().Columns, policyName) + policy, ok := e.getMaskingPolicyByNameForDDL(ctx, tbl.Meta().ID, tbl.Meta().Columns, policyName) if !ok { return errors.Errorf("masking policy %s doesn't exist", policyName.O) } @@ -6695,7 +6725,7 @@ func (e *executor) ModifyMaskingPolicyRestrictOn(ctx sessionctx.Context, ident a return errors.Trace(err) } policyName := spec.MaskingPolicyName - policy, ok := e.getMaskingPolicyByNameForTable(tbl.Meta().ID, tbl.Meta().Columns, policyName) + policy, ok := e.getMaskingPolicyByNameForDDL(ctx, tbl.Meta().ID, tbl.Meta().Columns, policyName) if !ok { return errors.Errorf("masking policy %s doesn't exist", policyName.O) } @@ -6739,7 +6769,7 @@ func (e *executor) createMaskingPolicyWithInfo( onExist OnExist, ) error { is := e.infoCache.GetLatest() - if existPolicy, ok := e.getMaskingPolicyByNameForTable(tableID, cols, policy.Name); ok { + if existPolicy, ok := e.getMaskingPolicyByNameForDDL(ctx, tableID, cols, policy.Name); ok { if existPolicy.ColumnID != policy.ColumnID { return errors.Errorf("masking policy %s already exists on another column", existPolicy.Name.O) } From ebf51ab648e0856be51f486b1c981ae854803125 Mon Sep 17 00:00:00 2001 From: tiancaiamao Date: Wed, 10 Jun 2026 21:50:02 +0800 Subject: [PATCH 09/13] make fmt --- pkg/ddl/executor.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/ddl/executor.go b/pkg/ddl/executor.go index 51530fdc2b717..b162c2306b8d0 100644 --- a/pkg/ddl/executor.go +++ b/pkg/ddl/executor.go @@ -6769,7 +6769,7 @@ func (e *executor) createMaskingPolicyWithInfo( onExist OnExist, ) error { is := e.infoCache.GetLatest() - if existPolicy, ok := e.getMaskingPolicyByNameForDDL(ctx, tableID, cols, policy.Name); ok { + if existPolicy, ok := e.getMaskingPolicyByNameForDDL(ctx, tableID, cols, policy.Name); ok { if existPolicy.ColumnID != policy.ColumnID { return errors.Errorf("masking policy %s already exists on another column", existPolicy.Name.O) } From 721a7449d58395fa5a342e9b5574b952da058408 Mon Sep 17 00:00:00 2001 From: tiancaiamao Date: Wed, 10 Jun 2026 22:12:17 +0800 Subject: [PATCH 10/13] fix(planner): clone MaskingExprs in FastClonePointGetForPlanCache The new MaskingExprs field on PointGetPlan was not being copied during fast clone, causing TestFastPointGetClone to fail. --- pkg/planner/core/plan_clone_utils.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pkg/planner/core/plan_clone_utils.go b/pkg/planner/core/plan_clone_utils.go index 2d97d11690bd8..c518ea5157ebc 100644 --- a/pkg/planner/core/plan_clone_utils.go +++ b/pkg/planner/core/plan_clone_utils.go @@ -67,7 +67,8 @@ func FastClonePointGetForPlanCache(newCtx base.PlanContext, src, dst *physicalop dst.Lock = src.Lock dst.SetOutputNames(src.OutputNames()) dst.LockWaitTime = src.LockWaitTime - dst.Columns = src.Columns + dst.Columns = src.Columns + dst.MaskingExprs = src.MaskingExprs // remaining fields are unnecessary to clone: // cost, planCostInit, planCost, planCostVer2, accessCols From 0a5f55635d3b2468f72e1c0a310a84cd1ae83c57 Mon Sep 17 00:00:00 2001 From: tiancaiamao Date: Wed, 10 Jun 2026 22:31:17 +0800 Subject: [PATCH 11/13] make fm --- pkg/planner/core/plan_clone_utils.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/planner/core/plan_clone_utils.go b/pkg/planner/core/plan_clone_utils.go index c518ea5157ebc..eae816d3e0aea 100644 --- a/pkg/planner/core/plan_clone_utils.go +++ b/pkg/planner/core/plan_clone_utils.go @@ -67,7 +67,7 @@ func FastClonePointGetForPlanCache(newCtx base.PlanContext, src, dst *physicalop dst.Lock = src.Lock dst.SetOutputNames(src.OutputNames()) dst.LockWaitTime = src.LockWaitTime - dst.Columns = src.Columns + dst.Columns = src.Columns dst.MaskingExprs = src.MaskingExprs // remaining fields are unnecessary to clone: From c37deb97d2a667d8b76d12fddf81f9b52b366a46 Mon Sep 17 00:00:00 2001 From: tiancaiamao Date: Mon, 15 Jun 2026 18:36:16 +0800 Subject: [PATCH 12/13] address comment --- pkg/planner/core/logical_plan_builder.go | 52 ++++++------------- .../core/masking_policy_at_result_test.go | 20 +++++-- .../physicalop/physical_batch_point_get.go | 6 +++ pkg/planner/core/planbuilder.go | 4 ++ 4 files changed, 43 insertions(+), 39 deletions(-) diff --git a/pkg/planner/core/logical_plan_builder.go b/pkg/planner/core/logical_plan_builder.go index 630a983d1274b..f7cfef00e298f 100644 --- a/pkg/planner/core/logical_plan_builder.go +++ b/pkg/planner/core/logical_plan_builder.go @@ -434,6 +434,11 @@ func (b *PlanBuilder) buildTableRefs(ctx context.Context, from *ast.TableRefsCla func (b *PlanBuilder) buildResultSetNode(ctx context.Context, node ast.ResultSetNode, isCTE bool) (p base.LogicalPlan, err error) { //If it is building the CTE queries, we will mark them. b.isCTE = isCTE + // Mark that we are inside a result-set node so that derived-table subqueries + // do not apply AT RESULT masking (the outer query needs original values). + oldBuildingResultSetNode := b.buildingResultSetNode + b.buildingResultSetNode = true + defer func() { b.buildingResultSetNode = oldBuildingResultSetNode }() switch x := node.(type) { case *ast.Join: return b.buildJoin(ctx, x) @@ -1566,13 +1571,9 @@ func (b *PlanBuilder) buildMaskingReplaceExprs(ctx context.Context, p base.Logic } expr, placeholder, err := getMaskingPolicyExpr(b.ctx.GetExprCtx(), sv, schemaVersion, policy, tblInfo, colInfo) if err != nil { - // Log and skip this column — this can happen during DDL transitions - // (e.g., column rename where the expression still references the old name). - logutil.BgLogger().Warn("failed to build masking policy expression, skipping", - zap.String("policy", policy.Name.L), - zap.Error(err), - ) - continue + // Fail-closed: returning an error is safer than skipping masking, + // which would leak the raw column value to the user. + return nil, errors.Trace(err) } if placeholder == nil { continue @@ -1974,7 +1975,7 @@ func (b *PlanBuilder) implicitProjectGroupingSetCols(projSchema *expression.Sche // buildProjection returns a Projection plan and non-aux columns length. func (b *PlanBuilder) buildProjection(ctx context.Context, p base.LogicalPlan, fields []*ast.SelectField, mapper map[*ast.AggregateFuncExpr]int, - windowMapper map[*ast.WindowFuncExpr]int, considerWindow bool, expandGenerateColumn bool, applyMasking bool) (base.LogicalPlan, []expression.Expression, int, error) { + windowMapper map[*ast.WindowFuncExpr]int, considerWindow bool, expandGenerateColumn bool) (base.LogicalPlan, []expression.Expression, int, error) { err := b.preprocessUserVarTypes(ctx, p, fields, mapper) if err != nil { return nil, nil, 0, err @@ -1985,20 +1986,6 @@ func (b *PlanBuilder) buildProjection(ctx context.Context, p base.LogicalPlan, f schema := expression.NewSchema(make([]*expression.Column, 0, len(fields))...) oldLen := 0 newNames := make([]*types.FieldName, 0, len(fields)) - var cachedMaskPlan base.LogicalPlan - var cachedMaskExprs []expression.Expression - getMaskExprs := func(p base.LogicalPlan) ([]expression.Expression, error) { - if cachedMaskPlan == p { - return cachedMaskExprs, nil - } - exprs, err := b.buildMaskingReplaceExprs(ctx, p) - if err != nil { - return nil, err - } - cachedMaskPlan = p - cachedMaskExprs = exprs - return exprs, nil - } for i, field := range fields { if !field.Auxiliary { oldLen++ @@ -2046,15 +2033,6 @@ func (b *PlanBuilder) buildProjection(ctx context.Context, p base.LogicalPlan, f } p = np - if applyMasking && !field.Auxiliary { - maskExprs, err := getMaskExprs(p) - if err != nil { - return nil, nil, 0, err - } - if maskExprs != nil { - newExpr = expression.ColumnSubstitute(b.ctx.GetExprCtx(), newExpr, p.Schema(), maskExprs) - } - } proj.Exprs = append(proj.Exprs, newExpr) col, name, err := b.buildProjectionField(ctx, p, field, newExpr) @@ -2313,6 +2291,7 @@ func (b *PlanBuilder) buildProjection4Union(_ context.Context, u *logicalop.Logi unionCols = append(unionCols, &expression.Column{ RetType: resultTp, UniqueID: b.ctx.GetSessionVars().AllocPlanColumnID(), + ID: col.ID, }) } u.SetSchema(expression.NewSchema(unionCols...)) @@ -2438,10 +2417,11 @@ func (b *PlanBuilder) buildSetOpr(ctx context.Context, setOpr *ast.SetOprStmt) ( setOprPlan = proj } - if b.buildingSetOprOperands == 1 && !b.isCTE && !b.buildingCTE { + if b.buildingSetOprOperands == 1 && !b.isCTE && !b.buildingCTE && !b.buildingResultSetNode { // Apply masking at the final result stage (AT RESULT semantics). // This ensures set operators (UNION/INTERSECT/EXCEPT) use original values. - // For CTEs, we skip masking here because CTE definitions should preserve original values. + // For CTEs and derived tables, we skip masking here because they should + // preserve original values for correct outer filtering/joining. setOprPlan, err = b.buildFinalProjectionWithMasking(ctx, setOprPlan, oldLen, nil) if err != nil { return nil, err @@ -4775,7 +4755,7 @@ func (b *PlanBuilder) buildSelect(ctx context.Context, sel *ast.SelectStmt) (p b var oldLen int // According to https://dev.mysql.com/doc/refman/8.0/en/window-functions-usage.html, // we can only process window functions after having clause, so `considerWindow` is false now. - p, projExprs, oldLen, err = b.buildProjection(ctx, p, sel.Fields.Fields, totalMap, nil, false, sel.OrderBy != nil, false) + p, projExprs, oldLen, err = b.buildProjection(ctx, p, sel.Fields.Fields, totalMap, nil, false, sel.OrderBy != nil) if err != nil { return nil, err } @@ -4813,7 +4793,7 @@ func (b *PlanBuilder) buildSelect(ctx context.Context, sel *ast.SelectStmt) (p b // In such case plan `p` is not changed, so we don't have to build another projection. if hasWindowFuncField { // Now we build the window function fields. - p, projExprs, oldLen, err = b.buildProjection(ctx, p, sel.Fields.Fields, windowAggMap, windowMapper, true, false, false) + p, projExprs, oldLen, err = b.buildProjection(ctx, p, sel.Fields.Fields, windowAggMap, windowMapper, true, false) if err != nil { return nil, err } @@ -4873,7 +4853,7 @@ func (b *PlanBuilder) buildSelect(ctx context.Context, sel *ast.SelectStmt) (p b // For CTEs, we skip masking here because: // 1. CTE definitions should preserve original values for correct filtering/joining // 2. Masking is applied when CTE results are materialized to the final output - if b.buildingSetOprOperands == 0 && !b.isCTE && !b.buildingCTE { + if b.buildingSetOprOperands == 0 && !b.isCTE && !b.buildingCTE && !b.buildingResultSetNode { p, err = b.buildFinalProjectionWithMasking(ctx, p, oldLen, originalFields) if err != nil { return nil, err diff --git a/pkg/planner/core/masking_policy_at_result_test.go b/pkg/planner/core/masking_policy_at_result_test.go index 2381116305c9e..2df99ee0d41a7 100644 --- a/pkg/planner/core/masking_policy_at_result_test.go +++ b/pkg/planner/core/masking_policy_at_result_test.go @@ -68,10 +68,10 @@ func TestMaskingPolicyAtResult(t *testing.T) { tk.MustExec("create masking policy p_union on t_union(val) as concat('[', val, ']') enable") // Two different original values should NOT be deduplicated - // even though they mask to the same pattern + // even though they mask to the same pattern. Also verify the output + // values are actually masked (not raw). result := tk.MustQuery("select val from t_union where id = 1 union distinct select val from t_union where id = 2") - rows := result.Rows() - require.Len(t, rows, 2) + result.Sort().Check(testkit.Rows("[value1]", "[value2]")) }) // Test aggregate with HAVING @@ -91,4 +91,18 @@ func TestMaskingPolicyAtResult(t *testing.T) { require.Contains(t, []string{"_B", "_C"}, rows[1][0]) require.NotEqual(t, rows[0][0], rows[1][0]) }) + + // Test derived table: outer WHERE uses original values, output is masked + t.Run("derived table outer filter uses original values", func(t *testing.T) { + tk.MustExec("drop table if exists t_derived") + tk.MustExec("create table t_derived(id int, val varchar(100))") + tk.MustExec("insert into t_derived values (1, 'apple'), (2, 'banana'), (3, 'cherry')") + tk.MustExec("create masking policy p_derived on t_derived(val) as concat('[', val, ']') enable") + + // Outer WHERE should compare original values, not masked values. + // If masking were applied inside the derived table, the WHERE clause + // would see '[apple]' > 'a' (false) and return no rows. + result := tk.MustQuery("select val from (select val from t_derived) sub where val > 'apple' order by val") + result.Check(testkit.Rows("[banana]", "[cherry]")) + }) } diff --git a/pkg/planner/core/operator/physicalop/physical_batch_point_get.go b/pkg/planner/core/operator/physicalop/physical_batch_point_get.go index b7d4fe4144590..825e673a25a63 100644 --- a/pkg/planner/core/operator/physicalop/physical_batch_point_get.go +++ b/pkg/planner/core/operator/physicalop/physical_batch_point_get.go @@ -333,6 +333,9 @@ func (p *PointGetPlan) MemoryUsage() (sum int64) { for _, col := range p.accessCols { sum += col.MemoryUsage() } + for _, expr := range p.MaskingExprs { + sum += expr.MemoryUsage() + } return } @@ -762,6 +765,9 @@ func (p *BatchPointGetPlan) MemoryUsage() (sum int64) { for _, col := range p.accessCols { sum += col.MemoryUsage() } + for _, expr := range p.MaskingExprs { + sum += expr.MemoryUsage() + } return } diff --git a/pkg/planner/core/planbuilder.go b/pkg/planner/core/planbuilder.go index 9cfffbcab671e..b074ff4aaffaa 100644 --- a/pkg/planner/core/planbuilder.go +++ b/pkg/planner/core/planbuilder.go @@ -289,6 +289,10 @@ type PlanBuilder struct { // When > 0, final masking projection should be skipped for inner SELECT/SET plans and // applied only once on the completed set-operator result. buildingSetOprOperands int + // buildingResultSetNode is true while building a subquery inside a FROM clause + // (i.e., inside buildResultSetNode). This prevents AT RESULT masking from being + // applied inside derived tables, where outer queries need to see original values. + buildingResultSetNode bool // Check whether the current building query is a CTE isCTE bool // CTE table name in lower case, it can be nil From 0fc9cf0ba5437d4fd3f1b1459d41fda86c871273 Mon Sep 17 00:00:00 2001 From: tiancaiamao Date: Mon, 15 Jun 2026 21:44:34 +0800 Subject: [PATCH 13/13] address comment --- pkg/infoschema/infoschema_v2.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pkg/infoschema/infoschema_v2.go b/pkg/infoschema/infoschema_v2.go index f8ba25d1123a7..feb399f0664f6 100644 --- a/pkg/infoschema/infoschema_v2.go +++ b/pkg/infoschema/infoschema_v2.go @@ -1603,6 +1603,11 @@ func (b *Builder) applyTableUpdateV2(m meta.Reader, diff *model.SchemaDiff) ([]i return nil, errors.Trace(err) } } + if needRefreshMaskingPoliciesForTableDiff(diff.Type) { + if err := refreshMaskingPoliciesForTableIDs(b, oldTableID, newTableID); err != nil { + return nil, errors.Trace(err) + } + } return tblIDs, nil }