feat(query): masking policy USING clause with multi-column support (#18884)

* feat(query): masking policy USING clause with multi-column support

Enable USING clause to reference multiple table columns instead of only
the masked column itself.

Move masking policy application from storage
read layer to SQL binder for unified handling across SELECT/WHERE/HAVING.

Core changes:

1.   Multi-column USING support with parameter substitution
- Add clone_expr_with_replacement() to walk AST and replace column refs
- Map policy parameters to USING columns by position
- Support N policy args mapped to N USING columns

2.   Move masking from read_plan to type checker (binder layer)
- Apply masking in TypeChecker::resolve() for all query paths
- Add in_masking_policy flag to prevent infinite recursion

3.   Meta protobuf schema fixes (v155)
- args: BTreeMap→Vec to preserve parameter order (was losing order)
- Only write first USING column to column_mask_policy_columns_ids.key

* fix

* debug

* optimize

* try fix

Author: TCeason

src/meta/api/src/schema_api_test_suite.rs

@@ -3238,7 +3238,6 @@ impl SchemaApiTestSuite {
                 tenant: tenant.clone(),
                 seq: MatchSeq::Exact(res.ident.seq),
                 table_id,
-                column: "number".to_string(),
                 action: SetSecurityPolicyAction::Set(mask1_id, vec![column_id]),
             };
             mt.set_table_column_mask_policy(req).await?;
@@ -3299,7 +3298,6 @@ impl SchemaApiTestSuite {
                 tenant: tenant.clone(),
                 seq: MatchSeq::Exact(res.ident.seq),
                 table_id,
-                column: "number".to_string(),
                 action: SetSecurityPolicyAction::Set(mask1_id, vec![column_id]),
             };
             mt.set_table_column_mask_policy(req).await?;
@@ -3353,12 +3351,10 @@ impl SchemaApiTestSuite {
                 .find(|f| f.name == "c1")
                 .unwrap()
                 .column_id;
-            println!("fields is {:?}", res.meta.schema.fields);
             let req = SetTableColumnMaskPolicyReq {
                 tenant: tenant.clone(),
                 seq: MatchSeq::Exact(res.ident.seq),
                 table_id: table_id_1,
-                column: "c1".to_string(),
                 action: SetSecurityPolicyAction::Set(mask2_id, vec![column_id]),
             };
             mt.set_table_column_mask_policy(req).await?;
@@ -3410,7 +3406,6 @@ impl SchemaApiTestSuite {
                 tenant: tenant.clone(),
                 seq: MatchSeq::Exact(res.ident.seq),
                 table_id: table_id_1,
-                column: "number".to_string(),
                 action: SetSecurityPolicyAction::Unset(mask2_id),
             };
             mt.set_table_column_mask_policy(req).await?;

src/meta/api/src/security_api.rs

@@ -113,7 +113,7 @@ where
                         columns_ids: column_ids.clone(),
                     };
 
-                    for column_id in column_ids {
+                    if let Some(column_id) = column_ids.first() {
                         new_table_meta
                             .column_mask_policy_columns_ids
                             .insert(*column_id, policy_map.clone());

src/meta/app/src/schema/table/mod.rs

@@ -173,6 +173,8 @@ pub struct TableMeta {
     pub shared_by: BTreeSet<u64>,
     // should be discard
     pub column_mask_policy: Option<BTreeMap<String, String>>,
+    // ColumnId always equals the first value in SecurityPolicyColumnMap::columns_ids
+    // One table can have multiple masking policies
     pub column_mask_policy_columns_ids: BTreeMap<ColumnId, SecurityPolicyColumnMap>,
     // One table only has an unique row access policy
     // should be discard
@@ -794,7 +796,6 @@ pub struct SetTableColumnMaskPolicyReq {
     pub tenant: Tenant,
     pub table_id: u64,
     pub seq: MatchSeq,
-    pub column: String,
     pub action: SetSecurityPolicyAction,
 }
 

src/meta/proto-conv/src/data_mask_from_to_protobuf_impl.rs

@@ -36,12 +36,20 @@ impl FromToProto for mt::DatamaskMeta {
     fn from_pb(p: pb::DatamaskMeta) -> Result<Self, Incompatible> {
         reader_check_msg(p.ver, p.min_reader_ver)?;
 
+        // Prioritize args_v2 (preserves order), fallback to args (backward compatibility)
+        let args: Vec<(String, String)> = if !p.args_v2.is_empty() {
+            p.args_v2
+                .into_iter()
+                .map(|arg| (arg.name, arg.r#type))
+                .collect()
+        } else {
+            // Backward compatibility: read from old args map
+            // Note: BTreeMap sorts keys alphabetically, so order may be lost for old data
+            p.args.into_iter().collect()
+        };
+
         let v = Self {
-            args: p
-                .args
-                .iter()
-                .map(|(arg_name, arg_type)| (arg_name.clone(), arg_type.clone()))
-                .collect::<Vec<_>>(),
+            args,
             return_type: p.return_type,
             body: p.body,
             comment: p.comment.clone(),
@@ -55,14 +63,21 @@ impl FromToProto for mt::DatamaskMeta {
     }
 
     fn to_pb(&self) -> Result<pb::DatamaskMeta, Incompatible> {
-        let mut args = BTreeMap::new();
-        for (arg_name, arg_type) in &self.args {
-            args.insert(arg_name.to_string(), arg_type.to_string());
-        }
+        // Write to args_v2 (new format that preserves order)
+        let args_v2: Vec<pb::DataMaskArg> = self
+            .args
+            .iter()
+            .map(|(arg_name, arg_type)| pb::DataMaskArg {
+                name: arg_name.clone(),
+                r#type: arg_type.clone(),
+            })
+            .collect();
+
         let p = pb::DatamaskMeta {
             ver: VER,
             min_reader_ver: MIN_READER_VER,
-            args,
+            args: BTreeMap::new(),
+            args_v2,
             return_type: self.return_type.clone(),
             body: self.body.clone(),
             comment: self.comment.clone(),

src/meta/proto-conv/src/util.rs

@@ -185,6 +185,7 @@ const META_CHANGE_LOG: &[(u64, &str)] = &[
     (153, "2025-10-16: Add: RowAccessPolicyColumnMap rename to SecurityColumnMap store mask/row_access policy name and column id"),
     (154, "2025-10-16: Add: VacuumWatermark"),
     (155, "2025-10-24: Add: RowAccessPolicyMeta::RowAccessPolicyArg"),
+    (156, "2025-10-22: Add: DataMaskMeta add DataMaskArg"),
     // Dear developer:
     //      If you're gonna add a new metadata version, you'll have to add a test for it.
     //      You could just copy an existing test file(e.g., `../tests/it/v024_table_meta.rs`)

src/meta/proto-conv/tests/it/main.rs

@@ -147,3 +147,4 @@ mod v152_external_udf;
 mod v153_security_column_map;
 mod v154_vacuum_watermark;
 mod v155_row_access_policy_args;
+mod v156_data_mask_args;

src/meta/proto-conv/tests/it/v156_data_mask_args.rs

@@ -0,0 +1,59 @@
+// Copyright 2023 Datafuse Labs.
+//
+// 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.
+
+use chrono::TimeZone;
+use chrono::Utc;
+use fastrace::func_name;
+
+use crate::common;
+
+// These bytes are built when a new version in introduced,
+// and are kept for backward compatibility test.
+//
+// *************************************************************
+// * These messages should never be updated,                   *
+// * only be added when a new version is added,                *
+// * or be removed when an old version is no longer supported. *
+// *************************************************************
+//
+// The message bytes are built from the output of `test_pb_from_to()`
+#[test]
+fn test_decode_v156_data_mask() -> anyhow::Result<()> {
+    let datamask_meta_v156 = vec![
+        18, 6, 83, 116, 114, 105, 110, 103, 26, 68, 67, 65, 83, 69, 32, 87, 72, 69, 78, 32, 99,
+        117, 114, 114, 101, 110, 116, 95, 114, 111, 108, 101, 40, 41, 32, 73, 78, 40, 39, 65, 78,
+        65, 76, 89, 83, 84, 39, 41, 32, 84, 72, 69, 78, 32, 86, 65, 76, 32, 69, 76, 83, 69, 32, 39,
+        42, 42, 42, 42, 42, 42, 42, 42, 42, 39, 32, 69, 78, 68, 34, 12, 115, 111, 109, 101, 32, 99,
+        111, 109, 109, 101, 110, 116, 42, 23, 50, 48, 49, 52, 45, 49, 49, 45, 50, 56, 32, 49, 50,
+        58, 48, 48, 58, 48, 57, 32, 85, 84, 67, 50, 23, 50, 48, 49, 52, 45, 49, 49, 45, 50, 56, 32,
+        49, 50, 58, 48, 48, 58, 48, 57, 32, 85, 84, 67, 58, 14, 10, 4, 116, 101, 115, 116, 18, 6,
+        83, 116, 114, 105, 110, 103, 58, 11, 10, 1, 97, 18, 6, 83, 116, 114, 105, 110, 103, 160, 6,
+        156, 1, 168, 6, 24,
+    ];
+
+    let want = || databend_common_meta_app::data_mask::DatamaskMeta {
+        args: vec![
+            ("test".to_string(), "String".to_string()),
+            ("a".to_string(), "String".to_string()),
+        ],
+        return_type: "String".to_string(),
+        body: "CASE WHEN current_role() IN('ANALYST') THEN VAL ELSE '*********' END".to_string(),
+        comment: Some("some comment".to_string()),
+        create_on: Utc.with_ymd_and_hms(2014, 11, 28, 12, 0, 9).unwrap(),
+        update_on: Some(Utc.with_ymd_and_hms(2014, 11, 28, 12, 0, 9).unwrap()),
+    };
+
+    common::test_pb_from_to(func_name!(), want())?;
+    common::test_load_old(func_name!(), datamask_meta_v156.as_slice(), 156, want())
+}

src/meta/protos/proto/data_mask.proto

@@ -16,16 +16,25 @@ syntax = "proto3";
 
 package databend_proto;
 
+message DataMaskArg {
+  string name = 1;
+  string type = 2;
+}
+
 message DatamaskMeta {
   uint64 ver = 100;
   uint64 min_reader_ver = 101;
 
+  // Deprecated: preserved for backward compatibility, use args_v2 instead
   map<string, string> args = 1;
   string return_type = 2;
   string body = 3;
   optional string comment = 4;
   string create_on = 5;
   optional string update_on = 6;
+
+  // New field that preserves argument order
+  repeated DataMaskArg args_v2 = 7;
 }
 
 message MaskpolicyTableIdList {

src/query/ast/src/ast/statements/table.rs

@@ -1099,9 +1099,20 @@ impl Display for ColumnComment {
 
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub enum ModifyColumnAction {
-    // (column name id, masking policy name)
-    SetMaskingPolicy(Identifier, String),
-    // column name id
+    // Set masking policy on a column with optional USING clause
+    // Syntax: ALTER TABLE ... MODIFY COLUMN <column> SET MASKING POLICY <policy_name> [USING (<columns>)]
+    //
+    // Fields:
+    // - Identifier: The column name to apply the masking policy to
+    // - String: The masking policy name
+    // - Option<Vec<Identifier>>: Optional USING clause columns
+    //
+    // USING clause behavior (follows Snowflake semantics):
+    // - First column: The column to mask/tokenize (must match the target column)
+    // - Additional columns: Condition columns to evaluate for conditional masking
+    // - If USING is omitted: Treats as normal (non-conditional) masking policy
+    SetMaskingPolicy(Identifier, String, Option<Vec<Identifier>>),
+    // Unset masking policy from a column
     UnsetMaskingPolicy(Identifier),
     // vec<ColumnDefinition>
     SetDataType(Vec<ColumnDefinition>),
@@ -1114,8 +1125,14 @@ pub enum ModifyColumnAction {
 impl Display for ModifyColumnAction {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
         match &self {
-            ModifyColumnAction::SetMaskingPolicy(column, name) => {
-                write!(f, "{} SET MASKING POLICY {}", column, name)?
+            ModifyColumnAction::SetMaskingPolicy(column, name, using_columns) => {
+                if let Some(using_columns) = using_columns {
+                    write!(f, "{} SET MASKING POLICY {} USING (", column, name)?;
+                    write_comma_separated_list(f, using_columns)?;
+                    write!(f, ")")?
+                } else {
+                    write!(f, "{} SET MASKING POLICY {}", column, name)?
+                }
             }
             ModifyColumnAction::UnsetMaskingPolicy(column) => {
                 write!(f, "{} UNSET MASKING POLICY", column)?

src/query/ast/src/parser/data_mask.rs

@@ -33,25 +33,10 @@ fn data_mask_arg(i: Input) -> IResult<DataMaskArg> {
     })(i)
 }
 
-fn data_mask_more_arg(i: Input) -> IResult<DataMaskArg> {
-    map(rule! { "," ~ #data_mask_arg }, |(_, arg)| arg)(i)
-}
-
-fn data_mask_arg_list(i: Input) -> IResult<Vec<DataMaskArg>> {
-    map(rule! { #data_mask_more_arg* }, |args| args)(i)
-}
-
 fn data_mask_args(i: Input) -> IResult<Vec<DataMaskArg>> {
     map(
-        rule! { AS ~ "(" ~ #data_mask_arg ~ #data_mask_arg_list ~ ")" },
-        |(_, _, arg_0, more_args, _)| {
-            let mut args = vec![];
-            args.push(arg_0);
-            for arg in more_args {
-                args.push(arg);
-            }
-            args
-        },
+        rule! { AS ~ "(" ~ #comma_separated_list1(data_mask_arg) ~ ")" },
+        |(_, _, args, _)| args,
     )(i)
 }