HHH-19925 - Locking root(s) should be based on select-clause, not from-clause

hibernate-core/src/main/java/org/hibernate/dialect/lock/internal/SqlAstBasedLockingStrategy.java

@@ -154,7 +154,7 @@ public void lock(
 			);
 		}
 
-		final var selectStatement = new SelectStatement( rootQuerySpec, List.of( idResult ) );
+		final var selectStatement = new SelectStatement( rootQuerySpec, List.of( idResult ), List.of( entityPath ) );
 		final JdbcSelect selectOperation =
 				session.getDialect().getSqlAstTranslatorFactory()
 						.buildSelectTranslator( factory, selectStatement )

hibernate-core/src/main/java/org/hibernate/loader/ast/internal/AbstractNaturalIdLoader.java

@@ -215,7 +215,7 @@ public Object resolveNaturalIdToId(Object naturalIdValue, SharedSessionContractI
 		return executeNaturalIdQuery(
 				naturalIdValue,
 				LockOptions.NONE,
-				new SelectStatement( rootQuerySpec, singletonList( domainResult ) ),
+				new SelectStatement( rootQuerySpec, singletonList( domainResult ), List.of() ),
 				rootTableGroup,
 				rootQuerySpec::applyPredicate,
 				sqlAstCreationState,

hibernate-core/src/main/java/org/hibernate/loader/ast/internal/DatabaseSnapshotExecutor.java

@@ -128,7 +128,7 @@ class DatabaseSnapshotExecutor {
 				}
 		);
 
-		final var selectStatement = new SelectStatement( rootQuerySpec, domainResults );
+		final var selectStatement = new SelectStatement( rootQuerySpec, domainResults, List.of() );
 		jdbcSelect =
 				sessionFactory.getJdbcServices().getJdbcEnvironment().getSqlAstTranslatorFactory()
 						.buildSelectTranslator( sessionFactory, selectStatement )

hibernate-core/src/main/java/org/hibernate/loader/ast/internal/LoaderSelectBuilder.java

@@ -177,7 +177,7 @@ public static SelectStatement createSelectBySingleArrayParameter(
 			builder.applyFiltering( rootQuerySpec, rootTableGroup, (Restrictable) loadable, sqlAstCreationState );
 		}
 
-		return new SelectStatement( rootQuerySpec, domainResults );
+		return new SelectStatement( rootQuerySpec, domainResults, List.of( rootNavigablePath ) );
 	}
 
 	private static void applyArrayParamRestriction(
@@ -505,7 +505,7 @@ else if ( cachedDomainResult != null ) {
 			applyFiltering( rootQuerySpec, rootTableGroup, (Restrictable) loadable, sqlAstCreationState );
 		}
 
-		return new SelectStatement( rootQuerySpec, domainResults );
+		return new SelectStatement( rootQuerySpec, domainResults, List.of( rootNavigablePath ) );
 	}
 
 	private List<DomainResult<?>> buildRequestedDomainResults(
@@ -1024,7 +1024,8 @@ private SelectStatement generateSelect(SubselectFetch subselect) {
 								rootTableGroup,
 								sqlAstCreationState
 						)
-				)
+				),
+				singletonList( rootNavigablePath )
 		);
 	}
 

hibernate-core/src/main/java/org/hibernate/query/sqm/mutation/internal/MatchingIdSelectionHelper.java

