HHH-19608 automatically add @PartitionKey to the primary key

hibernate-core/src/main/java/org/hibernate/dialect/Dialect.java

@@ -4622,6 +4622,16 @@ public boolean supportsPartitionBy() {
 		return false;
 	}
 
+	/**
+	 * Does this dialect require that the columns listed in
+	 * {@code partition by} also occur in the primary key?
+	 *
+	 * @since 7.1
+	 */
+	public boolean addPartitionKeyToPrimaryKey() {
+		return false;
+	}
+
 	/**
 	 * Override {@link DatabaseMetaData#supportsNamedParameters()}.
 	 *

hibernate-core/src/main/java/org/hibernate/dialect/MySQLDialect.java

@@ -1099,6 +1099,16 @@ public String getSelectGUIDString() {
 		return "select uuid()";
 	}
 
+	@Override
+	public boolean supportsPartitionBy() {
+		return true;
+	}
+
+	@Override
+	public boolean addPartitionKeyToPrimaryKey() {
+		return true;
+	}
+
 	@Override
 	public boolean supportsCommentOn() {
 		return true;

hibernate-core/src/main/java/org/hibernate/dialect/PostgreSQLDialect.java

@@ -821,6 +821,11 @@ public boolean supportsPartitionBy() {
 		return true;
 	}
 
+	@Override
+	public boolean addPartitionKeyToPrimaryKey() {
+		return true;
+	}
+
 	@Override
 	public boolean supportsNonQueryWithCTE() {
 		return true;

hibernate-core/src/main/java/org/hibernate/id/enhanced/ExportableColumnHelper.java

@@ -143,6 +143,11 @@ public boolean isColumnInsertable(int index) {
 			public boolean isColumnUpdateable(int index) {
 				return true;
 			}
+
+			@Override
+			public boolean isPartitionKey() {
+				return false;
+			}
 		} );
 		return column;
 	}

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

@@ -897,4 +897,9 @@ public Supplier<? extends Expectation> getDeleteAllExpectation() {
 	public void setDeleteAllExpectation(Supplier<? extends Expectation> deleteAllExpectation) {
 		this.deleteAllExpectation = deleteAllExpectation;
 	}
+
+	@Override
+	public boolean isPartitionKey() {
+		return false;
+	}
 }

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

@@ -234,4 +234,9 @@ public boolean isColumnUpdateable(int index) {
 	public String toString() {
 		return getClass().getSimpleName();
 	}
+
+	@Override
+	public boolean isPartitionKey() {
+		return false;
+	}
 }

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

@@ -427,10 +427,22 @@ public void createPrimaryKey() {
 		final PrimaryKey pk = new PrimaryKey( table );
 		pk.setName( PK_ALIAS.toAliasString( table.getName() ) );
 		pk.addColumns( getKey() );
-
+		if ( addPartitionKeyToPrimaryKey() ) {
+			for ( Property property : getProperties() ) {
+				if ( property.getValue().isPartitionKey() ) {
+					pk.addColumns( property.getValue() );
+				}
+			}
+		}
 		table.setPrimaryKey( pk );
 	}
 
+	private boolean addPartitionKeyToPrimaryKey() {
+		return metadataBuildingContext.getMetadataCollector()
+				.getDatabase().getDialect()
+				.addPartitionKeyToPrimaryKey();
+	}
+
 	public abstract String getWhere();
 
 	public int getBatchSize() {

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

@@ -123,6 +123,8 @@ private Type getIdType(EntityType entityType) {
 
 	boolean isAlternateUniqueKey();
 
+	boolean isPartitionKey();
+
 	boolean isNullable();
 
 	void createForeignKey();