proof of concept for SchemaManager.resynchronizeSequences()

Author: gavinking

hibernate-core/src/main/java/org/hibernate/boot/model/relational/InitCommand.java

@@ -4,6 +4,8 @@
  */
 package org.hibernate.boot.model.relational;
 
+import java.util.Arrays;
+
 /**
  * A general SQL command to be used while initializing a schema.
  *
@@ -14,4 +16,15 @@ public record InitCommand(String... initCommands) {
 	public String[] getInitCommands() {
 		return initCommands;
 	}
+
+	@Override
+	public boolean equals(Object object) {
+		return object instanceof InitCommand that
+			&& Arrays.equals( this.initCommands, that.initCommands );
+	}
+
+	@Override
+	public int hashCode() {
+		return Arrays.hashCode( initCommands );
+	}
 }

hibernate-core/src/main/java/org/hibernate/generator/GeneratorCreationContext.java

@@ -12,6 +12,7 @@
 import org.hibernate.mapping.PersistentClass;
 import org.hibernate.mapping.Property;
 import org.hibernate.mapping.RootClass;
+import org.hibernate.mapping.Value;
 import org.hibernate.service.ServiceRegistry;
 import org.hibernate.type.Type;
 
@@ -64,6 +65,11 @@ public interface GeneratorCreationContext {
 	 */
 	Property getProperty();
 
+	/**
+	 * The identifier.
+	 */
+	Value getValue();
+
 	/**
 	 * Mapping details for the identifier type.
 	 */

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

