planner,sessionctx: apply masking in projection

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)
 		}

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
 }
 

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()

pkg/executor/point_get.go

@@ -145,6 +145,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
 }
 
@@ -214,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()
@@ -448,6 +452,44 @@ 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 {
+	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
+	}
+	maskedChunk := chunk.NewChunkWithCapacity(fieldTypes, req.NumRows())
+	evalCtx := sctx.GetExprCtx().GetEvalCtx()
+	for rowIdx := range req.NumRows() {
+		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
 }
 

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
 }
 

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",