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;
     }

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

@@ -1,3 +1,19 @@
+/*
+ * 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;
@@ -158,6 +174,47 @@ public void work(EntityManager em) {
         });
     }
 
+
+    @Test
+    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.id").select("a.id")
+                        .bind("id").select("1")
+                        .unionAll()
+                        .from(MyEntity.class, "a")
+                        .where("a.attribute").eq("bogus")
+                        .bind("myEntity").select("a.id")
+                        .bind("id").select("1")
+                        .end()
+                        .from(Cte.class)
+                        .select("myEntity");
+
+                List<MyEntity> resultList = select
+                        .getResultList();
+
+                String sql = ((CustomSQLTypedQuery) select.getQuery()).getQuerySpecification().getSql();
+                String cteBase = sql.substring(sql.indexOf("\nselect") + 1, sql.lastIndexOf("UNION ALL") - 1);
+                String cteRecursivePart = sql.substring(sql.indexOf("UNION ALL")+10, sql.lastIndexOf("select")-3);
+                String queryPart = sql.substring(sql.lastIndexOf("select"));
+
+                assertEquals("Recursive base part should project 2 columns separated by a single comma",
+                        cteBase.indexOf(","), cteBase.lastIndexOf(","));
+
+                assertEquals("Recursive part should project 2 columns separated by a single comma",
+                        cteRecursivePart.indexOf(","), cteRecursivePart.lastIndexOf(","));
+
+                assertEquals("Final query should project 2 columns separated by a single comma",
+                        queryPart.indexOf(","), queryPart.lastIndexOf(","));
+
+            }
+        });
+    }
+
     @CTE
     @Entity
     @IdClass(Cte.CteIdClass.class)