From 0914fba014fa275d15d2bb36e1b8b61946efa228 Mon Sep 17 00:00:00 2001
From: Marco Belladelli <marco@hibernate.org>
Date: Wed, 2 Aug 2023 15:31:36 +0200
Subject: [PATCH] HHH-16616 Apply discriminator by pruning table group

Also, only apply the discriminator for loads since in queries the pruning is already handled by the entity name usage logic.
---
 .../persister/collection/OneToManyPersister.java | 16 ++++++++++------
 .../entity/AbstractEntityPersister.java          |  3 ++-
 .../query/sqm/sql/BaseSqmToSqlAstConverter.java  |  5 +++++
 .../sql/ast/spi/SqlAstCreationState.java         |  5 +++++
 .../sql/ast/tree/from/NamedTableReference.java   |  2 +-
 5 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/hibernate-core/src/main/java/org/hibernate/persister/collection/OneToManyPersister.java b/hibernate-core/src/main/java/org/hibernate/persister/collection/OneToManyPersister.java
index e909790090..742190207c 100644
--- a/hibernate-core/src/main/java/org/hibernate/persister/collection/OneToManyPersister.java
+++ b/hibernate-core/src/main/java/org/hibernate/persister/collection/OneToManyPersister.java
@@ -300,12 +300,16 @@ public class OneToManyPersister extends AbstractCollectionPersister {
 			SqlAstCreationState astCreationState) {
 		super.applyWhereFragments( predicateConsumer, alias, tableGroup, astCreationState );
 
-		getElementPersisterInternal().applyDiscriminator(
-				predicateConsumer,
-				alias,
-				tableGroup,
-				astCreationState
-		);
+		if ( !astCreationState.supportsEntityNameUsage() ) {
+			// We only need to apply discriminator for loads, since queries with joined
+			// inheritance subtypes are already filtered by the entity name usage logic
+			getElementPersisterInternal().applyDiscriminator(
+					predicateConsumer,
+					alias,
+					tableGroup,
+					astCreationState
+			);
+		}
 	}
 
 	@Override
diff --git a/hibernate-core/src/main/java/org/hibernate/persister/entity/AbstractEntityPersister.java b/hibernate-core/src/main/java/org/hibernate/persister/entity/AbstractEntityPersister.java
index 1bd3247016..d4363de54a 100644
--- a/hibernate-core/src/main/java/org/hibernate/persister/entity/AbstractEntityPersister.java
+++ b/hibernate-core/src/main/java/org/hibernate/persister/entity/AbstractEntityPersister.java
@@ -14,6 +14,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.BitSet;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.IdentityHashMap;
@@ -3013,7 +3014,7 @@ public abstract class AbstractEntityPersister
 			TableGroup tableGroup,
 			SqlAstCreationState creationState) {
 		if ( needsDiscriminator() ) {
-			predicateConsumer.accept( createDiscriminatorPredicate( alias, tableGroup, creationState ) );
+			pruneForSubclasses( tableGroup, Collections.singletonMap( getEntityName(), EntityNameUse.TREAT ) );
 		}
 	}
 
diff --git a/hibernate-core/src/main/java/org/hibernate/query/sqm/sql/BaseSqmToSqlAstConverter.java b/hibernate-core/src/main/java/org/hibernate/query/sqm/sql/BaseSqmToSqlAstConverter.java
index 6e9ae8194a..467bdf3a71 100644
--- a/hibernate-core/src/main/java/org/hibernate/query/sqm/sql/BaseSqmToSqlAstConverter.java
+++ b/hibernate-core/src/main/java/org/hibernate/query/sqm/sql/BaseSqmToSqlAstConverter.java
@@ -3042,6 +3042,11 @@ public abstract class BaseSqmToSqlAstConverter<T extends Statement> extends Base
 		}
 	}
 
+	@Override
+	public boolean supportsEntityNameUsage() {
+		return true;
+	}
+
 	@Override
 	public void registerEntityNameUsage(
 			TableGroup tableGroup,
diff --git a/hibernate-core/src/main/java/org/hibernate/sql/ast/spi/SqlAstCreationState.java b/hibernate-core/src/main/java/org/hibernate/sql/ast/spi/SqlAstCreationState.java
index 89a46ff154..5a599ad10e 100644
--- a/hibernate-core/src/main/java/org/hibernate/sql/ast/spi/SqlAstCreationState.java
+++ b/hibernate-core/src/main/java/org/hibernate/sql/ast/spi/SqlAstCreationState.java
@@ -42,4 +42,9 @@ public interface SqlAstCreationState {
 			String hibernateEntityName) {
 		// No-op
 	}
+
+	@Internal
+	default boolean supportsEntityNameUsage() {
+		return false;
+	}
 }
diff --git a/hibernate-core/src/main/java/org/hibernate/sql/ast/tree/from/NamedTableReference.java b/hibernate-core/src/main/java/org/hibernate/sql/ast/tree/from/NamedTableReference.java
index 33e4dc2f62..d78c1a3c7e 100644
--- a/hibernate-core/src/main/java/org/hibernate/sql/ast/tree/from/NamedTableReference.java
+++ b/hibernate-core/src/main/java/org/hibernate/sql/ast/tree/from/NamedTableReference.java
@@ -72,7 +72,7 @@ public class NamedTableReference extends AbstractTableReference {
 
 	@Override
 	public boolean containsAffectedTableName(String requestedName) {
-		return isEmpty( requestedName ) || getTableExpression().equals( requestedName );
+		return isEmpty( requestedName ) || getTableExpression().contains( requestedName );
 	}
 
 	@Override