Fields that are polymorphic fail due to first type being used.

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 String getFieldPath(String fieldName){
+	private String getFieldPath(String fieldName, Schema.sObjectType relatedSObjectType){
 		if(!fieldName.contains('.')){ //single field
 			Schema.SObjectField token = fflib_SObjectDescribe.getDescribe(table).getField(fieldName.toLowerCase());
 			if(token == null)
@@ -108,7 +108,22 @@ public class fflib_QueryFactory { //No explicit sharing declaration - inherit fr
 			}
 
 			if(token != null && i.hasNext() && tokenDescribe.getSOAPType() == Schema.SOAPType.ID){
-				lastSObjectType = tokenDescribe.getReferenceTo()[0]; //if it's polymorphic doesn't matter which one we get
+				if(relatedSObjectType == null) {
+					lastSObjectType = tokenDescribe.getReferenceTo()[0]; //if it's polymorphic get the first one - user did not specify the one to use.
+				}else{
+					List<Schema.sObjectType> relatedObjs = tokenDescribe.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;
+							}
+						}
+					}else{
+						lastSObjectType = relatedOBjs[0]; //only one type returned, use it.
+					}
+				}
 				fieldPath.add(tokenDescribe.getRelationshipName());
 			}else if(token != null && !i.hasNext()){
 				fieldPath.add(tokenDescribe.getName());
@@ -122,6 +137,10 @@ public class fflib_QueryFactory { //No explicit sharing declaration - inherit fr
 
 		return String.join(fieldPath,'.');
 	}
+
+	private String getFieldPath(String fieldName) {
+		return this.getFieldPath(fieldName, null);
+	}
 
 	@TestVisible
 	private static String getFieldTokenPath(Schema.SObjectField field){
@@ -197,16 +216,27 @@ public class fflib_QueryFactory { //No explicit sharing declaration - inherit fr
 		this.sortSelectFields = doSort;
 		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.
 	 **/
 	public fflib_QueryFactory selectField(String fieldName){ 		
-		fields.add( getFieldPath(fieldName) );
+		fields.add( getFieldPath(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( getFieldPath(fieldName, relatedObjectType) );
+        	return this;
+    	}
+
 	/**
 	 * Selects a field, avoiding the possible ambiguitiy of String API names.
 	 * @see #selectField(String)
@@ -760,6 +790,18 @@ 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 \''+fieldname+'\'' );
+                }
+        }
+
 	public class InvalidFieldSetException extends Exception{}
 	public class NonReferenceFieldException extends Exception{}
 	public class InvalidSubqueryRelationshipException extends Exception{}	

fflib/src/classes/fflib_SObjectSelector.cls

@@ -329,7 +329,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

@@ -300,5 +300,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;
+		}
 	}	
 }