TEIID-5443 fix for creating an improper view when adding masking

engine/src/main/java/org/teiid/query/optimizer/relational/ColumnMaskingHelper.java

@@ -46,6 +46,7 @@
 import org.teiid.query.sql.navigator.PreOrPostOrderNavigator;
 import org.teiid.query.sql.symbol.ElementSymbol;
 import org.teiid.query.sql.symbol.Expression;
+import org.teiid.query.sql.symbol.ExpressionSymbol;
 import org.teiid.query.sql.symbol.GroupSymbol;
 import org.teiid.query.sql.symbol.SearchedCaseExpression;
 import org.teiid.query.sql.visitor.ExpressionMappingVisitor;
@@ -75,6 +76,9 @@ public int compareTo(WhenThen arg0) {
 		}
 	}
 
+	/**
+	 * Create a masked version of the col.  Will return either the col if no masking or an {@link ExpressionSymbol} with the masked expression
+	 */
 	private static Expression maskColumn(ElementSymbol col, GroupSymbol unaliased, QueryMetadataInterface metadata, ExpressionMappingVisitor emv, Map<String, DataPolicy> policies, CommandContext cc) throws TeiidComponentException, TeiidProcessingException {
 		Object metadataID = col.getMetadataID();
 		String fullName = metadata.getFullName(metadataID);
@@ -152,7 +156,7 @@ private static Expression maskColumn(ElementSymbol col, GroupSymbol unaliased, Q
 		sce.setElseExpression(col);
 		sce.setType(expectedType);
 		Expression mask = QueryRewriter.rewriteExpression(sce, cc, metadata, true);
-		return mask;
+		return new ExpressionSymbol(col.getShortName(), mask);
 	}
 
 	public static List<? extends Expression> maskColumns(List<ElementSymbol> cols,

engine/src/main/java/org/teiid/query/optimizer/relational/rules/RuleApplySecurity.java

@@ -103,7 +103,7 @@ public PlanNode execute(PlanNode plan, QueryMetadataInterface metadata,
 		        	if (mapping == null) {
 		        		mapping = new HashMap<ElementSymbol, Expression>();
 		        	}
-		        	mapping.put(cols.get(i), maskedCol);
+		        	mapping.put(cols.get(i), SymbolMap.getExpression(maskedCol));
 		        }
         	    PlanNode parentJoin = NodeEditor.findParent(sourceNode.getParent(), NodeConstants.Types.JOIN, NodeConstants.Types.SOURCE);
 		        if (mapping != null) {
@@ -131,7 +131,7 @@ public PlanNode execute(PlanNode plan, QueryMetadataInterface metadata,
 			            		root = sourceNode.getParent();
 			            	}
 	        				root.addAsParent(project);
-	        				addView(metadata, context, group, cols, masked, project);
+	        				addView(metadata, context, group, cols, project);
 	        				parentJoin = null;
 	        			}
 		        	}
@@ -164,7 +164,7 @@ public PlanNode execute(PlanNode plan, QueryMetadataInterface metadata,
         	    	PlanNode project = RelationalPlanner.createProjectNode(cols);
         	    	parent.addAsParent(project);
         	    	//a view is needed to keep the logical placement of the criteria
-    				addView(metadata, context, group, cols, cols, project);
+    				addView(metadata, context, group, cols, project);
         	    }
 			}
 		} catch (TeiidProcessingException e) {
@@ -175,13 +175,12 @@ public PlanNode execute(PlanNode plan, QueryMetadataInterface metadata,
 
 	private void addView(QueryMetadataInterface metadata,
 			CommandContext context, GroupSymbol group,
-			List<ElementSymbol> cols, List<? extends Expression> old,
-			PlanNode viewRoot) throws TeiidComponentException,
+			List<ElementSymbol> cols, PlanNode viewRoot) throws TeiidComponentException,
 			QueryMetadataException, QueryPlannerException {
 		GroupSymbol securityVeiw = new GroupSymbol("sec"); //$NON-NLS-1$
 		Set<String> groups = context.getGroups();
 		securityVeiw = RulePlaceAccess.recontextSymbol(securityVeiw, groups);
-		List<ElementSymbol> newCols = RulePushAggregates.defineNewGroup(securityVeiw, old, metadata);
+		List<ElementSymbol> newCols = RulePushAggregates.defineNewGroup(securityVeiw, cols, metadata);
 		PlanNode newSourceNode = RuleDecomposeJoin.createSource(securityVeiw, viewRoot, newCols);
 		Map<ElementSymbol, Expression> upperMapping = SymbolMap.createSymbolMap(cols, newCols).asMap();
 		FrameUtil.convertFrame(newSourceNode.getParent(), group, Collections.singleton(securityVeiw), upperMapping, metadata);

engine/src/test/java/org/teiid/query/processor/TestColumnMasking.java

@@ -59,7 +59,7 @@ private static CommandContext createContext() {
 		pmd1.setResourceName("pm1.g1.e2");
 		pmd1.setMask("case when e1 = 'a' then null else e2 end");
 
-		policy.addPermission(pmd, pmd1);
+        policy.addPermission(pmd, pmd1);
 		policy.setName("some-role");
 		policies.put("some-role", policy);
 
@@ -163,6 +163,23 @@ private static CommandContext createContext() {
 		helpProcess(plan, context, dataManager, expectedResults);
 	}
 
+	@Test public void testMultipleTableMaskWithPreservedProjection() throws Exception {
+	    DataPolicyMetadata policy1 = new DataPolicyMetadata();
+        PermissionMetaData pmd11 = new PermissionMetaData();
+        pmd11.setResourceName("pm1.g1.e3");
+        pmd11.setMask("rand() > .5");
+
+        policy1.addPermission(pmd11);
+        policy1.setName("other-role");
+        context.getAllowedDataPolicies().put("other-role", policy1);
+
+        HardcodedDataManager dataManager = new HardcodedDataManager();
+        dataManager.addData("SELECT pm1.g1.e1, pm1.g1.e2 FROM pm1.g1", new List<?>[] {Arrays.asList("a", 1), Arrays.asList("b", 2)});
+        ProcessorPlan plan = helpGetPlan(helpParse("select e2 from pm1.g1"), RealMetadataFactory.example1Cached(), new DefaultCapabilitiesFinder(), context);
+        List<?>[] expectedResults = new List<?>[] {Collections.singletonList(null), Arrays.asList(2)};
+        helpProcess(plan, context, dataManager, expectedResults);
+    }
+
 	@Test public void testColumnSubstitution() throws Exception {
 		DataPolicyMetadata policy1 = new DataPolicyMetadata();
 		PermissionMetaData pmd11 = new PermissionMetaData();