@@ -65,15 +65,11 @@
 public class MatchingIdSelectionHelper {
 	private static final Logger LOG = Logger.getLogger( MatchingIdSelectionHelper.class );
 
-	/**
-	 * @asciidoc
-	 *
-	 * Generates a query-spec for selecting all ids matching the restriction defined as part
-	 * of the user's update/delete query.  This query-spec is generally used:
-	 *
-	 * 		* to select all the matching ids via JDBC - see {@link MatchingIdSelectionHelper#selectMatchingIds}
-	 * 		* as a sub-query restriction to insert rows into an "id table"
-	 */
+	/// Generates a query-spec for selecting all ids matching the restriction defined as part
+	/// of the user's update/delete query.  This query-spec is generally used:
+	///
+	///  * to select all the matching ids via JDBC - see {@link MatchingIdSelectionHelper#selectMatchingIds}
+	///  * as a sub-query restriction to insert rows into an "id table"
 	public static SelectStatement generateMatchingIdSelectStatement(
 			EntityMappingType targetEntityDescriptor,
 			SqmDeleteOrUpdateStatement<?> sqmStatement,
@@ -109,14 +105,12 @@ public static SelectStatement generateMatchingIdSelectStatement(
 				mutatingTableGroup.getNavigablePath(),
 				mutatingTableGroup,
 				sqmConverter,
-				(selection, jdbcMapping) ->
-						domainResults.add(
-								new BasicResult<>(
-										selection.getValuesArrayPosition(),
-										null,
-										jdbcMapping
-								)
-						)
+				(selection, jdbcMapping) -> domainResults.add(
+					new BasicResult<>(
+							selection.getValuesArrayPosition(),
+							null,
+							jdbcMapping
+					) )
 		);
 		sqmConverter.getProcessingStateStack().pop();
 
@@ -130,18 +124,14 @@ public static SelectStatement generateMatchingIdSelectStatement(
 				sqmConverter
 		);
 
-		return new SelectStatement( idSelectionQuery, domainResults );
+		return new SelectStatement( idSelectionQuery, domainResults, List.of() );
 	}
 
-	/**
-	 * @asciidoc
-	 *
-	 * Generates a query-spec for selecting all ids matching the restriction defined as part
-	 * of the user's update/delete query.  This query-spec is generally used:
-	 *
-	 * 		* to select all the matching ids via JDBC - see {@link MatchingIdSelectionHelper#selectMatchingIds}
-	 * 		* as a sub-query restriction to insert rows into an "id table"
-	 */
+	/// Generates a query-spec for selecting all ids matching the restriction defined as part
+	/// of the user's update/delete query.  This query-spec is generally used:
+	///
+	/// 	* to select all the matching ids via JDBC - see {@link MatchingIdSelectionHelper#selectMatchingIds}
+	/// 	* as a sub-query restriction to insert rows into an "id table"
 	public static SqmSelectStatement<?> generateMatchingIdSelectStatement(
 			SqmDeleteOrUpdateStatement<?> sqmStatement,
 			EntityMappingType entityDescriptor) {
@@ -167,59 +157,6 @@ public static SqmSelectStatement<?> generateMatchingIdSelectStatement(
 				nodeBuilder
 		);
 	}
-//
-//	/**
-//	 * @asciidoc
-//	 *
-//	 * Generates a query-spec for selecting all ids matching the restriction defined as part
-//	 * of the user's update/delete query.  This query-spec is generally used:
-//	 *
-//	 * 		* to select all the matching ids via JDBC - see {@link MatchingIdSelectionHelper#selectMatchingIds}
-//	 * 		* as a sub-query restriction to insert rows into an "id table"
-//	 */
-//	public static QuerySpec generateMatchingIdSelectQuery(
-//			EntityMappingType targetEntityDescriptor,
-//			SqmDeleteOrUpdateStatement sqmStatement,
-//			DomainParameterXref domainParameterXref,
-//			Predicate restriction,
-//			MultiTableSqmMutationConverter sqmConverter,
-//			SessionFactoryImplementor sessionFactory) {
-//		final EntityDomainType entityDomainType = sqmStatement.getTarget().getModel();
-//		if ( LOG.isTraceEnabled() ) {
-//			LOG.tracef(
-//					"Starting generation of entity-id SQM selection - %s",
-//					entityDomainType.getHibernateEntityName()
-//			);
-//		}
-//
-//		final QuerySpec idSelectionQuery = new QuerySpec( true, 1 );
-//
-//		final TableGroup mutatingTableGroup = sqmConverter.getMutatingTableGroup();
-//		idSelectionQuery.getFromClause().addRoot( mutatingTableGroup );
-//
-//		targetEntityDescriptor.getIdentifierMapping().forEachSelectable(
-//				(position, selection) -> {
-//					final TableReference tableReference = mutatingTableGroup.resolveTableReference(
-//							mutatingTableGroup.getNavigablePath(),
-//							selection.getContainingTableExpression()
-//					);
-//					final Expression expression = sqmConverter.getSqlExpressionResolver().resolveSqlExpression(
-//							tableReference,
-//							selection
-//					);
-//					idSelectionQuery.getSelectClause().addSqlSelection(
-//							new SqlSelectionImpl(
-//									position,
-//									expression
-//							)
-//					);
-//				}
-//		);
-//
-//		idSelectionQuery.applyPredicate( restriction );
-//
-//		return idSelectionQuery;
-//	}
 
 	/**
 	 * Centralized selection of ids matching the restriction of the DELETE

hibernate-core/src/main/java/org/hibernate/query/sqm/mutation/internal/cte/AbstractCteMutationHandler.java

@@ -138,7 +138,7 @@ public AbstractCteMutationHandler(
 		// Create the main query spec that will return the count of
 		final QuerySpec querySpec = new QuerySpec( true, 1 );
 		final List<DomainResult<?>> domainResults = new ArrayList<>( 1 );
-		final SelectStatement statement = new SelectStatement( querySpec, domainResults );
+		final SelectStatement statement = new SelectStatement( querySpec, domainResults, List.of() );
 		final JdbcServices jdbcServices = factory.getJdbcServices();
 		final SqlAstTranslator<JdbcOperationQuerySelect> translator = jdbcServices.getJdbcEnvironment()
 				.getSqlAstTranslatorFactory()

hibernate-core/src/main/java/org/hibernate/query/sqm/mutation/internal/cte/CteInsertHandler.java

@@ -291,7 +291,7 @@ public CteInsertHandler(
 		// Create the main query spec that will return the count of rows
 		final QuerySpec querySpec = new QuerySpec( true, 1 );
 		final List<DomainResult<?>> domainResults = new ArrayList<>( 1 );
-		final SelectStatement statement = new SelectStatement( querySpec, domainResults );
+		final SelectStatement statement = new SelectStatement( querySpec, domainResults, List.of() );
 
 		final CteStatement entityCte;
 		if ( additionalInsertValues.requiresRowNumberIntermediate() ) {

hibernate-core/src/main/java/org/hibernate/query/sqm/mutation/internal/temptable/TableBasedInsertHandler.java

@@ -593,7 +593,8 @@ private RootTableInserter createRootTableInserter(
 									null,
 									false
 							)
-					)
+					),
+					List.of()
 			);
 			temporaryTableIdentitySelect = jdbcServices.getJdbcEnvironment()
 					.getSqlAstTranslatorFactory()

hibernate-core/src/main/java/org/hibernate/query/sqm/sql/BaseSqmToSqlAstConverter.java

@@ -1640,11 +1640,12 @@
 		final QueryPart queryPart = visitQueryPart( statement.getQueryPart() );
 		final List<DomainResult<?>> domainResults = queryPart.isRoot() ? this.domainResults : emptyList();
 		try {
-			return new SelectStatement( cteContainer, queryPart, domainResults );
+			return new SelectStatement( cteContainer, queryPart, domainResults, rootPathsForLocking );
 		}
 		finally {
 			this.currentSqmStatement = oldSqmStatement;
 			this.cteContainer = oldCteContainer;
+			rootPathsForLocking = null;
 		}
 	}
 
@@ -1753,7 +1754,7 @@
 
 								final CteStatement cteStatement = new CteStatement(
 										cteTable,
-										new SelectStatement( subCteContainer, group, emptyList() ),
+										new SelectStatement( subCteContainer, group, emptyList(), emptyList() ),
 										sqmCteStatement.getMaterialization(),
 										sqmCteStatement.getSearchClauseKind(),
 										visitSearchBySpecifications( cteTable, sqmCteStatement.getSearchBySpecifications() ),
@@ -2205,11 +2206,22 @@
 		return getFromClauseAccess().getTableGroup( navigablePath );
 	}
 
+	private List<NavigablePath> rootPathsForLocking;
+	private Consumer<NavigablePath> rootPathsForLockingCollector;
+
 	@Override
 	public SelectClause visitSelectClause(SqmSelectClause selectClause) {
 		currentClauseStack.push( Clause.SELECT );
 		try {
 			final SelectClause sqlSelectClause = currentQuerySpec().getSelectClause();
+			if ( sqmQueryPartStack.depth() == 1 ) {
+				rootPathsForLockingCollector = (navigablePath) -> {
+					if ( rootPathsForLocking == null ) {
+						rootPathsForLocking = new ArrayList<>();
+					}
+					rootPathsForLocking.add(  navigablePath );
+				};
+			}
 			if ( selectClause == null ) {
 				final SqmFrom<?, ?> implicitSelection = determineImplicitSelection( (SqmQuerySpec<?>) getCurrentSqmQueryPart() );
 				visitSelection( 0, new SqmSelection<>( implicitSelection, implicitSelection.nodeBuilder() ) );
@@ -2224,6 +2236,7 @@
 			return sqlSelectClause;
 		}
 		finally {
+			rootPathsForLockingCollector = null;
 			currentClauseStack.pop();
 		}
 	}
@@ -2243,6 +2256,8 @@
 	}
 
 	private void visitSelection(int index, SqmSelection<?> sqmSelection) {
+		collectRootPathsForLocking( sqmSelection );
+
 		inferTargetPath( index );
 		callResultProducers( resultProducers( sqmSelection ) );
 		if ( statement instanceof SqmInsertSelectStatement<?>
@@ -2251,6 +2266,52 @@
 		}
 	}
 
+	private void collectRootPathsForLocking(SqmSelection<?> sqmSelection) {
+		if ( rootPathsForLockingCollector == null ) {
+			return;
+		}
+
+		collectRootPathsForLocking( sqmSelection.getSelectableNode() );
+	}
+
+	private void collectRootPathsForLocking(SqmSelectableNode<?> selectableNode) {
+		// roughly speaking we only care about 2 cases here:
+		//		1) entity path - the entity will be locked
+		//		2) scalar path - the entity from which the path originates will be locked
+		//
+		// note, however, that we need to account for both cases as the argument to a dynamic instantiation
+
+		if ( selectableNode instanceof SqmPath<?> selectedPath ) {
+			collectRootPathsForLocking( selectedPath );
+		}
+		else if ( selectableNode instanceof SqmDynamicInstantiation<?> dynamicInstantiation ) {
+			collectRootPathsForLocking( dynamicInstantiation );
+		}
+	}
+
+	private void collectRootPathsForLocking(SqmPath<?> selectedPath) {
+		assert rootPathsForLockingCollector != null;
+
+		if ( selectedPath == null ) {
+			// typically this comes from paths rooted in a CTE.
+			// regardless, without a path we cannot evaluate so just return.
+			return;
+		}
+
+		if ( selectedPath.getNodeType() instanceof EntityTypeImpl ) {
+			rootPathsForLockingCollector.accept( selectedPath.getNavigablePath() );
+		}
+		else {
+			collectRootPathsForLocking( selectedPath.getLhs() );
+		}
+	}
+
+	private void collectRootPathsForLocking(SqmDynamicInstantiation<?> dynamicInstantiation) {
+		dynamicInstantiation.getArguments().forEach( ( argument ) -> {
+			collectRootPathsForLocking( argument.getSelectableNode() );
+		} );
+	}
+
 	private void inferTargetPath(int index) {
 		// Only infer the type on the "top level" select clauses
 		// todo: add WriteExpression handling
@@ -7129,7 +7190,7 @@
 		final QueryPart queryPart = visitQueryPart( sqmSubQuery.getQueryPart() );
 		currentlyProcessingJoin = oldJoin;
 		this.cteContainer = oldCteContainer;
-		return new SelectStatement( cteContainer, queryPart, emptyList() );
+		return new SelectStatement( cteContainer, queryPart, emptyList(), emptyList() );
 	}
 
 	@Override

hibernate-core/src/main/java/org/hibernate/sql/ast/spi/AbstractSqlAstTranslator.java

@@ -6492,6 +6492,7 @@ protected Predicate determineLateralEmulationPredicate(TableGroup tableGroup) {
 						new SelectStatement(
 								statement,
 								new QueryGroup( false, SetOperator.INTERSECT, queryParts ),
+								emptyList(),
 								emptyList()
 						),
 						false,
@@ -6540,7 +6541,7 @@ protected Predicate determineLateralEmulationPredicate(TableGroup tableGroup) {
 				);
 
 				return new ExistsPredicate(
-					new SelectStatement( statement, existsQuery, emptyList() ),
+					new SelectStatement( statement, existsQuery, emptyList(), emptyList() ),
 					false,
 						booleanType
 				);
@@ -6583,7 +6584,7 @@ protected Predicate determineLateralEmulationPredicate(TableGroup tableGroup) {
 			);
 
 			final ExistsPredicate existsPredicate = new ExistsPredicate(
-					new SelectStatement( statement, existsQuery, emptyList() ),
+					new SelectStatement( statement, existsQuery, emptyList(), emptyList() ),
 					false,
 					booleanType
 			);
@@ -6701,7 +6702,7 @@ protected Predicate determineLateralEmulationPredicate(TableGroup tableGroup) {
 					List.of(
 							existsPredicate,
 							new BetweenPredicate(
-									new SelectStatement( statement, countQuery, emptyList() ),
+									new SelectStatement( statement, countQuery, emptyList(), emptyList() ),
 									countLower,
 									countUpper,
 									false,
@@ -6764,6 +6765,7 @@ private SelectStatement stripToSelectClause(SelectStatement statement) {
 		return new SelectStatement(
 				statement,
 				stripToSelectClause( statement.getQueryPart() ),
+				emptyList(),
 				emptyList()
 		);
 	}