Implementing a Custom ID generator for Hibernate entities in Java

If you have used Hibernate before you know that you can specify that an Entity class should have it's @Id fields set via @GeneratedValue - typically you use either a sequence or identity strategy to automatically generate, say, auto-incrementing integer IDs (also take note that Hibernate also provides a @UuidGenerator).

However, Hibernate also allows you to create custom implementations of a generator which you can use to create a custom ID generation scheme - which can be useful for generating Transaction IDs for example.

In this article, I will show you how to implement such a custom generator that will be used to generate user-friendly, somewhat predictable, IDs. The IDs will have the following form, FAN-00YY<Month-Letter><Day-Letter>#<counter> which would generate a value that looks something like thisFAN-0023JL#3642. Of course, you can amend the solution presented here to match your own needs for the format of the IDs.

Let's get coding. To enable us to test the solution easily, we will use hibernate withing a Spring Boot project - so create a new Spring Boot project with Web, JPA and Postgres as dependencies. Done? Okay, great.

Let's assume we have an SQL table to keep track of Fans (whether the electric kind or the eccentric kind is up to you). As you can see, we have defined its id column as text and as a primary key which means we must provide unique values for that field every time we want to insert/create a record.

CREATE TABLE fans (
  id text not null primary key,
  name text not null,
  created_at timestamptz
);

The associated entity class may look like this:

import jakarta.persistence.Column;
import jakarta.persistence.Entity;
import jakarta.persistence.Id;
import jakarta.persistence.Table;

import java.time.OffsetDateTime;

@Entity
@Table(name = "fans")
public class Fan {
    @Id
    @FanId
    private String id;

    @Column(name = "name")
    private String name;

    @Column(name="created_at")
    private OffsetDateTime createdAt;

    // getters and setters omitted...
}

The attentive reader will observe that we have added a non-standard annotation @FanId to the id field. This is our custom annotation that will be used to drive the custom ID generation. Lets create the custom annotation as shown below:-

import org.hibernate.annotations.IdGeneratorType;

import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

import static java.lang.annotation.ElementType.FIELD;
import static java.lang.annotation.ElementType.METHOD;

@IdGeneratorType( FanIdGenerator.class )
@Retention(RetentionPolicy.RUNTIME)
@Target({ FIELD, METHOD })
public @interface FanId {
    boolean usePersistentCounter() default false;
}

The ID Generator class

This is the main class that handles the heavy lifting of the ID generation process. It is implemented as a class that implements the BeforeExecutionGenerator which means we can perform some actions before lifecycle events like Entity#save.

The generator defines an interface CounterService to which it delegates the generation of an automatically incrementing counter which forms the last part of our custom ID. In the constructor, we check if the annotation we are processing is set to use persistent counters, if so - we use a specific implementation of the counter service that stores the counters in the database - more on that below.

As you can see in the code below, we specify that we only want this class to be executed on INSERT operations/events.

import org.hibernate.HibernateException;
import org.hibernate.engine.spi.SharedSessionContractImplementor;
import org.hibernate.generator.BeforeExecutionGenerator;
import org.hibernate.generator.EventType;
import org.hibernate.generator.EventTypeSets;
import org.hibernate.id.factory.spi.CustomIdGeneratorCreationContext;

import java.lang.reflect.Member;
import java.time.OffsetDateTime;
import java.util.EnumSet;

import static org.hibernate.internal.util.ReflectHelper.getPropertyType;

public class FanIdGenerator implements BeforeExecutionGenerator {

    private static final CounterService DEFAULT_IN_PROCESS_COUNTER = new InMemoryCounterService();
    private static final char[] CHARS = new char[]{
            'A', 'B', 'C', 'D', 'E', 'F', 'G',
            'H', 'I', 'J', 'K', 'L', 'M', 'N',
            'O', 'P', 'Q', 'R', 'S', 'T', 'U',
            'V', 'W', 'X', 'Y', 'Z', '1', '2',
            '3', '4', '5', '6', '7', '8', '9'
    };

    private final CounterService counterService;

    public FanIdGenerator(
            FanId config,
            Member idMember,
            CustomIdGeneratorCreationContext creationContext) {

        final Class<?> propertyType = getPropertyType(idMember);

        if (config.usePersistentCounter()) {
            counterService = new PersistentCounterService();
        } else {
            counterService = DEFAULT_IN_PROCESS_COUNTER;
        }

        if (!String.class.isAssignableFrom(propertyType)) {
            throw new HibernateException("Unanticipated return type [" + propertyType.getName() + "] for FanId conversion");
        }
    }

    private static String formatYear(int year) {
        return String.format("00%d", year % 100);
    }

    private static char formatMonth(int dayValue) {
        return CHARS[dayValue - 1];
    }

    private static char formatDay(int monthValue) {
        return CHARS[monthValue - 1];
    }

    @Override
    public Object generate(SharedSessionContractImplementor session, Object owner, Object currentValue, EventType eventType) {

        OffsetDateTime now = OffsetDateTime.now();
        return String.format("FAN-%s%s%s#%s",
                formatYear(now.getYear()),
                formatMonth(now.getMonthValue()),
                formatDay(now.getDayOfMonth()),
                counterService.nextInteger(session, now)
        );
    }

    @Override
    public EnumSet<EventType> getEventTypes() {
        return EventTypeSets.INSERT_ONLY;
    }

    interface CounterService {
        int nextInteger(SharedSessionContractImplementor session, OffsetDateTime onDate);
    }
}

