HHH-19602 - Adjust JdbcOperation to allow more-than-one statement

Author: sebersole

hibernate-core/src/main/java/org/hibernate/sql/exec/internal/AbstractDatabaseOperation.java

@@ -7,6 +7,7 @@
 import org.hibernate.internal.util.collections.CollectionHelper;
 import org.hibernate.sql.exec.spi.ExecutionContext;
 import org.hibernate.sql.exec.spi.DatabaseOperation;
+import org.hibernate.sql.exec.spi.JdbcOperation;
 import org.hibernate.sql.exec.spi.PostAction;
 import org.hibernate.sql.exec.spi.PreAction;
 import org.hibernate.sql.exec.spi.SecondaryAction;
@@ -18,13 +19,18 @@
 import java.util.List;
 
 /**
+ * Abstract support for DatabaseOperation implementations, mainly
+ * managing {@linkplain PreAction pre-} and {@linkplain PostAction post-}
+ * actions.
+ *
  * @author Steve Ebersole
  */
-@SuppressWarnings("ALL")
-public abstract class AbstractDatabaseOperation implements DatabaseOperation {
+public abstract class AbstractDatabaseOperation<P extends JdbcOperation>
+		implements DatabaseOperation<P> {
 	protected final PreAction[] preActions;
 	protected final PostAction[] postActions;
 
+	@SuppressWarnings("unused")
 	public AbstractDatabaseOperation() {
 		this( null, null );
 	}
@@ -77,7 +83,8 @@ protected abstract static class Builder<T extends Builder<T>> {
 		protected abstract T getThis();
 
 		/**
-		 * Appends the {@code actions} to the growing list of pre-actions
+		 * Appends the {@code actions} to the growing list of pre-actions,
+		 * executed (in order) after all currently registered actions.
 		 *
 		 * @return {@code this}, for method chaining.
 		 */
@@ -98,9 +105,8 @@ public T prependPreAction(PreAction... actions) {
 			if ( preActions == null ) {
 				preActions = new ArrayList<>();
 			}
-			for ( int i = actions.length - 1; i >= 0; i-- ) {
-				preActions.add( 0, actions[i] );
-			}
+			// todo (DatabaseOperation) : should we invert the order of the incoming actions?
+			Collections.addAll( preActions, actions );
 			return getThis();
 		}
 
@@ -126,15 +132,17 @@ public T prependPostAction(PostAction... actions) {
 			if ( postActions == null ) {
 				postActions = new ArrayList<>();
 			}
-			for ( int i = actions.length - 1; i >= 0; i-- ) {
-				postActions.add( 0, actions[i] );
-			}
+			// todo (DatabaseOperation) : should we invert the order of the incoming actions?
+			Collections.addAll( postActions, actions );
 			return getThis();
 		}
 
 		/**
-		 * Adds a secondary action.  Assumes the action implements both
-		 * {@linkplain PreAction} and {@linkplain PostAction}.
+		 * Adds a secondary action pair.
+		 * Assumes the {@code action} implements both {@linkplain PreAction} and {@linkplain PostAction}.
+		 *
+		 * @apiNote Prefer {@linkplain #addSecondaryActionPair(PreAction, PostAction)} to avoid
+		 * the casts needed here.
 		 *
 		 * @see #prependPreAction
 		 * @see #appendPostAction

hibernate-core/src/main/java/org/hibernate/sql/exec/internal/DatabaseOperationSelectImpl.java

@@ -27,7 +27,7 @@
  * @author Steve Ebersole
  */
 public class DatabaseOperationSelectImpl
-		extends AbstractDatabaseOperation
+		extends AbstractDatabaseOperation<JdbcOperationQuerySelect>
 		implements DatabaseOperationSelect {
 	private final JdbcOperationQuerySelect primaryOperation;
 

hibernate-core/src/main/java/org/hibernate/sql/exec/spi/DatabaseOperation.java

@@ -9,18 +9,18 @@
 import java.util.Set;
 
 /**
- * An operation against the database, comprised of a single
- * {@linkplain #getPrimaryOperation primary operation} and zero-or-more
- * before/after {@linkplain SecondaryAction secondary actions}.
+ * An operation against the database, comprised of a
+ * {@linkplain #getPrimaryOperation primary operation} and
+ * zero-or-more {@linkplain SecondaryAction secondary actions}.
  *
  * @author Steve Ebersole
  */
 @Incubating
-public interface DatabaseOperation {
+public interface DatabaseOperation<P extends JdbcOperation> {
 	/**
 	 * The primary operation for the group.
 	 */
-	JdbcOperation getPrimaryOperation();
+	P getPrimaryOperation();
 
 	/**
 	 * The names of tables referenced or affected by this operation.

hibernate-core/src/main/java/org/hibernate/sql/exec/spi/DatabaseOperationSelect.java

@@ -12,15 +12,12 @@
 import java.sql.PreparedStatement;
 
 /**
- * {@linkplain DatabaseOperation} whose primary operation is a selection.
+ * {@linkplain DatabaseOperation} whose primary operation is a {@linkplain JdbcOperationQuerySelect selection}.
  *
  * @author Steve Ebersole
  */
 @Incubating
-public interface DatabaseOperationSelect extends DatabaseOperation {
-	@Override
-	JdbcOperationQuerySelect getPrimaryOperation();
-
+public interface DatabaseOperationSelect extends DatabaseOperation<JdbcOperationQuerySelect> {
 	/**
 	 * Execute the underlying statements and return the result(s).
 	 *

hibernate-core/src/test/java/org/hibernate/orm/test/sql/exec/spi/DatabaseOperationSmokeTest.java

@@ -23,6 +23,7 @@
 import org.hibernate.sql.ast.tree.expression.JdbcParameter;
 import org.hibernate.sql.ast.tree.select.SelectStatement;
 import org.hibernate.sql.exec.internal.BaseExecutionContext;
+import org.hibernate.sql.exec.internal.CallbackImpl;
 import org.hibernate.sql.exec.internal.JdbcParameterBindingImpl;
 import org.hibernate.sql.exec.internal.JdbcParameterBindingsImpl;
 import org.hibernate.sql.exec.internal.StandardStatementCreator;
@@ -35,14 +36,20 @@
 import org.hibernate.sql.exec.spi.PreAction;
 import org.hibernate.sql.exec.spi.StatementAccess;
 import org.hibernate.sql.results.spi.SingleResultConsumer;
+import org.hibernate.testing.orm.junit.DialectFeatureChecks;
 import org.hibernate.testing.orm.junit.DomainModel;
+import org.hibernate.testing.orm.junit.RequiresDialectFeature;
 import org.hibernate.testing.orm.junit.SessionFactory;
 import org.hibernate.testing.orm.junit.SessionFactoryScope;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import java.sql.Connection;
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.assertj.core.api.Assertions.assertThat;
 
 
 /**
@@ -96,14 +103,13 @@ void testSimpleSelect(SessionFactoryScope factoryScope) {
 	}
 
 	@Test
+	@RequiresDialectFeature( feature = DialectFeatureChecks.SupportsConnectionLockTimeouts.class)
 	void testConnectionLockTimeout(SessionFactoryScope factoryScope) {
 		final SessionFactoryImplementor sessionFactory = factoryScope.getSessionFactory();
 
 		final LockingSupport lockingSupport = sessionFactory.getJdbcServices().getDialect().getLockingSupport();
 		final ConnectionLockTimeoutStrategy lockTimeoutStrategy = lockingSupport.getConnectionLockTimeoutStrategy();
-		if ( lockTimeoutStrategy.getSupportedLevel() == ConnectionLockTimeoutStrategy.Level.NONE ) {
-			return;
-		}
+		assert lockTimeoutStrategy.getSupportedLevel() != ConnectionLockTimeoutStrategy.Level.NONE;
 
 		final EntityPersister entityDescriptor = sessionFactory.getMappingMetamodel().findEntityDescriptor( Person.class );
 
@@ -138,6 +144,66 @@ void testConnectionLockTimeout(SessionFactoryScope factoryScope) {
 		} );
 	}
 
+	@Test
+	void testFollowOnLockingParadigm(SessionFactoryScope factoryScope) {
+		// NOTE: this just tests the principle -
+		// for now, just collect the values loaded.
+		// ultimately, this can be used to apply smarter follow-on locking.
+
+		final SessionFactoryImplementor sessionFactory = factoryScope.getSessionFactory();
+		final EntityPersister entityDescriptor = sessionFactory.getMappingMetamodel().findEntityDescriptor( Person.class );
+
+		final PersonQuery personQuery = createPersonQuery( entityDescriptor, sessionFactory );
+		final JdbcOperationQuerySelect jdbcOperation = personQuery.jdbcOperation();
+		final JdbcParameterBindings jdbcParameterBindings = personQuery.jdbcParameterBindings();
+
+		factoryScope.inTransaction( (session) -> {
+			final LoadedValueCollector loadedValueCollector = new LoadedValueCollector();
+
+			final Callback callback = new CallbackImpl();
+			callback.registerAfterLoadAction( (entity, entityMappingType, session1) ->  {
+				loadedValueCollector.loadedValues.add( entity );
+			} );
+
+			final SingleIdExecutionContext executionContext = new SingleIdExecutionContext(
+					session,
+					null,
+					1,
+					entityDescriptor,
+					QueryOptions.NONE,
+					callback
+			);
+
+
+			final DatabaseOperationSelectImpl.Builder operationBuilder = DatabaseOperationSelectImpl
+					.builder( jdbcOperation )
+					.appendPostAction( loadedValueCollector );
+
+			final ConnectionLockTimeoutStrategy lockTimeoutStrategy = session
+					.getDialect()
+					.getLockingSupport()
+					.getConnectionLockTimeoutStrategy();
+			if ( lockTimeoutStrategy.getSupportedLevel() != ConnectionLockTimeoutStrategy.Level.NONE ) {
+				final LockTimeoutHandler lockTimeoutHandler = new LockTimeoutHandler( Timeout.seconds( 2 ), lockTimeoutStrategy );
+				operationBuilder.addSecondaryActionPair( lockTimeoutHandler, lockTimeoutHandler );
+			}
+
+			final DatabaseOperationSelectImpl databaseOperation = operationBuilder.build();
+			final Person person = databaseOperation.execute(
+					Person.class,
+					1,
+					StandardStatementCreator.getStatementCreator( ScrollMode.FORWARD_ONLY ),
+					jdbcParameterBindings,
+					row -> (Person) row[0],
+					SingleResultConsumer.instance(),
+					executionContext
+			);
+
+			assertThat( loadedValueCollector.loadedValues ).hasSize( 1 );
+		} );
+
+	}
+
 	private static class LockTimeoutHandler implements PreAction, PostAction {
 		private final ConnectionLockTimeoutStrategy lockTimeoutStrategy;
 		private final Timeout timeout;
@@ -172,6 +238,17 @@ public void performPostAction(StatementAccess jdbcStatementAccess, Connection jd
 		}
 	}
 
+	private static class LoadedValueCollector implements PostAction {
+		private final List<Object> loadedValues = new ArrayList<>();
+
+		@Override
+		public void performPostAction(StatementAccess jdbcStatementAccess, Connection jdbcConnection, ExecutionContext executionContext) {
+			loadedValues.forEach( (value) -> {
+				System.out.printf( "Loaded value: %s\n", value );
+			} );
+		}
+	}
+
 
 	private PersonQuery createPersonQuery(
 			EntityPersister entityDescriptor,