Replace entities with ids in CTEs for Id attributes and FK keys, fixes …

…Blazebit#571

core/impl/src/main/java/com/blazebit/persistence/impl/AbstractCTECriteriaBuilder.java

@@ -210,22 +210,24 @@ protected List<String> prepareAndGetAttributes() {
             attributes.add(attributeName);
 
             if (JpaMetamodelUtils.isJoinable(attributePath.get(attributePath.size() - 1))) {
-                // We have to map *-to-one relationships to their ids
-                EntityType<?> type = mainQuery.metamodel.entity(attributeEntry.getElementClass());
-                Attribute<?, ?> idAttribute = JpaMetamodelUtils.getSingleIdAttribute(type);
+                List<String> idOrUniqueKeyProps = cbf.getJpaProvider().getIdentifierOrUniqueKeyEmbeddedPropertyNames(cteType, attributeName);
+                // We have to map *-to-one relationships to their id or unique props
                 // NOTE: Since we are talking about *-to-ones, the expression can only be a path to an object
                 // so it is safe to just append the id to the path
                 Expression selectExpression = selectManager.getSelectInfos().get(bindingEntry.getValue()).getExpression();
 
                 // TODO: Maybe also allow Treat, Case-When, Array?
                 if (selectExpression instanceof NullExpression) {
                     // When binding null, we don't have to adapt anything
-                } else if (selectExpression instanceof PathExpression) {
+                } else if (selectExpression instanceof PathExpression && idOrUniqueKeyProps.size() == 1) {
                     PathExpression pathExpression = (PathExpression) selectExpression;
+                    String idOrUniqueKeyProp = idOrUniqueKeyProps.get(0);
                     // Only append the id if it's not already there
-                    if (!idAttribute.getName().equals(pathExpression.getExpressions().get(pathExpression.getExpressions().size() - 1).toString())) {
-                        pathExpression.getExpressions().add(new PropertyExpression(idAttribute.getName()));
+                    if (!idOrUniqueKeyProp.equals(pathExpression.getExpressions().get(pathExpression.getExpressions().size() - 1).toString())) {
+                        pathExpression.getExpressions().add(new PropertyExpression(idOrUniqueKeyProp));
                     }
+
+                    // TODO: Handle >2 case
                 } else {
                     throw new IllegalArgumentException("Illegal expression '" + selectExpression.toString() + "' for binding relation '" + attributeName + "'!");
                 }

core/parser/src/main/java/com/blazebit/persistence/parser/PathTargetResolvingExpressionVisitor.java

@@ -239,6 +239,8 @@ public void visit(PropertyExpression expression) {
                     keyType = metamodel.type(typeArguments[0]);
                 }
             }
+        } else if (attribute instanceof SingularAttribute<?,?>) {
+            valueType = type = ((SingularAttribute) attribute).getType();
         } else {
             valueType = type;
         }

core/parser/src/main/java/com/blazebit/persistence/parser/util/JpaMetamodelUtils.java

@@ -546,9 +546,10 @@ public static boolean isJoinable(Attribute<?, ?> attr) {
         SingularAttribute<?, ?> singularAttribute = (SingularAttribute<?, ?>) attr;
         // This is a special case for datanucleus... apparently an embedded id is an ONE_TO_ONE association although I think it should be an embedded
         // TODO: create a test case for datanucleus and report the problem
-        if (singularAttribute.isId()) {
-            return false;
-        }
+        // TODO: Workaround should either be DN specific, because an ID prop can in fact be joinable (i.e. OneToOne's)
+//        if (singularAttribute.isId()) {
+//            return false;
+//        }
         return attr.getPersistentAttributeType() == Attribute.PersistentAttributeType.MANY_TO_ONE
             || attr.getPersistentAttributeType() == Attribute.PersistentAttributeType.ONE_TO_ONE;
     }
@@ -564,9 +565,10 @@ public static boolean isCompositeNode(Attribute<?, ?> attr) {
         SingularAttribute<?, ?> singularAttribute = (SingularAttribute<?, ?>) attr;
         // This is a special case for datanucleus... apparently an embedded id is an ONE_TO_ONE association although I think it should be an embedded
         // TODO: create a test case for datanucleus and report the problem
-        if (singularAttribute.isId()) {
-            return false;
-        }
+        // TODO: Workaround should either be DN specific, because an ID prop can in fact be joinable (i.e. OneToOne's)
+//        if (singularAttribute.isId()) {
+//            return false;
+//        }
         return attr.getPersistentAttributeType() == Attribute.PersistentAttributeType.MANY_TO_ONE
                 || attr.getPersistentAttributeType() == Attribute.PersistentAttributeType.ONE_TO_ONE;
     }

core/testsuite/src/test/java/com/blazebit/persistence/testsuite/Issue571Test.java

@@ -1,11 +1,33 @@
+/*
+ * Copyright 2014 - 2018 Blazebit.
+ *
+ * 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.
+ */
+
 package com.blazebit.persistence.testsuite;
 
 import com.blazebit.persistence.CTE;
 import com.blazebit.persistence.CriteriaBuilder;
 import com.blazebit.persistence.impl.query.CustomSQLTypedQuery;
+import com.blazebit.persistence.testsuite.base.jpa.category.NoDatanucleus;
+import com.blazebit.persistence.testsuite.base.jpa.category.NoEclipselink;
+import com.blazebit.persistence.testsuite.base.jpa.category.NoMySQL;
+import com.blazebit.persistence.testsuite.base.jpa.category.NoOpenJPA;
+import com.blazebit.persistence.testsuite.base.jpa.category.NoOracle;
 import com.blazebit.persistence.testsuite.tx.TxVoidWork;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.experimental.categories.Category;
 
 import javax.persistence.Entity;
 import javax.persistence.EntityManager;
@@ -39,9 +61,9 @@ public void work(EntityManager em) {
         });
     }
 
