Use @IdClass as identifier for IdClasses with a single attribute for repository operations.

We now use the IdClass type for lookups when the entity defines a `@IdClass`. Previously, we uses the type of the defined singular identifier attribute which lead to invalid queries.

Closes #2330

src/main/java/org/springframework/data/jpa/provider/PersistenceProvider.java

@@ -20,15 +20,19 @@ import static org.springframework.data.jpa.provider.PersistenceProvider.Constant
 
 import java.util.Collections;
 import java.util.NoSuchElementException;
+import java.util.Set;
 
 import javax.persistence.EntityManager;
 import javax.persistence.Query;
+import javax.persistence.metamodel.IdentifiableType;
 import javax.persistence.metamodel.Metamodel;
+import javax.persistence.metamodel.SingularAttribute;
 
 import org.eclipse.persistence.jpa.JpaQuery;
 import org.eclipse.persistence.queries.ScrollableCursor;
 import org.hibernate.ScrollMode;
 import org.hibernate.ScrollableResults;
+import org.hibernate.metamodel.model.domain.spi.IdentifiableTypeDescriptor;
 import org.hibernate.proxy.HibernateProxy;
 import org.springframework.data.util.CloseableIterator;
 import org.springframework.lang.Nullable;
@@ -93,6 +97,17 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor {
 			return ((HibernateProxy) entity).getHibernateLazyInitializer().getIdentifier();
 		}
 
+		/*
+		 * (non-Javadoc)
+		 * @see org.springframework.data.jpa.provider.PersistenceProvider#getIdClassAttributes(javax.persistence.metamodel.IdentifiableType)
+		 */
+		@Override
+		public <T> Set<SingularAttribute<? super T, ?>> getIdClassAttributes(IdentifiableType<T> type) {
+			return type instanceof IdentifiableTypeDescriptor && ((IdentifiableTypeDescriptor<T>) type).hasIdClass()
+					? super.getIdClassAttributes(type)
+					: Collections.emptySet();
+		}
+
 		/*
 		 * (non-Javadoc)
 		 * @see org.springframework.data.jpa.provider.PersistenceProvider#executeQueryWithResultStream(javax.persistence.Query)
@@ -291,6 +306,20 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor {
 		return true;
 	}
 
+	/**
+	 * @param type the entity type.
+	 * @return the set of identifier attributes used in a {@code @IdClass} for {@code type}. Empty when {@code type} does
+	 *         not use {@code @IdClass}.
+	 * @since 2.5.6
+	 */
+	public <T> Set<SingularAttribute<? super T, ?>> getIdClassAttributes(IdentifiableType<T> type) {
+		try {
+			return type.getIdClassAttributes();
+		} catch (IllegalArgumentException e) {
+			return Collections.emptySet();
+		}
+	}
+
 	/**
 	 * Holds the PersistenceProvider specific interface names.
 	 *

src/main/java/org/springframework/data/jpa/repository/support/JpaMetamodelEntityInformation.java

@@ -86,7 +86,7 @@ public class JpaMetamodelEntityInformation<T, ID> extends JpaEntityInformationSu
 
 		IdentifiableType<T> identifiableType = (IdentifiableType<T>) type;
 
-		this.idMetadata = new IdMetadata<>(identifiableType);
+		this.idMetadata = new IdMetadata<>(identifiableType, PersistenceProvider.fromMetamodel(metamodel));
 		this.versionAttribute = findVersionAttribute(identifiableType, metamodel);
 	}
 
@@ -260,20 +260,22 @@ public class JpaMetamodelEntityInformation<T, ID> extends JpaEntityInformationSu
 	private static class IdMetadata<T> implements Iterable<SingularAttribute<? super T, ?>> {
 
 		private final IdentifiableType<T> type;
+		private final Set<SingularAttribute<? super T, ?>> idClassAttributes;
 		private final Set<SingularAttribute<? super T, ?>> attributes;
 		private @Nullable Class<?> idType;
 
 		@SuppressWarnings("unchecked")
-		IdMetadata(IdentifiableType<T> source) {
+		IdMetadata(IdentifiableType<T> source, PersistenceProvider persistenceProvider) {
 
 			this.type = source;
+			this.idClassAttributes = persistenceProvider.getIdClassAttributes(source);
 			this.attributes = (Set<SingularAttribute<? super T, ?>>) (source.hasSingleIdAttribute()
 					? Collections.singleton(source.getId(source.getIdType().getJavaType()))
 					: source.getIdClassAttributes());
 		}
 
 		boolean hasSimpleId() {
-			return attributes.size() == 1;
+			return idClassAttributes.isEmpty() && attributes.size() == 1;
 		}
 
 		public Class<?> getType() {
@@ -296,18 +298,26 @@ public class JpaMetamodelEntityInformation<T, ID> extends JpaEntityInformationSu
 		private Class<?> tryExtractIdTypeWithFallbackToIdTypeLookup() {
 
 			try {
+
+				Class<?> idClassType = lookupIdClass(type);
+				if (idClassType != null) {
+					return idClassType;
+				}
+
 				Type<?> idType = type.getIdType();
-				return idType == null ? fallbackIdTypeLookup(type) : idType.getJavaType();
+				return idType == null ? null : idType.getJavaType();
 			} catch (IllegalStateException e) {
 				// see https://hibernate.onjira.com/browse/HHH-6951
-				return fallbackIdTypeLookup(type);
+				return null;
 			}
 		}
 
 		@Nullable
-		private static Class<?> fallbackIdTypeLookup(IdentifiableType<?> type) {
+		private static Class<?> lookupIdClass(IdentifiableType<?> type) {
 
-			IdClass annotation = AnnotationUtils.findAnnotation(type.getJavaType(), IdClass.class);
+			IdClass annotation = type.getJavaType() != null
+					? AnnotationUtils.findAnnotation(type.getJavaType(), IdClass.class)
+					: null;
 			return annotation == null ? null : annotation.value();
 		}
 

src/test/java/org/springframework/data/jpa/domain/sample/PersistableWithSingleIdClass.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2021 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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 org.springframework.data.jpa.domain.sample;
+
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import javax.persistence.IdClass;
+
+/**
+ * Sample entity using {@link IdClass} annotation to demarcate ids.
+ *
+ * @author Mark Paluch
+ */
+@Entity
+@IdClass(PersistableWithSingleIdClassPK.class)
+public class PersistableWithSingleIdClass {
+
+	private static final long serialVersionUID = 1L;
+
+	@Id private Long first;
+
+	protected PersistableWithSingleIdClass() {
+
+	}
+
+	public PersistableWithSingleIdClass(Long first) {
+		this.first = first;
+	}
+
+	/**
+	 * @return the first
+	 */
+	public Long getFirst() {
+		return first;
+	}
+
+}

