BatchSequenceGenerator.java

package com.vladmihalcea.hibernate.id;

import org.hibernate.HibernateException;
import org.hibernate.MappingException;
import org.hibernate.cfg.ObjectNameNormalizer;
import org.hibernate.dialect.Dialect;
import org.hibernate.engine.jdbc.spi.JdbcCoordinator;
import org.hibernate.engine.spi.SessionImplementor;
import org.hibernate.id.BulkInsertionCapableIdentifierGenerator;
import org.hibernate.id.Configurable;
import org.hibernate.id.IdentifierGenerationException;
import org.hibernate.id.PersistentIdentifierGenerator;
import org.hibernate.id.enhanced.DatabaseStructure;
import org.hibernate.id.enhanced.SequenceStructure;
import org.hibernate.internal.util.config.ConfigurationHelper;
import org.hibernate.mapping.Table;
import org.hibernate.type.Type;

import java.io.Serializable;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * A sequence generator that uses a recursive query to fetch multiple
 * values from a sequence in a single database access.
 *
 * <h2>Parameters</h2>
 * The following configuration parameters are supported:
 * <dl>
 * <dt>{@value #SEQUENCE_PARAM}</dt>
 * <dd><strong>mandatory</strong>, name of the sequence to use</dd>
 * <dt>{@value #FETCH_SIZE_PARAM}</dt>
 * <dd>optional, how many sequence numbers should be fetched at a time,
 * default is {@value #DEFAULT_FETCH_SIZE}</dd>
 * </dl>
 *
 * <h2>SQL</h2>
 * Per default the generated SELECT will look something like this
 * <pre><code>
 * WITH RECURSIVE t(n) AS (
 *   SELECT 1
 *     UNION ALL
 *   SELECT n + 1
 *   FROM t
 *   WHERE n &lt; ?)
 * SELECT nextval(seq_xxx)
 *   FROM t;
 * </code></pre>
 *
 * <h3>DB2</h3>
 * For DB2 the generated SELECT will look something like this
 * <pre><code>
 * WITH t(n) AS (
 *   SELECT 1 as n
 *     FROM (VALUES 1)
 *       UNION ALL
 *     SELECT n + 1 as n
 *       FROM t
 *      WHERE n &lt; ?)
 * SELECT next value for SEQ_CHILD_ID as n
 *   FROM t;
 * </code></pre>
 *
 * <h3>HSQLDB</h3>
 * For HSQLDB the generated SELECT will look something like this
 * <pre><code>
 * SELECT next value for seq_parent_id
 *   FROM UNNEST(SEQUENCE_ARRAY(1, ?, 1));
 * </code></pre>
 *
 * <h3>Oracle</h3>
 * For Oracle the generated SELECT will look something like this
 * because Oracle does not support using recursive common table
 * expressions to fetch multiple values from a sequence.
 * <pre><code>
 * SELECT seq_xxx.nextval
 * FROM dual
 * CONNECT BY rownum &lt;= ?
 * </code></pre>
 *
 * <h3>SQL Server</h3>
 * For SQL Server the generated SELECT will look something like this
 * <pre><code>
 * WITH t(n) AS (
 *   SELECT 1 as n
 *     UNION ALL
 *   SELECT n + 1 as n
 *     FROM t
 *    WHERE n &lt; ?)
 * SELECT NEXT VALUE FOR seq_xxx as n
 *   FROM t
 * </code></pre>
 *
 * <h3>Firebird</h3>
 * For Firebird the generated SELECT will look something like this
 * <pre><code>
 * WITH RECURSIVE t(n, level_num) AS (
 *   SELECT NEXT VALUE FOR seq_xxx as n, 1 as level_num
 *   FROM rdb$database
 *     UNION ALL
 *   SELECT NEXT VALUE FOR seq_xxx as n, level_num + 1 as level_num
 *     FROM t
 *    WHERE level_num &lt; ?)
 * SELECT n
 *   FROM t
 * </code></pre>
 *
 * <h2>Database Support</h2>
 * The following RDBMS have been verified to work
 * <ul>
 *  <li>DB2</li>
 *  <li>Firebird</li>
 *  <li>Oracle</li>
 *  <li>H2</li>
 *  <li>HSQLDB</li>
 *  <li>MariaDB</li>
 *  <li>Postgres</li>
 *  <li>SQL Sever</li>
 * </ul>
 * <p>
 * In theory any RDBMS that supports {@code WITH RECURSIVE} and
 * sequences is supported.
 * <p>
 * For more details about how to use it, check out <a href="https://vladmihalcea.com/hibernate-batch-sequence-generator/">this article</a> on <a href="https://vladmihalcea.com/">vladmihalcea.com</a>.
 *
 * @author Philippe Marschall
 * @since 2.14.0
 */
public class BatchSequenceGenerator implements BulkInsertionCapableIdentifierGenerator, PersistentIdentifierGenerator, Configurable {

    /**
     * Indicates the name of the sequence to use, mandatory.
     */
    public static final String SEQUENCE_PARAM = "sequence";

    /**
     * Indicates how many sequence values to fetch at once. The default value is {@link #DEFAULT_FETCH_SIZE}.
     */
    public static final String FETCH_SIZE_PARAM = "fetch_size";

    /**
     * The default value for {@link #FETCH_SIZE_PARAM}.
     */
    public static final int DEFAULT_FETCH_SIZE = 10;

    private final Lock lock = new ReentrantLock();

    private String select;
    private int fetchSize;
    private IdentifierPool identifierPool;
    private IdentifierExtractor identifierExtractor;
    private DatabaseStructure databaseStructure;

    @Override
    public void configure(Type type, Properties params, Dialect dialect)
        throws MappingException {

        String sequenceName = determineSequenceName(params, dialect);
        this.fetchSize = determineFetchSize(params);

        this.select = buildSelect(sequenceName, dialect);
        this.identifierExtractor = IdentifierExtractor.getIdentifierExtractor(type.getReturnedClass());
        this.identifierPool = IdentifierPool.empty();

        this.databaseStructure = this.buildDatabaseStructure(type, sequenceName, dialect);
    }

    private static String buildSelect(String sequenceName, Dialect dialect) {
        if (dialect instanceof org.hibernate.dialect.Oracle8iDialect) {
            return "SELECT " + dialect.getSelectSequenceNextValString(sequenceName) + " FROM dual CONNECT BY rownum <= ?";
        }
        if (dialect instanceof org.hibernate.dialect.SQLServerDialect) {
            // No RECURSIVE
            return "WITH t(n) AS ( "
                + "SELECT 1 as n "
                + "UNION ALL "
                +"SELECT n + 1 as n FROM t WHERE n < ?) "
                // sequence generation outside of WITH
                + "SELECT " + dialect.getSelectSequenceNextValString(sequenceName) + " as n FROM t";
        }
        if (dialect instanceof org.hibernate.dialect.DB2Dialect) {
            // No RECURSIVE
            return "WITH t(n) AS ( "
                + "SELECT 1 as n "
                // difference
                + "FROM (VALUES 1) "
                + "UNION ALL "
                +"SELECT n + 1 as n FROM t WHERE n < ?) "
                // sequence generation outside of WITH
                + "SELECT " + dialect.getSelectSequenceNextValString(sequenceName) + " as n FROM t";
        }
        if (dialect instanceof org.hibernate.dialect.HSQLDialect) {
            // https://stackoverflow.com/questions/44472280/cte-based-sequence-generation-with-hsqldb/52329862
            return "SELECT " + dialect.getSelectSequenceNextValString(sequenceName) + " FROM UNNEST(SEQUENCE_ARRAY(1, ?, 1))";
        }
        if (dialect instanceof org.hibernate.dialect.FirebirdDialect) {
            return "WITH RECURSIVE t(n, level_num) AS ("
                + "SELECT " + dialect.getSelectSequenceNextValString(sequenceName) + " as n, 1 as level_num "
                // difference
                + " FROM rdb$database "
                + "UNION ALL "
                + "SELECT " + dialect.getSelectSequenceNextValString(sequenceName) + " as n, level_num + 1 as level_num "
                + " FROM t "
                + " WHERE level_num < ?) "
                + "SELECT n FROM t";
        }
        return "WITH RECURSIVE t(n) AS ("
            + "SELECT 1 "
            + "UNION ALL "
            + "SELECT n + 1"
            + " FROM t "
            + " WHERE n < ?) "
            + "SELECT " + dialect.getSelectSequenceNextValString(sequenceName) + " FROM t";
    }

    private SequenceStructure buildDatabaseStructure(Type type, String sequenceName, Dialect dialect) {
        return new SequenceStructure(dialect, sequenceName, 1, 1, type.getReturnedClass());
    }

    private static String determineSequenceName(Properties params, Dialect dialect) {
        ObjectNameNormalizer normalizer = (ObjectNameNormalizer) params.get(IDENTIFIER_NORMALIZER);
        String sequenceName = normalizer.normalizeIdentifierQuoting(
            ConfigurationHelper.getString(SEQUENCE_PARAM, params));

        if (sequenceName == null) {
            throw new MappingException("no squence name specified");
        }
        if (sequenceName.indexOf('.') == -1) {
            sequenceName = normalizer.normalizeIdentifierQuoting(sequenceName);
            String schemaName = params.getProperty(SCHEMA);
            String catalogName = params.getProperty(CATALOG);
            sequenceName = Table.qualify(dialect.quote(catalogName),
                dialect.quote(schemaName),
                dialect.quote(sequenceName));
        }
        return sequenceName;
    }

    private static int determineFetchSize(Properties params) {
        int fetchSize = ConfigurationHelper.getInt(FETCH_SIZE_PARAM, params, DEFAULT_FETCH_SIZE);
        if (fetchSize <= 0) {
            throw new MappingException("fetch size must be positive");
        }
        return fetchSize;
    }

    @Override
    public boolean supportsBulkInsertionIdentifierGeneration() {
        return true;
    }

    @Override
    public String determineBulkInsertionIdentifierGenerationSelectFragment(Dialect dialect) {
        return dialect.getSelectSequenceNextValString(this.getSequenceName());
    }

    @Override
    public Serializable generate(SessionImplementor session, Object object) throws HibernateException {
        this.lock.lock();
        try {
            if (this.identifierPool.isEmpty()) {
                this.identifierPool = this.replenishIdentifierPool(session);
            }
            return this.identifierPool.next();
        } finally {
            this.lock.unlock();
        }
    }

    @Override
    public Object generatorKey() {
        return this.getSequenceName();
    }

    private String getSequenceName() {
        return this.databaseStructure.getName();
    }

    @Override
    @Deprecated
    public String[] sqlCreateStrings(Dialect dialect) {
        return this.databaseStructure.sqlCreateStrings(dialect);
    }

    @Override
    @Deprecated
    public String[] sqlDropStrings(Dialect dialect) {
        return this.databaseStructure.sqlDropStrings(dialect);
    }

    private IdentifierPool replenishIdentifierPool(SessionImplementor session) throws HibernateException {
        List<Serializable> identifiers = new ArrayList<Serializable>(this.fetchSize);
        JdbcCoordinator coordinator = session.getTransactionCoordinator().getJdbcCoordinator();
        PreparedStatement statement = coordinator.getStatementPreparer().prepareStatement(this.select);
        try {
            statement.setFetchSize(this.fetchSize);
            statement.setInt(1, this.fetchSize);

            ResultSet resultSet = coordinator.getResultSetReturn().extract(statement);
            try {
                while (resultSet.next()) {
                    identifiers.add(this.identifierExtractor.extractIdentifier(resultSet));
                }
            } finally {
                coordinator.release(resultSet, statement);
            }
        } catch (SQLException e) {
            throw session.getFactory().getSQLExceptionHelper()
                .convert(e, "could not get next sequence value", this.select);
        } finally {
            coordinator.release(statement);
        }
        if (identifiers.size() != this.fetchSize) {
            throw new IdentifierGenerationException(
                "expected " + this.fetchSize + " values from " + this.getSequenceName() + " but got " + identifiers.size());
        }
        return IdentifierPool.forList(identifiers);
    }

    /**
     * Holds a number of prefetched identifiers.
     */
    static final class IdentifierPool {

        private final Iterator<Serializable> iterator;

        private IdentifierPool (List<Serializable> identifiers) {
            this.iterator = identifiers.iterator();
        }

        static IdentifierPool forList(List<Serializable> identifiers) {
            return new IdentifierPool(identifiers);
        }

        static IdentifierPool empty() {
            return new IdentifierPool(Collections.<Serializable>emptyList());
        }

        boolean isEmpty() {
            return !this.iterator.hasNext();
        }

        Serializable next() {
            return this.iterator.next();
        }

    }

    /**
     * Extracts a {@link Serializable} identifier from a {@link ResultSet}.
     *
     * @see org.hibernate.id.IntegralDataTypeHolder
     */
    enum IdentifierExtractor {

        INTEGER_IDENTIFIER_EXTRACTOR {
            @Override
            Serializable extractIdentifier(ResultSet resultSet) throws SQLException {
                int intValue = resultSet.getInt(1);
                if (resultSet.wasNull()) {
                    throw new IdentifierGenerationException("sequence returned null");
                }
                return intValue;
            }
        },

        LONG_IDENTIFIER_EXTRACTOR {
            @Override
            Serializable extractIdentifier(ResultSet resultSet) throws SQLException {
                long longValue = resultSet.getLong(1);
                if (resultSet.wasNull()) {
                    throw new IdentifierGenerationException("sequence returned null");
                }
                return longValue;
            }
        },

        BIG_INTEGER_IDENTIFIER_EXTRACTOR {
            @Override
            Serializable extractIdentifier(ResultSet resultSet) throws SQLException {
                BigDecimal bigDecimal = resultSet.getBigDecimal(1);
                if (resultSet.wasNull()) {
                    throw new IdentifierGenerationException("sequence returned null");
                }
                return bigDecimal.setScale(0, BigDecimal.ROUND_UNNECESSARY).toBigInteger();
            }
        },

        BIG_DECIMAL_IDENTIFIER_EXTRACTOR {
            @Override
            Serializable extractIdentifier(ResultSet resultSet) throws SQLException {
                BigDecimal bigDecimal = resultSet.getBigDecimal(1);
                if (resultSet.wasNull()) {
                    throw new IdentifierGenerationException("sequence returned null");
                }
                return bigDecimal;
            }
        };

        abstract Serializable extractIdentifier(ResultSet resultSet) throws SQLException;

        static IdentifierExtractor getIdentifierExtractor(Class<?> integralType) {
            if ((integralType == Integer.class) || (integralType == int.class)) {
                return INTEGER_IDENTIFIER_EXTRACTOR;
            }
            if ((integralType == Long.class) || (integralType == long.class)) {
                return LONG_IDENTIFIER_EXTRACTOR;
            }
            if (integralType == BigInteger.class) {
                return BIG_INTEGER_IDENTIFIER_EXTRACTOR;
            }
            if (integralType == BigDecimal.class) {
                return BIG_DECIMAL_IDENTIFIER_EXTRACTOR;
            }
            throw new IdentifierGenerationException("unsupported integral type: " + integralType);
        }

    }

    @Override
    public String toString() {
        // for debugging only
        return this.getClass().getSimpleName() + '(' + this.getSequenceName() + ')';
    }

}