-
     @Test
-    public void testImplicitJoinAttributeOfCteAssociation() {
+    @Category({ NoDatanucleus.class, NoEclipselink.class, NoOpenJPA.class, NoMySQL.class, NoOracle.class })
+    public void testBindingCteAssociationToEntity() {
         transactional(new TxVoidWork() {
             @Override
             public void work(EntityManager em) {
@@ -52,7 +74,7 @@ public void work(EntityManager em) {
                         .bind("id").select("1")
                         .unionAll()
                         .from(MyEntity.class, "a")
-                        .where("a.attribute").notIn().from(Cte.class, "b").select("b.myEntity.attribute").end()
+                        .where("a.attribute").eq("bogus")
                         .bind("myEntity").select("a")
                         .bind("id").select("1")
                         .end()
@@ -80,19 +102,20 @@ public void work(EntityManager em) {
     }
 
     @Test
-    public void testBindingCteAssociationToEntity() {
+    @Category({ NoDatanucleus.class, NoEclipselink.class, NoOpenJPA.class, NoMySQL.class, NoOracle.class })
+    public void testBindingCteAssociationToEntityId() {
         transactional(new TxVoidWork() {
             @Override
             public void work(EntityManager em) {
                 CriteriaBuilder<MyEntity> select = cbf.create(em, MyEntity.class)
                         .withRecursive(Cte.class)
                         .from(MyEntity.class, "a")
-                        .bind("myEntity").select("a")
+                        .bind("myEntity").select("a.id")
                         .bind("id").select("1")
                         .unionAll()
                         .from(MyEntity.class, "a")
                         .where("a.attribute").eq("bogus")
-                        .bind("myEntity").select("a")
+                        .bind("myEntity").select("a.id")
                         .bind("id").select("1")
                         .end()
                         .from(Cte.class)
@@ -114,19 +137,23 @@ public void work(EntityManager em) {
 
                 assertEquals("Final query should project 2 columns separated by a single comma",
                         queryPart.indexOf(","), queryPart.lastIndexOf(","));
+
             }
         });
     }
 
+
     @Test
-    public void testBindingCteAssociationToEntityId() {
+    @Category({ NoDatanucleus.class, NoEclipselink.class, NoOpenJPA.class, NoMySQL.class, NoOracle.class })
+
+    public void testBindingCteAssociationToEntityId2() {
         transactional(new TxVoidWork() {
             @Override
             public void work(EntityManager em) {
                 CriteriaBuilder<MyEntity> select = cbf.create(em, MyEntity.class)
                         .withRecursive(Cte.class)
                         .from(MyEntity.class, "a")
-                        .bind("myEntity").select("a.id")
+                        .bind("myEntity.id").select("a.id")
                         .bind("id").select("1")
                         .unionAll()
                         .from(MyEntity.class, "a")