The ID generation scheme was inspired by some code shared and implemented by a friend for a health information system used widely in Malawi, the code for that is here.

Integer Counter generator using in-memory counter

With this simplest implementation, the counter service uses a variable to keep track of the values generated. We use an AtomicInteger for this to enable some kind of thread safety, in case that's required. The implementation is fairly straightforward:-

import org.hibernate.engine.spi.SharedSessionContractImplementor;

import java.time.OffsetDateTime;
import java.util.concurrent.atomic.AtomicInteger;

class InMemoryCounterService implements FanIdGenerator.CounterService {
    private static final AtomicInteger id = new AtomicInteger(1000);

    @Override
    public int nextInteger(SharedSessionContractImplementor session, OffsetDateTime onDate) {
        return id.incrementAndGet();
    }
}

The above approach works but has the limitation that the counter of the IDs will be reset every time the service is restarted which may lead to errors since that introduces the potential of generating duplicate keys (remember we want the IDs to be unique to use them as primary keys). We can do better by storing the counter's latest value in the database, indexed by date - this way we can be sure that we can always continue from where we left off despite restarts.

Integer Counter Generator using a value stored in a database table

Below is the implementation of the CounterService interface, named PersistentCounterService, which stores counter values in a database table to allow us to fetch easily the latest counter value to try to avoid collisions. Each date has its counter that gets incremented every time we receive a request to generate an ID for that day.

As you can see in the code, the table that keeps track of the counters is named fan_id_counters and we try to create it every time, this allows our ID Generator to be re-usable without any upfront configuration - the only implication is that the user connecting to the database should have appropriate permissions to create tables.

You could always change the implementation to remove the create call and manually create your counter-tracking table ahead of time.

import org.hibernate.engine.spi.SharedSessionContractImplementor;

import java.time.LocalDate;
import java.time.OffsetDateTime;

public class PersistentCounterService implements FanIdGenerator.CounterService {
    private static final int DEFAULT_COUNTER_START = 1000;
    private static final String SQL_CREATE_COUNTER_TABLE = """
            CREATE TABLE IF NOT EXISTS fan_id_counters (
                entry_date date primary key not null,
                counter_value integer not null default 1000
            );
            """;
    @Override
    public int nextInteger(SharedSessionContractImplementor session, OffsetDateTime onDate) {
        final var today = LocalDate.now();
        var statelessSession = session.isStatelessSession() ? session.asStatelessSession() : session.getSession();
        statelessSession.createNativeMutationQuery(SQL_CREATE_COUNTER_TABLE)
                .executeUpdate();

        var list = statelessSession.createNativeQuery("SELECT counter_value FROM fan_id_counters WHERE entry_date = :date FOR UPDATE LIMIT 1;", Integer.class)
                .setParameter("date", today)
                .setFetchSize(1)
                .getResultList();
        if (list.isEmpty()) {
            statelessSession.createNativeMutationQuery("INSERT INTO fan_id_counters(entry_date, counter_value) VALUES (:date , :counter)")
                    .setParameter("date", today)
                    .setParameter("counter", DEFAULT_COUNTER_START + 1)
                    .executeUpdate();
            return DEFAULT_COUNTER_START;
        }

        statelessSession.createNativeMutationQuery("UPDATE fan_id_counters SET counter_value = counter_value + 1 WHERE entry_date = :date ;")
                .setParameter("date", today)
                .executeUpdate();

        return list.get(0);
    }
}

Let's test this thing

Now that we have the ID generator setup, we can implement a simple controller to allow us to create Fan entities to ensure that the generator does its job and also do a little bit of load testing to ensure that the generator doesn't trip up when there is a large number of requests, as one use-case would be to use it to generate Transaction IDs or Transaction Reference numbers for example.

First, let's create a simple controller that allows us to list and create Fans

@RestController
@RequestMapping("/api/v1/fans")
public class FanController {

    private final FanRepository fanRepository;

    public FanController(FanRepository fanRepository) {
        this.fanRepository = fanRepository;
    }

    @GetMapping
    public ResponseEntity<List<Fan>> createFan(Pageable pageable) {
        return ResponseEntity.ok(fanRepository.findAll(pageable).toList());
    }

    @PostMapping
    public ResponseEntity<Fan> createFan(@RequestParam("name") String name) {
        Fan fan = new Fan();
        fan.setName(name);
        fan.setCreatedAt(OffsetDateTime.now());
        fanRepository.saveAndFlush(fan);
        return ResponseEntity.ok(fan);
    }
}

We can use the following cURL request to create a Fan

curl -X POST "http://localhost:8080/api/v1/fans?name=James"

This could return something like this

{
    "id": "FAN-0023JL#1000",
    "name": "James",
    "createdAt": "2023-08-12T21:04:44.135036Z"
},

Load testing the ID Generator

We can now use a load testing tool like Hey to load test the server - which will send about 1000 requests with the command below.

hey -n 1000 -m POST "http://localhost:8080/api/v1/fans?name=James"

Here is a screenshot that shows a sample run on my environment, not so bad :) -

We can check in the database and verify that each fan has a uniquely generated ID.

Conclusion

In this article, we have seen how to implement a custom ID generator with Hibernate that uses a custom format that includes an auto-incrementing counter that has two implementations, one in-memory and another that stores the most recent counter in a database table indexed by calendar day. You can adapt the code here to implement your ID generator which can be useful for things like Transaction IDs or Transaction References that are somewhat more user friendly.

Thank you for reading.