src/test/java/org/springframework/data/jpa/domain/sample/PersistableWithSingleIdClassPK.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2021 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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 org.springframework.data.jpa.domain.sample;
+
+import static org.springframework.util.ObjectUtils.*;
+
+import java.io.Serializable;
+
+/**
+ * @author Mark Paluch
+ */
+public class PersistableWithSingleIdClassPK implements Serializable {
+
+	private static final long serialVersionUID = 23126782341L;
+
+	private Long first;
+
+	public PersistableWithSingleIdClassPK() {
+
+	}
+
+	public PersistableWithSingleIdClassPK(Long first) {
+		this.first = first;
+	}
+
+	public void setFirst(Long first) {
+		this.first = first;
+	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see java.lang.Object#equals(java.lang.Object)
+	 */
+	@Override
+	public boolean equals(Object obj) {
+
+		if (this == obj) {
+			return true;
+		}
+
+		if (obj == null || !(obj.getClass().equals(getClass()))) {
+			return false;
+		}
+
+		PersistableWithSingleIdClassPK that = (PersistableWithSingleIdClassPK) obj;
+
+		return nullSafeEquals(this.first, that.first);
+	}
+
+	/*
+	 * (non-Javadoc)
+	 * @see java.lang.Object#hashCode()
+	 */
+	@Override
+	public int hashCode() {
+
+		int result = 17;
+
+		result += nullSafeHashCode(this.first);
+
+		return result;
+	}
+}

src/test/java/org/springframework/data/jpa/repository/support/JpaMetamodelEntityInformationIntegrationTests.java

@@ -92,6 +92,18 @@ public class JpaMetamodelEntityInformationIntegrationTests {
 		assertThat(id).isEqualTo(new PersistableWithIdClassPK(2L, 4L));
 	}
 
+	@Test // GH-2330
+	void returnsIdOfSingleAttributeIdClassCorrectly() {
+
+		PersistableWithSingleIdClass entity = new PersistableWithSingleIdClass(2L);
+
+		JpaEntityInformation<PersistableWithSingleIdClass, ?> information = getEntityInformation(
+				PersistableWithSingleIdClass.class, em);
+		Object id = information.getId(entity);
+
+		assertThat(id).isEqualTo(new PersistableWithSingleIdClassPK(2L));
+	}
+
 	@Test // DATAJPA-413
 	void returnsIdOfEntityWithIdClassCorrectly() {
 

src/test/resources/META-INF/persistence.xml

@@ -32,6 +32,8 @@
 		<class>org.springframework.data.jpa.domain.sample.Order</class>
 		<class>org.springframework.data.jpa.domain.sample.Parent</class>
 		<class>org.springframework.data.jpa.domain.sample.PersistableWithIdClass</class>
+		<class>org.springframework.data.jpa.domain.sample.PersistableWithSingleIdClass
+		</class>
 		<class>org.springframework.data.jpa.domain.sample.PrimitiveVersionProperty</class>
 		<class>org.springframework.data.jpa.domain.sample.Product</class>
 		<class>org.springframework.data.jpa.domain.sample.Role</class>
@@ -89,9 +91,9 @@
 			<property name="hibernate.dialect" value="org.hibernate.dialect.HSQLDialect" />
 		</properties>
 	</persistence-unit>
-	
+
 	<!--  Custom PUs for metadata tests -->
-	
+
 	<persistence-unit name="metadata">
 		<provider>org.hibernate.jpa.HibernatePersistenceProvider</provider>
 		<class>org.springframework.data.jpa.domain.sample.CustomAbstractPersistable</class>