test: add userAttributes shorthand in mask.sql test cases

packages/cubejs-schema-compiler/src/adapter/BaseQuery.js

@@ -2846,14 +2846,19 @@ export class BaseQuery {
     return R.pipe(
       R.map(f => f.getMembers()),
       R.flatten,
-      R.map(s => (
-        (cache || this.compilerCache).cache(
+      R.map(s => {
+        const memberPath = s.path() ? s.path().join('.') : null;
+        const hasSqlMask = memberPath &&
+          this.maskedMembers && this.maskedMembers.size > 0 &&
+          this.maskedMembers.has(memberPath) &&
+          s.definition()?.mask && typeof s.definition().mask === 'object' && s.definition().mask.sql;
+        return (cache || (hasSqlMask ? this.queryCache : this.compilerCache)).cache(
           ['collectFrom'].concat(methodCacheKey).concat(
-            s.path() ? [s.path().join('.')] : [s.cube().name, s.expression?.toString() || s.expressionName || s.definition().sql]
+            memberPath ? [memberPath] : [s.cube().name, s.expression?.toString() || s.expressionName || s.definition().sql]
           ),
           () => fn(() => this.traverseSymbol(s))
-        )
-      )),
+        );
+      }),
       R.unnest,
       R.uniq,
       R.filter(R.identity)

packages/cubejs-schema-compiler/test/unit/transpilers.test.ts

@@ -353,6 +353,70 @@ describe('Transpilers', () => {
     expect(transpiledSql!.toString()).toMatch('SECURITY_CONTEXT.cubeCloud.userAttributes');
   });
 
+  it('CubePropContextTranspiler with userAttributes shorthand in mask.sql should transpile to SECURITY_CONTEXT', async () => {
+    const { cubeEvaluator, compiler } = prepareJsCompiler(`
+        cube(\`Test\`, {
+          sql: 'SELECT * FROM users',
+          dimensions: {
+            userId: {
+              sql: \`userId\`,
+              type: 'string'
+            },
+            masked_dim: {
+              sql: \`price\`,
+              type: 'number',
+              mask: {
+                sql: \`CAST(\${userAttributes.tenantId} AS INTEGER)\`,
+              }
+            }
+          }
+        })
+    `);
+
+    await compiler.compile();
+
+    const transpiledMaskSql = (cubeEvaluator.cubeFromPath('Test').dimensions.masked_dim as any).mask.sql;
+    expect(transpiledMaskSql!.toString()).toMatch('SECURITY_CONTEXT.cubeCloud.userAttributes');
+  });
+
+  it('CubePropContextTranspiler mask.sql with CUBE reference should resolve correctly', async () => {
+    const compilers = prepareJsCompiler(`
+        cube(\`Test\`, {
+          sql_table: 'public.test',
+          dimensions: {
+            id: {
+              sql: \`id\`,
+              type: 'number',
+              primary_key: true,
+            },
+            secret: {
+              sql: \`secret_val\`,
+              type: 'string',
+              mask: {
+                sql: \`CONCAT('***', RIGHT(CAST(\${CUBE}.secret_val AS TEXT), 2))\`,
+              }
+            }
+          },
+          measures: {
+            count: { type: 'count' }
+          }
+        })
+    `);
+
+    await compilers.compiler.compile();
+
+    const query = new PostgresQuery(
+      compilers,
+      {
+        measures: ['Test.count'],
+        dimensions: ['Test.secret'],
+        maskedMembers: ['Test.secret'],
+      }
+    );
+    const sql = query.buildSqlAndParams();
+    expect(sql[0]).toContain('"test".secret_val');
+  });
+
   it('CubePropContextTranspiler should not transform groups shorthand when a cube member named groups exists', async () => {
     const { cubeEvaluator, compiler } = prepareJsCompiler(`
         cube(\`Test\`, {

packages/cubejs-schema-compiler/test/unit/yaml-schema.test.ts

@@ -1,4 +1,5 @@
 import { prepareYamlCompiler } from './PrepareCompiler';
+import { PostgresQuery } from '../../src';
 
 describe('Yaml Schema Testing', () => {
   describe('Duplicate member detection', () => {
@@ -1114,4 +1115,122 @@ cubes:
       }
     });
   });
+
+  describe('Mask SQL with shorthand', () => {
+    it('userAttributes shorthand in mask sql should compile and resolve', async () => {
+      const compilers = prepareYamlCompiler(`
+cubes:
+  - name: orders
+    sql_table: public.orders
+    dimensions:
+      - name: id
+        sql: id
+        type: number
+        primary_key: true
+      - name: status
+        sql: status
+        type: string
+        mask:
+          sql: "CASE WHEN { userAttributes.hasStatusAccess } THEN {CUBE}.status ELSE '***' END"
+    measures:
+      - name: count
+        type: count
+    access_policy:
+      - role: "*"
+        member_level:
+          includes: []
+        member_masking:
+          includes: "*"
+      `);
+
+      await compilers.compiler.compile();
+
+      const dim = compilers.cubeEvaluator.cubeFromPath('orders').dimensions.status;
+      const maskSql = (dim as any).mask.sql.toString();
+      expect(maskSql).toContain('SECURITY_CONTEXT.cubeCloud.userAttributes.hasStatusAccess');
+      expect(maskSql).toContain('CUBE');
+      expect(maskSql).not.toMatch(/[^.}]userAttributes\.hasStatusAccess/);
+
+      const query = new PostgresQuery(
+        compilers,
+        {
+          measures: ['orders.count'],
+          dimensions: ['orders.status'],
+          maskedMembers: ['orders.status'],
+          contextSymbols: {
+            securityContext: { cubeCloud: { userAttributes: { hasStatusAccess: true } } }
+          }
+        }
+      );
+      const sql = query.buildSqlAndParams();
+      expect(sql[0]).toContain('"orders".status');
+      expect(sql[0]).toContain('CASE WHEN');
+    });
+
+    it('user_attributes shorthand in mask sql should compile and resolve', async () => {
+      const compilers = prepareYamlCompiler(`
+cubes:
+  - name: orders
+    sql_table: public.orders
+    dimensions:
+      - name: id
+        sql: id
+        type: number
+        primary_key: true
+      - name: status
+        sql: status
+        type: string
+        mask:
+          sql: "CASE WHEN { user_attributes.hasStatusAccess } THEN {CUBE}.status ELSE '***' END"
+    measures:
+      - name: count
+        type: count
+    access_policy:
+      - role: "*"
+        member_level:
+          includes: []
+        member_masking:
+          includes: "*"
+      `);
+
+      await compilers.compiler.compile();
+
+      const dim = compilers.cubeEvaluator.cubeFromPath('orders').dimensions.status;
+      const maskSql = (dim as any).mask.sql.toString();
+      expect(maskSql).toContain('SECURITY_CONTEXT.cubeCloud.userAttributes.hasStatusAccess');
+    });
+
+    it('groups shorthand in mask sql should compile and resolve', async () => {
+      const compilers = prepareYamlCompiler(`
+cubes:
+  - name: orders
+    sql_table: public.orders
+    dimensions:
+      - name: id
+        sql: id
+        type: number
+        primary_key: true
+      - name: secret
+        sql: price
+        type: number
+        mask:
+          sql: "CASE WHEN {CUBE}.product_id IN ({groups}) THEN {CUBE}.price ELSE -1 END"
+    measures:
+      - name: count
+        type: count
+    access_policy:
+      - role: "*"
+        member_level:
+          includes: []
+        member_masking:
+          includes: "*"
+      `);
+
+      await compilers.compiler.compile();
+
+      const dim = compilers.cubeEvaluator.cubeFromPath('orders').dimensions.secret;
+      const maskSql = (dim as any).mask.sql.toString();
+      expect(maskSql).toContain('SECURITY_CONTEXT.cubeCloud.groups');
+    });
+  });
 });

packages/cubejs-testing/birdbox-fixtures/rbac/model/cubes/masking_test.yaml

@@ -106,6 +106,32 @@ cubes:
         member_level:
           includes: []
 
+  - name: yaml_ua_mask_test
+    sql_table: public.line_items
+
+    dimensions:
+      - name: id
+        sql: id
+        type: number
+        primary_key: true
+
+      - name: masked_status
+        sql: product_id
+        type: number
+        mask:
+          sql: "CASE WHEN { userAttributes.tenantId } = '1' THEN {CUBE}.product_id ELSE -1 END"
+
+    measures:
+      - name: count
+        type: count
+
+    access_policy:
+      - role: "*"
+        member_level:
+          includes: []
+        member_masking:
+          includes: "*"
+
 views:
   # View with full access at view level - but cube masking still applies (RLS pattern)
   # Excludes members with {CUBE} references in SQL (secret_string, secret_boolean)

packages/cubejs-testing/birdbox-fixtures/rbac/model/cubes/security_context_test.js

@@ -66,6 +66,132 @@ cube('sc_interpolation_test', {
   },
 });
 
+cube('sc_ua_mask_test', {
+  sql_table: 'public.line_items',
+
+  dimensions: {
+    id: {
+      sql: 'id',
+      type: 'number',
+      primary_key: true,
+    },
+    product_id: {
+      sql: 'product_id',
+      type: 'number',
+    },
+    masked_price: {
+      sql: 'price',
+      type: 'number',
+      mask: {
+        sql: `CAST(${userAttributes.tenantId} AS INTEGER)`,
+      },
+    },
+  },
+
+  measures: {
+    count: {
+      type: 'count',
+    },
+  },
+
+  accessPolicy: [
+    {
+      role: '*',
+      memberLevel: {
+        includes: [],
+      },
+      memberMasking: {
+        includes: '*',
+      },
+    },
+  ],
+});
+
+cube('sc_cube_mask_test', {
+  sql_table: 'public.line_items',
+
+  dimensions: {
+    id: {
+      sql: 'id',
+      type: 'number',
+      primary_key: true,
+    },
+    masked_product: {
+      sql: `${CUBE}.product_id`,
+      type: 'number',
+      mask: {
+        sql: `${CUBE}.product_id * -1`,
+      },
+    },
+  },
+
+  measures: {
+    count: {
+      type: 'count',
+    },
+  },
+
+  accessPolicy: [
+    {
+      role: '*',
+      memberLevel: {
+        includes: [],
+      },
+      memberMasking: {
+        includes: '*',
+      },
+    },
+  ],
+});
+
+cube('sc_joined_mask_test', {
+  sql_table: 'public.line_items',
+
+  joins: {
+    orders: {
+      relationship: 'many_to_one',
+      sql: `${orders}.id = ${sc_joined_mask_test}.order_id`,
+    },
+  },
+
+  dimensions: {
+    id: {
+      sql: 'id',
+      type: 'number',
+      primary_key: true,
+    },
+    order_id: {
+      sql: `${CUBE}.order_id`,
+      type: 'number',
+    },
+    masked_order_id: {
+      sql: `${CUBE}.order_id`,
+      type: 'number',
+      mask: {
+        sql: `${orders.id}`,
+      },
+    },
+  },
+
+  measures: {
+    count: {
+      type: 'count',
+    },
+  },
+
+  accessPolicy: [
+    {
+      role: '*',
+      memberLevel: {
+        includes: [],
+      },
+      memberMasking: {
+        includes: '*',
+      },
+    },
+  ],
+});
+
 cube('sc_groups_shorthand_test', {
   sql_table: 'public.line_items',