/atlassian-hibernate5-extras/src/main/java/com/atlassian/hibernate/extras/ResetableHiLoGeneratorHelper.java

package com.atlassian.hibernate.extras;

import org.hibernate.JDBCException;
import org.hibernate.Session;
import org.hibernate.SessionFactory;
import org.hibernate.engine.spi.SessionFactoryImplementor;
import org.hibernate.engine.spi.SessionImplementor;
import org.hibernate.jdbc.AbstractReturningWork;
import org.hibernate.jdbc.AbstractWork;
import org.hibernate.metadata.ClassMetadata;
import org.hibernate.persister.entity.EntityPersister;
import org.hibernate.persister.entity.SingleTableEntityPersister;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.List;
import java.util.Map;

/**
 * This class is used to set the hibernate_unique_key.next_hi value to the high bits of the highest id in the DB, +1
 * after an import (for now) to maintain state in the database.
 *
 * @since 6.1
 */
@SuppressWarnings("unused")
public class ResetableHiLoGeneratorHelper {

    public static final String HIBERNATE_UNIQUE_KEY_COLUMN = "next_hi";
    public static final String HIBERNATE_UNIQUE_KEY_TABLE = "hibernate_unique_key";

    private static final Logger log = LoggerFactory.getLogger(ResetableHiLoGeneratorHelper.class);

    private SessionFactory sessionFactory;

    /**
     * @return long the next hi value
     */
    public long getHiValue() throws SQLException {
        Session session = sessionFactory.getCurrentSession();
        try {
            return ((SessionImplementor) session).getTransactionCoordinator()
                    .createIsolationDelegate()
                    .delegateWork(new NextHiGetter(), true);
        } catch (JDBCException e) {
            log.error("Error finding maximum next_hi value", e.getSQLException());

            throw e.getSQLException();
        }
    }

    public void setNextHiValue(List<String> errors) {
        Session session = sessionFactory.getCurrentSession();
        ((SessionImplementor) session).getTransactionCoordinator()
                .createIsolationDelegate()
                .delegateWork(new NextHiSetter(errors, sessionFactory), true);
    }

    public void setSessionFactory(SessionFactory sessionFactory) {
        this.sessionFactory = sessionFactory;
    }

    private static class NextHiGetter extends AbstractReturningWork<Long> {
        private static final String SQL = "select " + HIBERNATE_UNIQUE_KEY_COLUMN + " from " + HIBERNATE_UNIQUE_KEY_TABLE;

        @Override
        public Long execute(Connection connection) throws SQLException {
            try (Statement statement = connection.createStatement();
                 ResultSet resultSet = statement.executeQuery(SQL)) {
                if (resultSet.next()) {
                    return resultSet.getLong(1); // might be null, but that's ok, will return 0;
                }
            }

            return 0L;
        }
    }

    private static class NextHiSetter extends AbstractWork {
        private final List<String> errors;
        private final SessionFactoryImplementor sessionFactory;

        private NextHiSetter(List<String> errors, SessionFactory sessionFactory) {
            this.errors = errors;
            this.sessionFactory = (SessionFactoryImplementor) sessionFactory;
        }

        @Override
        public void execute(Connection connection) {
            Map<String, ClassMetadata> metadataMap = sessionFactory.getAllClassMetadata();

            try (Statement statement = connection.createStatement()) {
                long maximumId = 0;
                int maxLo = 0;
                for (String className : metadataMap.keySet()) {
                    EntityPersister persister = sessionFactory.getEntityPersister(className);
                    if (persister instanceof SingleTableEntityPersister &&
                            persister.getIdentifierGenerator() instanceof ResettableTableHiLoGenerator) {
                        SingleTableEntityPersister entityPersister = (SingleTableEntityPersister) persister;
                        ResettableTableHiLoGenerator generator = (ResettableTableHiLoGenerator) persister.getIdentifierGenerator();

                        // TODO need to work out why we have this logic.
                        if (maxLo == 0) {
                            maxLo = generator.getMaxLo();
                        } else if (maxLo != generator.getMaxLo()) {
                            //throw new UpgradeException("One generator uses " + maxLo + " for maxLo, generator for " +
                            //        c.getName() + " uses " + generator.getMaxLo());
                            errors.add("One generator uses " + maxLo + " for maxLo, generator for " +
                                    className + " uses " + generator.getMaxLo());
                        }

                        String[] idColumnNames = entityPersister.getIdentifierColumnNames();
                        if (idColumnNames.length != 1) {
                            errors.add("Expected a single ID column for " + className + " found " + idColumnNames.length);
                        }

                        String sql = "select max(" + idColumnNames[0] + ") from " + entityPersister.getTableName();
                        try (ResultSet resultSet = statement.executeQuery(sql)) {
                            if (resultSet.next()) {
                                long value = resultSet.getLong(1); // might be null, but that's ok, will return 0;
                                if (value > maximumId) {
                                    maximumId = value;
                                }
                            } else {
                                errors.add("No maximum ID returned for " + className);
                            }
                        }
                    }
                }

                setNextHi(statement, (int) (maximumId / (maxLo + 1)) + 1);
            } catch (Exception e) {
                log.error("Error finding maximum next_hi value", e);
                errors.add(e.getMessage());
            } finally {
                log.info("Completed database update: HiLoIdRepairUpgradeTask");
            }
        }

        private void setNextHi(Statement statement, int nextHi) throws SQLException {
            log.info("Setting new next_hi to " + nextHi);
            if (statement.executeUpdate("update " + HIBERNATE_UNIQUE_KEY_TABLE + " set " + HIBERNATE_UNIQUE_KEY_COLUMN + " = " + nextHi) == 0) {
                // no row, need to insert one
                if (statement.executeUpdate("insert into " + HIBERNATE_UNIQUE_KEY_TABLE + " values(" + nextHi + ")") == 0) {
                    errors.add("failed to insert initial next_hi value");
                }
            }
        }
    }
}