@@ -9,6 +9,7 @@
 import org.hibernate.boot.model.relational.QualifiedName;
 import org.hibernate.boot.model.relational.SqlStringGenerationContext;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.mapping.Table;
 
 /**
  * Encapsulates definition of the underlying data structure backing a
@@ -88,6 +89,21 @@ default void configure(Optimizer optimizer) {
 	@Override
 	void registerExportables(Database database);
 
+	/**
+	 * Register additional database objects which need to be aware of the
+	 * table for which this structure is used to generate values. Used to
+	 * deal with automatic sequence resynchronization after data import.
+	 *
+	 * @param table The table for which this structure is used to generate values
+	 * @param database The database instance
+	 *
+	 * @see org.hibernate.relational.SchemaManager#resynchronizeSequences()
+	 *
+	 * @since 7.2
+	 */
+	default void registerExtraExportables(Table table, Database database) {
+	}
+
 	/**
 	 * Initializes this structure, in particular pre-generates SQL as necessary.
 	 * <p>

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

@@ -8,11 +8,13 @@
 
 import org.hibernate.AssertionFailure;
 import org.hibernate.boot.model.relational.Database;
+import org.hibernate.boot.model.relational.InitCommand;
 import org.hibernate.boot.model.relational.QualifiedName;
 import org.hibernate.boot.model.relational.Sequence;
 import org.hibernate.boot.model.relational.SqlStringGenerationContext;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.id.IntegralDataTypeHolder;
+import org.hibernate.mapping.Table;
 
 import static org.hibernate.engine.jdbc.JdbcLogging.JDBC_LOGGER;
 import static org.hibernate.id.IdentifierGeneratorHelper.getIntegralDataTypeHolder;
@@ -164,6 +166,20 @@ public void initialize(SqlStringGenerationContext context) {
 				.getSequenceNextValString( context.format( physicalSequenceName ) );
 	}
 
+	@Override
+	public void registerExtraExportables(Table table, Database database) {
+		table.addResyncCommand( context -> {
+			final String sequenceName = context.format( physicalSequenceName );
+			final String tableName = context.format( table.getQualifiedTableName() );
+			final String primaryKeyColumnName = table.getPrimaryKey().getColumn( 0 ).getName();
+			final String sequenceCurrentValue =
+					database.getDialect().getSequenceSupport()
+							.getSelectSequenceNextValString( sequenceName );
+			return new InitCommand( "alter sequence %s restart with select greatest(max(%s)+%d,%s) from %s"
+					.formatted( sequenceName, primaryKeyColumnName, incrementSize, sequenceCurrentValue, tableName ) );
+		} );
+	}
+
 	@Override
 	public boolean isPhysicalSequence() {
 		return true;

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

@@ -24,6 +24,7 @@
 import org.hibernate.id.IdentifierGenerator;
 import org.hibernate.id.PersistentIdentifierGenerator;
 import org.hibernate.id.SequenceMismatchStrategy;
+import org.hibernate.mapping.Table;
 import org.hibernate.service.ServiceRegistry;
 import org.hibernate.tool.schema.Action;
 import org.hibernate.tool.schema.extract.spi.SequenceInformation;
@@ -151,6 +152,7 @@ public class SequenceStyleGenerator
 	private DatabaseStructure databaseStructure;
 	private Optimizer optimizer;
 	private Type identifierType;
+	private Table table;
 
 	/**
 	 * Getter for property 'databaseStructure'.
@@ -190,6 +192,7 @@ public void configure(GeneratorCreationContext creationContext, Properties param
 		final var dialect = jdbcEnvironment.getDialect();
 
 		identifierType = creationContext.getType();
+		table = creationContext.getValue().getTable();
 
 		final var sequenceName = determineSequenceName( parameters, jdbcEnvironment, serviceRegistry );
 		final int initialValue = determineInitialValue( parameters );
@@ -329,6 +332,7 @@ private boolean isSchemaToBeRecreated(String contributor, ConfigurationService c
 	@Override
 	public void registerExportables(Database database) {
 		databaseStructure.registerExportables( database );
+		databaseStructure.registerExtraExportables( table, database );
 	}
 
 	@Override
@@ -380,7 +384,7 @@ private static QualifiedName sequenceName(
 
 	/**
 	 * Determine the name of the column used to store the generator value in
-	 * the db.
+	 * the database.
 	 * <p>
 	 * Called during {@linkplain #configure configuration} <b>when resolving to a
 	 * physical table</b>.
@@ -409,8 +413,8 @@ protected int determineInitialValue(Properties params) {
 	}
 
 	/**
-	 * Determine the increment size to be applied.  The exact implications of
-	 * this value depends on the {@linkplain #getOptimizer() optimizer} being used.
+	 * Determine the increment size to be applied. The exact implications of
+	 * this value depend on the {@linkplain #getOptimizer() optimizer} in use.
 	 * <p>
 	 * Called during {@linkplain #configure configuration}.
 	 *
@@ -438,7 +442,7 @@ protected OptimizerDescriptor determineOptimizationStrategy(Properties params, i
 
 	/**
 	 * In certain cases we need to adjust the increment size based on the
-	 * selected optimizer.  This is the hook to achieve that.
+	 * selected optimizer. This is the hook to achieve that.
 	 *
 	 * @param optimizationStrategy The optimizer strategy (name)
 	 * @param incrementSize The {@link #determineIncrementSize determined increment size}

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

@@ -333,4 +333,16 @@ public void initialize(SqlStringGenerationContext context) {
 				" set " + valueColumnNameText + "= ?" +
 				" where " + valueColumnNameText + "=?";
 	}
+
+	@Override
+	public void registerExtraExportables(Table table, Database database) {
+		table.addResyncCommand( context -> {
+			final String sequenceTableName = context.format( physicalTableName );
+			final String tableName = context.format( table.getQualifiedTableName() );
+			final String primaryKeyColumnName = table.getPrimaryKey().getColumn( 0 ).getName();
+			return new InitCommand( "update %s set %s = select greatest(max(%s)+%d,(select %s from %s)) from %s"
+					.formatted( sequenceTableName, valueColumnNameText, primaryKeyColumnName,
+							incrementSize, valueColumnNameText, sequenceTableName, tableName ) );
+		} );
+	}
 }

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

@@ -602,6 +602,11 @@ public Property getProperty() {
 			return Property.this;
 		}
 
+		@Override
+		public Value getValue() {
+			return value;
+		}
+
 		@Override
 		public SqlStringGenerationContext getSqlStringGenerationContext() {
 			return context.getSqlStringGenerationContext();

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

@@ -422,7 +422,7 @@ public Generator createGenerator(
 			Property property,
 			GeneratorSettings defaults) {
 		if ( customIdGeneratorCreator != null ) {
-			final var context = new IdGeneratorCreationContext( rootClass, property, defaults );
+			final var context = new IdGeneratorCreationContext( this, rootClass, property, defaults );
 			final var generator = customIdGeneratorCreator.createGenerator( context );
 			if ( generator.allowAssignedIdentifiers() && nullValue == null ) {
 				setNullValueUndefined();
@@ -1074,11 +1074,13 @@ public Long[] getColumnLengths() {
 	}
 
 	private class IdGeneratorCreationContext implements GeneratorCreationContext {
+		private final SimpleValue identifier;
 		private final RootClass rootClass;
 		private final Property property;
 		private final GeneratorSettings defaults;
 
-		public IdGeneratorCreationContext(RootClass rootClass, Property property, GeneratorSettings defaults) {
+		public IdGeneratorCreationContext(SimpleValue identifier, RootClass rootClass, Property property, GeneratorSettings defaults) {
+			this.identifier = identifier;
 			this.rootClass = rootClass;
 			this.property = property;
 			this.defaults = defaults;
@@ -1124,6 +1126,11 @@ public Property getProperty() {
 			return property;
 		}
 
+		@Override
+		public Value getValue() {
+			return identifier;
+		}
+
 		@Override
 		public Type getType() {
 			return SimpleValue.this.getType();

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

@@ -72,6 +72,7 @@ public class Table implements Serializable, ContributableDatabaseObject {
 	private String options;
 
 	private List<Function<SqlStringGenerationContext, InitCommand>> initCommandProducers;
+	private List<Function<SqlStringGenerationContext, InitCommand>> resyncCommandProducers;
 
 	@Deprecated(since="6.2", forRemoval = true)
 	public Table() {
@@ -808,16 +809,28 @@ public void addInitCommand(Function<SqlStringGenerationContext, InitCommand> com
 	}
 
 	public List<InitCommand> getInitCommands(SqlStringGenerationContext context) {
-		if ( initCommandProducers == null ) {
-			return emptyList();
-		}
-		else {
-			final List<InitCommand> initCommands = new ArrayList<>();
-			for ( Function<SqlStringGenerationContext, InitCommand> producer : initCommandProducers ) {
-				initCommands.add( producer.apply( context ) );
-			}
-			return unmodifiableList( initCommands );
+		return initCommandProducers == null
+				? emptyList()
+				: initCommandProducers.stream()
+						.map( producer -> producer.apply( context ) )
+						.toList();
+	}
+
+	public void addResyncCommand(Function<SqlStringGenerationContext, InitCommand> commandProducer) {
+		if ( resyncCommandProducers == null ) {
+			resyncCommandProducers = new ArrayList<>();
 		}
+		resyncCommandProducers.add( commandProducer );
+	}
+
+	public List<InitCommand> getResyncCommands(SqlStringGenerationContext context) {
+		return resyncCommandProducers == null
+				? emptyList()
+				: resyncCommandProducers.stream()
+						.map( producer -> producer.apply( context ) )
+						.distinct()
+						.toList();
+
 	}
 
 	public String getOptions() {

hibernate-core/src/main/java/org/hibernate/relational/SchemaManager.java

@@ -19,7 +19,8 @@
  * {@link jakarta.persistence.SchemaManager}, which it now inherits,
  * with a minor change to the naming of its operations. It is retained
  * for backward compatibility and as a place to define additional
- * functionality such as {@link #populate} and {@link #forSchema}.
+ * operations like {@link #populate}, {@link #resynchronizeSequences},
+ * and {@link #forSchema}.
  *
  * @since 6.2
  * @author Gavin King
@@ -95,6 +96,22 @@ public interface SchemaManager extends jakarta.persistence.SchemaManager {
 	@Incubating
 	void populate();
 
+	/**
+	 * Resynchronize sequences after importing test data.
+	 * <p>
+	 * When data is imported to the database without the use of a Hibernate session,
+	 * a database sequence might become stale with respect to the data in the table for
+	 * which it is used to generate unique keys. This operation restarts every sequence
+	 * so that the next generated unique key will be larger than the largest key
+	 * currently in use.
+	 * <p>
+	 * Programmatic way to run {@link org.hibernate.tool.schema.spi.SequenceSynchronizer}.
+	 *
+	 * @since 7.2
+	 */
+	@Incubating
+	void resynchronizeSequences();
+
 	/**
 	 * Obtain an instance which targets the given schema.
 	 * <p>