introduce 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/cfg/MappingSettings.java

@@ -118,7 +118,7 @@ public interface MappingSettings {
 	String KEYWORD_AUTO_QUOTING_ENABLED = "hibernate.auto_quote_keyword";
 
 	/**
-	 * When a generator specifies an increment-size and an optimizer was not explicitly
+	 * When a generator specifies an increment size and an optimizer was not explicitly
 	 * specified, which of the "pooled" optimizers should be preferred? Can specify an
 	 * optimizer short name or the name of a class which implements
 	 * {@link org.hibernate.id.enhanced.Optimizer}.

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 optimizer The {@link Optimizer} for this generator
+	 *
+	 * @see org.hibernate.relational.SchemaManager#resynchronizeSequences()
+	 *
+	 * @since 7.2
+	 */
+	default void registerExtraExportables(Table table, Optimizer optimizer) {
+	}
+
 	/**
 	 * Initializes this structure, in particular pre-generates SQL as necessary.
 	 * <p>

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

@@ -58,7 +58,7 @@ public interface Optimizer {
 	 *
 	 * @return True if the values in the source are to be incremented
 	 * according to the defined increment size; false otherwise, in which
-	 * case the increment is totally an in memory construct.
+	 * case the increment size is a completely in-memory construct.
 	 */
 	boolean applyIncrementSizeToSourceValues();
 
@@ -75,4 +75,11 @@ public interface Optimizer {
 	 * @since 7.1
 	 */
 	Expression createLowValueExpression(Expression databaseValue, SessionFactoryImplementor sessionFactory);
+
+	/**
+	 * @since 7.2
+	 */
+	default int getAdjustment() {
+		return 1;
+	}
 }

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

@@ -59,7 +59,7 @@ public Serializable generate(AccessCallback callback) {
 				generationState.lastSourceValue = callback.getNextValue();
 				generationState.upperLimitValue = generationState.lastSourceValue.copy().add( incrementSize );
 				generationState.value = generationState.lastSourceValue.copy();
-				// handle cases where initial-value is less that one (hsqldb for instance).
+				// handle cases where the initial value is less than one (hsqldb, for instance)
 				while ( generationState.value.lt( 1 ) ) {
 					generationState.value.increment();
 				}

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

@@ -61,8 +61,7 @@ public PooledOptimizer(Class<?> returnClass, int incrementSize) {
 	public Serializable generate(AccessCallback callback) {
 		lock.lock();
 		try {
-			final GenerationState generationState = locateGenerationState( callback.getTenantIdentifier() );
-
+			final var generationState = locateGenerationState( callback.getTenantIdentifier() );
 			if ( generationState.hiValue == null ) {
 				generationState.hiValue = callback.getNextValue();
 				// unfortunately not really safe to normalize this
@@ -86,7 +85,6 @@ else if ( generationState.value.gt( generationState.hiValue ) ) {
 				generationState.hiValue = callback.getNextValue();
 				generationState.value = generationState.hiValue.copy().subtract( incrementSize - 1 );
 			}
-
 			return generationState.value.makeValueThenIncrement();
 		}
 		finally {
@@ -175,4 +173,9 @@ public Expression createLowValueExpression(Expression databaseValue, SessionFact
 				integerType
 		);
 	}
+
+	@Override
+	public int getAdjustment() {
+		return incrementSize;
+	}
 }

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

@@ -4,15 +4,20 @@
  */
 package org.hibernate.id.enhanced;
 
+import java.sql.Connection;
 import java.sql.SQLException;
 
 import org.hibernate.AssertionFailure;
+import org.hibernate.HibernateException;
 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.dialect.Dialect;
 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 +169,51 @@ public void initialize(SqlStringGenerationContext context) {
 				.getSequenceNextValString( context.format( physicalSequenceName ) );
 	}
 
+	@Override
+	public void registerExtraExportables(Table table, Optimizer optimizer) {
+		table.addResyncCommand( (context, connection) -> {
+			final String sequenceName = context.format( physicalSequenceName );
+			final String tableName = context.format( table.getQualifiedTableName() );
+			final String primaryKeyColumnName = table.getPrimaryKey().getColumn( 0 ).getName();
+			final int adjustment = optimizer.getAdjustment();
+			final long max = getMax( connection, primaryKeyColumnName, tableName );
+			final long current = getCurrent( connection, sequenceName, context.getDialect() );
+			final long newValue = Math.max( max + adjustment, current );
+			final String restart = "alter sequence " + sequenceName + " restart with " + newValue;
+			return new InitCommand( restart );
+		} );
+	}
+
+
+	private long getCurrent(Connection connection, String sequenceName, Dialect dialect) {
+		final String sequenceCurrentValue =
+				dialect.getSequenceSupport()
+						.getSequenceNextValString( sequenceName );
+		try ( var select = connection.prepareStatement( sequenceCurrentValue ) ) {
+			try ( var resultSet = select.executeQuery() ) {
+				resultSet.next();
+				return resultSet.getLong(1);
+			}
+		}
+		catch (SQLException e) {
+			throw new HibernateException( "Could not fetch the current sequence value from the database", e );
+		}
+	}
+
+	private static long getMax(Connection connection, String primaryKeyColumnName, String tableName) {
+		final String selectMax =
+				"select max(" + primaryKeyColumnName + ") from " + tableName;
+		try ( var select = connection.prepareStatement( selectMax ) ) {
+			try ( var resultSet = select.executeQuery() ) {
+				resultSet.next();
+				return resultSet.getLong(1);
+			}
+		}
+		catch (SQLException e) {
+			throw new HibernateException( "Could not fetch the max primary key from the database", e );
+		}
+	}
+
 	@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, optimizer );
 	}
 
 	@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

@@ -10,6 +10,7 @@
 import java.sql.SQLException;
 
 import org.hibernate.AssertionFailure;
+import org.hibernate.HibernateException;
 import org.hibernate.LockOptions;
 import org.hibernate.boot.model.naming.Identifier;
 import org.hibernate.boot.model.relational.Database;
@@ -333,4 +334,53 @@ public void initialize(SqlStringGenerationContext context) {
 				" set " + valueColumnNameText + "= ?" +
 				" where " + valueColumnNameText + "=?";
 	}
+
+	@Override
+	public void registerExtraExportables(Table table, Optimizer optimizer) {
+		table.addResyncCommand( (context, connection) -> {
+			final String sequenceTableName = context.format( physicalTableName );
+			final String tableName = context.format( table.getQualifiedTableName() );
+			final String primaryKeyColumnName = table.getPrimaryKey().getColumn( 0 ).getName();
+			final int adjustment = optimizer.getAdjustment();
+			final long max = getMax( connection, primaryKeyColumnName, tableName );
+			final long current = getCurrent( connection, sequenceTableName );
+			if ( max + adjustment > current ) {
+				final String update =
+						"update " + sequenceTableName
+						+ " set " + valueColumnNameText + " = " + (max + adjustment);
+				return new InitCommand( update );
+			}
+			else {
+				return new InitCommand();
+			}
+		} );
+	}
+
+	private long getCurrent(Connection connection, String sequenceTableName) {
+		final String selectCurrent =
+				"select " + valueColumnNameText + " from " + sequenceTableName;
+		try ( var select = connection.prepareStatement( selectCurrent ) ) {
+			try ( var resultSet = select.executeQuery() ) {
+				resultSet.next();
+				return resultSet.getLong(1);
+			}
+		}
+		catch (SQLException e) {
+			throw new HibernateException( "Could not fetch the current sequence value from the database", e );
+		}
+	}
+
+	private static long getMax(Connection connection, String primaryKeyColumnName, String tableName) {
+		final String selectMax =
+				"select max(" + primaryKeyColumnName + ") from " + tableName;
+		try ( var select = connection.prepareStatement( selectMax ) ) {
+			try ( var resultSet = select.executeQuery() ) {
+				resultSet.next();
+				return resultSet.getLong(1);
+			}
+		}
+		catch (SQLException e) {
+			throw new HibernateException( "Could not fetch the max primary key from the database", e );
+		}
+	}
 }