HHH-18830 fallout from fix

improve handling of @MapKey and add tests

Signed-off-by: Gavin King <gavin@hibernate.org>

hibernate-core/src/main/java/org/hibernate/boot/model/internal/MapBinder.java

@@ -88,9 +88,10 @@ public class MapBinder extends CollectionBinder {
 
 	@Override
 	SecondPass getSecondPass() {
-		return new CollectionSecondPass( MapBinder.this.collection ) {
+		return new CollectionSecondPass( collection ) {
 			public void secondPass(java.util.Map<String, PersistentClass> persistentClasses)
 					throws MappingException {
+				getMap().setHasMapKeyProperty( hasMapKeyProperty );
 				bindStarToManySecondPass( persistentClasses );
 				bindKeyFromAssociationTable(
 						getElementType(),

hibernate-core/src/main/java/org/hibernate/mapping/IndexedCollection.java

@@ -47,6 +47,10 @@ public abstract class IndexedCollection extends Collection {
 		return true;
 	}
 
+	public boolean hasMapKeyProperty() {
+		return false;
+	}
+
 	@Override
 	public boolean isSame(Collection other) {
 		return other instanceof IndexedCollection

hibernate-core/src/main/java/org/hibernate/mapping/Map.java

@@ -22,6 +22,7 @@ import org.hibernate.usertype.UserCollectionType;
 public class Map extends IndexedCollection {
 
 	private String mapKeyPropertyName;
+	private boolean hasMapKeyProperty;
 
 	public Map(MetadataBuildingContext buildingContext, PersistentClass owner) {
 		super( buildingContext, owner );
@@ -75,4 +76,13 @@ public class Map extends IndexedCollection {
 	public Object accept(ValueVisitor visitor) {
 		return visitor.accept(this);
 	}
+
+	@Override
+	public boolean hasMapKeyProperty() {
+		return hasMapKeyProperty;
+	}
+
+	public void setHasMapKeyProperty(boolean hasMapKeyProperty) {
+		this.hasMapKeyProperty = hasMapKeyProperty;
+	}
 }

hibernate-core/src/main/java/org/hibernate/persister/collection/AbstractCollectionPersister.java

@@ -400,14 +400,14 @@ public abstract class AbstractCollectionPersister
 
 		// INDEX AND ROW SELECT
 
-		final boolean hasIndex = collectionBootDescriptor.isIndexed();
-		if ( hasIndex ) {
+		if ( collectionBootDescriptor instanceof IndexedCollection indexedCollection ) {
+			assert collectionBootDescriptor.isIndexed();
 			// NativeSQL: collect index column and auto-aliases
-			final IndexedCollection indexedCollection = (IndexedCollection) collectionBootDescriptor;
-			indexType = indexedCollection.getIndex().getType();
-			final int indexSpan = indexedCollection.getIndex().getColumnSpan();
-			final boolean[] indexColumnInsertability = indexedCollection.getIndex().getColumnInsertability();
-			final boolean[] indexColumnUpdatability = indexedCollection.getIndex().getColumnUpdateability();
+			final Value index = indexedCollection.getIndex();
+			indexType = index.getType();
+			final int indexSpan = index.getColumnSpan();
+			final boolean[] indexColumnInsertability = index.getColumnInsertability();
+			final boolean[] indexColumnUpdatability = index.getColumnUpdateability();
 			indexColumnNames = new String[indexSpan];
 			indexFormulaTemplates = new String[indexSpan];
 			indexFormulas = new String[indexSpan];
@@ -416,7 +416,7 @@ public abstract class AbstractCollectionPersister
 			indexColumnAliases = new String[indexSpan];
 			int i = 0;
 			boolean hasFormula = false;
-			for ( Selectable selectable: indexedCollection.getIndex().getSelectables() ) {
+			for ( Selectable selectable: index.getSelectables() ) {
 				indexColumnAliases[i] = selectable.getAlias( dialect );
 				if ( selectable.isFormula() ) {
 					final Formula indexForm = (Formula) selectable;
@@ -428,19 +428,16 @@ public abstract class AbstractCollectionPersister
 					indexFormulas[i] = indexForm.getFormula();
 					hasFormula = true;
 				}
-				// Treat a mapped-by index like a formula to avoid trying to set it in insert/update
-				// Previously this was a sub-query formula, but was changed to represent the proper mapping
-				// which enables optimizations for queries. The old insert/update code wasn't adapted yet though.
-				// For now, this is good enough, because the formula is never used anymore,
-				// since all read paths go through the new code that can properly handle this case
-				else if ( indexedCollection instanceof org.hibernate.mapping.Map map
-						&& map.getMapKeyPropertyName() != null ) {
-					final Column indexCol = (Column) selectable;
-					indexFormulaTemplates[i] = Template.TEMPLATE + indexCol.getQuotedName( dialect );
-					indexFormulas[i] = indexCol.getQuotedName( dialect );
-					hasFormula = true;
-				}
 				else {
+					if ( indexedCollection.hasMapKeyProperty() ) {
+						// If the Map key is set via @MapKey, it should not be written
+						// since it is a reference to a field of the associated entity.
+						// (Note that the analogous situation for Lists never arises
+						// because @OrderBy is treated as defining a sorted bag.)
+						indexColumnInsertability[i] = false;
+						indexColumnUpdatability[i] = false;
+						hasFormula = true; // this is incorrect, but needed for some reason
+					}
 					final Column indexCol = (Column) selectable;
 					indexColumnNames[i] = indexCol.getQuotedName( dialect );
 					indexColumnIsGettable[i] = true;

hibernate-core/src/test/java/org/hibernate/orm/test/collection/mapkey/Book.java

@@ -0,0 +1,53 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.collection.mapkey;
+
+import jakarta.persistence.Column;
+import jakarta.persistence.Entity;
+import jakarta.persistence.Id;
+import jakarta.persistence.MapKey;
+import jakarta.persistence.OneToMany;
+import jakarta.validation.constraints.Size;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static jakarta.persistence.CascadeType.PERSIST;
+
+@Entity
+class Book {
+	@Id
+	@Size(min=10, max = 13)
+	String isbn;
+
+	@OneToMany(cascade = PERSIST,
+			mappedBy = "isbn")
+	@MapKey(name = "name")
+	Map<String,Chapter> chapters;
+
+	Book(String isbn) {
+		this.isbn = isbn;
+		chapters = new HashMap<>();
+	}
+
+	Book() {
+	}
+}
+
+@Entity
+class Chapter {
+	@Id String isbn;
+	@Id @Column(name = "chapter_name") String name;
+	String text;
+
+	Chapter(String isbn, String name, String text) {
+		this.isbn = isbn;
+		this.name = name;
+		this.text = text;
+	}
+
+	Chapter() {
+	}
+}

hibernate-core/src/test/java/org/hibernate/orm/test/collection/mapkey/MapKeyMappedByTest.java

@@ -0,0 +1,35 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.collection.mapkey;
+
+import org.hibernate.testing.jdbc.SQLStatementInspector;
+import org.hibernate.testing.orm.junit.EntityManagerFactoryScope;
+import org.hibernate.testing.orm.junit.Jpa;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+@Jpa(annotatedClasses = {Chapter.class, Book.class},
+		useCollectingStatementInspector = true)
+public class MapKeyMappedByTest {
+	@Test void test(EntityManagerFactoryScope scope) {
+		scope.inTransaction( em -> {
+			Book book = new Book( "XXX" );
+			em.persist( book );
+			book.chapters.put( "one", new Chapter("XXX", "one", "Lorem ipsum") );
+		} );
+		scope.inTransaction( em -> {
+			Book book = em.find( Book.class, "XXX" );
+			book.chapters.put( "two", new Chapter("XXX", "two", "Lorem ipsum") );
+		} );
+		List<String> queries = ((SQLStatementInspector) scope.getStatementInspector()).getSqlQueries();
+		assertEquals( 5, queries.size() ); // we can't put in a Map without initializing it
+		scope.inTransaction( em -> {
+			assertEquals(  em.find( Book.class, "XXX" ).chapters.size(), 2 );
+		} );
+	}
+}

hibernate-core/src/test/java/org/hibernate/orm/test/collection/mapkeycolumn/Book.java

@@ -0,0 +1,53 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.collection.mapkeycolumn;
+
+import jakarta.persistence.Column;
+import jakarta.persistence.Entity;
+import jakarta.persistence.Id;
+import jakarta.persistence.MapKeyColumn;
+import jakarta.persistence.OneToMany;
+import jakarta.validation.constraints.Size;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static jakarta.persistence.CascadeType.PERSIST;
+
+@Entity
+class Book {
+	@Id
+	@Size(min=10, max = 13)
+	String isbn;
+
+	@OneToMany(cascade = PERSIST,
+			mappedBy = "isbn")
+	@MapKeyColumn(name = "chapter_name")
+	Map<String,Chapter> chapters;
+
+	Book(String isbn) {
+		this.isbn = isbn;
+		chapters = new HashMap<>();
+	}
+
+	Book() {
+	}
+}
+
+@Entity
+class Chapter {
+	@Id String isbn;
+	@Id @Column(name = "chapter_name") String name;
+	String text;
+
+	Chapter(String isbn, String name, String text) {
+		this.isbn = isbn;
+		this.name = name;
+		this.text = text;
+	}
+
+	Chapter() {
+	}
+}

hibernate-core/src/test/java/org/hibernate/orm/test/collection/mapkeycolumn/MapKeyColumnMappedByTest.java

@@ -0,0 +1,35 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.collection.mapkeycolumn;
+
+import org.hibernate.testing.jdbc.SQLStatementInspector;
+import org.hibernate.testing.orm.junit.EntityManagerFactoryScope;
+import org.hibernate.testing.orm.junit.Jpa;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+@Jpa(annotatedClasses = {Chapter.class, Book.class},
+		useCollectingStatementInspector = true)
+public class MapKeyColumnMappedByTest {
+	@Test void test(EntityManagerFactoryScope scope) {
+		scope.inTransaction( em -> {
+			Book book = new Book( "XXX" );
+			em.persist( book );
+			book.chapters.put( "one", new Chapter("XXX", "one", "Lorem ipsum") );
+		} );
+		scope.inTransaction( em -> {
+			Book book = em.find( Book.class, "XXX" );
+			book.chapters.put( "two", new Chapter("XXX", "two", "Lorem ipsum") );
+		} );
+		List<String> queries = ((SQLStatementInspector) scope.getStatementInspector()).getSqlQueries();
+		assertEquals( 5, queries.size() ); // we can't put in a Map without initializing it
+		scope.inTransaction( em -> {
+			assertEquals(  em.find( Book.class, "XXX" ).chapters.size(), 2 );
+		} );
+	}
+}