Fields that are polymorphic fail when the chosen sObjectType is not the type specified in the selector.

Add the ability to set the sObjectType when validating the field exists on the object.  Use the type set in the selector to determine which polymorphic type to use.

Example: Lead.Owner.UserRoleID
Owner is GROUP | USER
Group does not contain UserRoleID and fails when the getField call sets the sObjectField token to null.

Author: tstone

fflib/src/classes/fflib_QueryFactory.cls

@@ -84,7 +84,7 @@ public class fflib_QueryFactory { //No explicit sharing declaration - inherit fr
 	private Schema.ChildRelationship relationship;
 	private Map<Schema.ChildRelationship, fflib_QueryFactory> subselectQueryMap;
 
-	private QueryField getFieldToken(String fieldName){
+	private QueryField getFieldToken(String fieldName, Schema.sObjectType relatedSObjectType){
 		QueryField result;
 		if(!fieldName.contains('.')){ //single field
 			Schema.SObjectField token = fflib_SObjectDescribe.getDescribe(table).getField(fieldName.toLowerCase());
@@ -103,7 +103,25 @@ public class fflib_QueryFactory { //No explicit sharing declaration - inherit fr
 				if (token != null && enforceFLS) 
 					fflib_SecurityUtils.checkFieldIsReadable(lastSObjectType, token);
 				if(token != null && i.hasNext() && token.getDescribe().getSOAPType() == Schema.SOAPType.ID){
-					lastSObjectType = token.getDescribe().getReferenceTo()[0]; //if it's polymorphic doesn't matter which one we get
+					if(relatedSObjectType == null) {
+						lastSObjectType = token.getDescribe().getReferenceTo()[0]; //if it's polymorphic get the first one - user did not specify the one to use.
+					}else{
+						lastSObjectType = null;
+						List<Schema.sObjectType> relatedObjs = token.getDescribe().getReferenceTo(); //if it's polymorphic, it matters which one we use - i.e. Lead.Owner is GROUP|USER and each has different fields.
+						if(relatedObjs.size() > 1) {
+							String relatedSObjectName = fflib_SObjectDescribe.getDescribe(relatedSObjectType).getDescribe().getName();
+							for(Schema.sObjectType sot : relatedObjs) {
+								if(fflib_SObjectDescribe.getDescribe(sot).getDescribe().getName() == relatedSObjectName){
+									lastSObjectType = sot;
+									break;
+								}
+							}
+							if(lastSObjectType == null)
+								throw new InvalidRelationshipException(fieldName, relatedSObjectType);
+						}else{
+							lastSObjectType = relatedOjs[0]; //only one type returned, use it.
+						}
+					}
 					fieldPath.add(token);
 				}else if(token != null && !i.hasNext()){
 					fieldPath.add(token);
@@ -118,6 +136,9 @@ public class fflib_QueryFactory { //No explicit sharing declaration - inherit fr
 		}
 		return result;
 	}
+	private QueryField getFieldToken(String fieldName) {
+		return this.getFieldToken(fieldName, null);
+	}
 
 	/**
 	 * fflib_QueryFactory instances will be considered equal if they produce the same SOQL query.
@@ -192,9 +213,19 @@ public class fflib_QueryFactory { //No explicit sharing declaration - inherit fr
 	 * @param fieldName the API name of the field to add to the query's SELECT clause.
 	 **/
 	public fflib_QueryFactory selectField(String fieldName){ 		
-		fields.add( getFieldToken(fieldName) );
+		fields.add( getFieldToken(fieldName, null) );
 		return this;
 	} 
+    /**
+     * Selects a single field from the SObject specified in {@link #table}.
+     * Selecting fields is idempotent, if this field is already selected calling this method will have no additional impact.
+     * @param fieldName the API name of the field to add to the query's SELECT clause.
+	 * @param relatedSObjectType the related sObjectType to resolve polymorphic object fields.
+     **/
+    public fflib_QueryFactory selectField(String fieldName, Schema.sOBjectType relatedObjectType){        
+        fields.add( getFieldToken(fieldName, relatedObjectType) );
+        return this;
+    }
 	/**
 	 * Selects a field, avoiding the possible ambiguitiy of String API names.
 	 * @see #selectField(String)
@@ -772,7 +803,16 @@ public class fflib_QueryFactory { //No explicit sharing declaration - inherit fr
 			this.setMessage( 'Invalid field \''+fieldName+'\' for object \''+objectType+'\'' );
 		}
 	}
+	public class InvalidRelationshipException extends Exception{
+		private String fieldName;
+        private Schema.SObjectType objectType;
+        public InvalidRelationshipException(String fieldname, Schema.SObjectType objectType) {
+			this.fieldname = fieldname;
+            this.objectType = objectType;
+            this.setMessage( 'Not able to find related sObjectType for object \''+objectType+'\' denoted by field \''+field+'\'' );
+        }
+    }
 	public class InvalidFieldSetException extends Exception{}
 	public class NonReferenceFieldException extends Exception{}
-	public class InvalidSubqueryRelationshipException extends Exception{}	
+	public class InvalidSubqueryRelationshipException extends Exception{}
 }

fflib/src/classes/fflib_SObjectSelector.cls

@@ -330,7 +330,7 @@ public abstract with sharing class fflib_SObjectSelector
 	{
 		// Add fields from selector prefixing the relationship path		
 		for(SObjectField field : getSObjectFieldList())		
-        	queryFactory.selectField(relationshipFieldPath + '.' + field.getDescribe().getName());
+        	queryFactory.selectField(relationshipFieldPath + '.' + field.getDescribe().getName(), this.getSObjectType());
         // Automatically select the CurrencyIsoCode for MC orgs (unless the object is a known exception to the rule)
         if(Userinfo.isMultiCurrencyOrganization() && CURRENCY_ISO_CODE_ENABLED)
             queryFactory.selectField(relationshipFieldPath+'.CurrencyIsoCode');		

fflib/src/classes/fflib_SObjectSelectorTest.cls

@@ -291,5 +291,97 @@ private with sharing class fflib_SObjectSelectorTest
 			return null;
 		}		
 		return testUser;
+	}
+	
+	@isTest
+	static void testPolyMorphicSelectWithRelatedType()
+	{
+		//Given		
+		
+		Testfflib_CampaignMemberSelector cmSelector = new Testfflib_CampaignMemberSelector();
+		fflib_QueryFactory qf = cmSelector.newQueryFactory();
+		new Testfflib_LeadSelector().configureQueryFactoryFields(qf, 'Lead');
+        new Testfflib_UserSelector().configureQueryFactoryFields(qf, 'Lead.Owner');
+
+		
+		Set<String> expectedSelectFields = new Set<String>{ 'Id', 'Status', 'Lead.Id', 'Lead.OwnerId', 'Lead.Owner.Id', 'Lead.Owner.UserRoleId' };
+		if (UserInfo.isMultiCurrencyOrganization())
+		{
+			expectedSelectFields.add('CurrencyIsoCode');
+		}
+
+		//When
+		String soql = qf.toSOQL();
+
+		//Then
+		Pattern soqlPattern = Pattern.compile('SELECT (.*) FROM CampaignMember ORDER BY CreatedDate ASC NULLS FIRST ');
+		Matcher soqlMatcher = soqlPattern.matcher(soql);
+		soqlMatcher.matches();
+		
+		List<String> actualSelectFields = soqlMatcher.group(1).deleteWhiteSpace().split(',');
+		System.assertEquals(expectedSelectFields, new Set<String>(actualSelectFields));
+	}
+	
+	private class Testfflib_CampaignMemberSelector extends fflib_SObjectSelector
+	{
+		public Testfflib_CampaignMemberSelector()
+		{
+			super();
+		}
+		
+		public List<Schema.SObjectField> getSObjectFieldList()
+		{
+			return new List<Schema.SObjectField> {
+				CampaignMember.Id,
+				CampaignMember.Status
+			};
+		}
+		
+		public Schema.SObjectType getSObjectType()
+		{
+			return CampaignMember.sObjectType;
+		}
+	}
+	
+	private class Testfflib_UserSelector extends fflib_SObjectSelector
+	{
+		public Testfflib_UserSelector()
+		{
+			super();
+		}
+		
+		public List<Schema.SObjectField> getSObjectFieldList()
+		{
+			return new List<Schema.SObjectField> {
+				User.UserRoleId,
+				User.Id
+			};
+		}
+		
+		public Schema.SObjectType getSObjectType()
+		{
+			return User.sObjectType;
+		}
+	}
+	
+	private class Testfflib_LeadSelector extends fflib_SObjectSelector
+	{
+		public Testfflib_LeadSelector()
+		{
+			super();
+		}
+		
+		public List<Schema.SObjectField> getSObjectFieldList()
+		{
+			return new List<Schema.SObjectField> {
+				Lead.OwnerId,
+				Lead.Id
+			};
+		}
+		
+		public Schema.SObjectType getSObjectType()
+		{
+			return Lead.sObjectType;
+		}
 	}	
 }