Obtaining Auto-generated Keys in Spring JDBC

1. Introduction

In this quick tutorial, we’ll explore the possibility of getting the auto-generated key after inserting entities when working with Spring JDBC.

2. Maven Dependencies

At first, we need to have spring-boot-starter-jdbc and H2 dependencies defined in our pom.xml:

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-jdbc</artifactId>
</dependency>
<dependency>
    <groupId>com.h2database</groupId>
    <artifactId>h2</artifactId>
    <scope>runtime</scope>
</dependency>

We can check out the latest version of those two dependencies on Maven Central: spring-boot-starter-jdbc and h2.

3. Getting the Auto-Generated Key

3.1. The Scenario

Let’s define a sys_message table which has 2 columns: id (auto-generated key) and message:

CREATE TABLE IF NOT EXISTS sys_message (
    id bigint(20) NOT NULL AUTO_INCREMENT,
    message varchar(100) NOT NULL,
    PRIMARY KEY (id)
);

3.2. Using the JdbcTemplate

Now, let’s implement a method which will use JDBCTemplate to insert the new record and return the auto-generated id. 

Therefore, we’ll use the JDBCTemplate update() method which supports the retrieval of primary keys generated by the database. This method takes an instance of the PrepareStatementCreator interface as the first argument and the other argument is the KeyHolder. 

Since the PrepareStatementCreator interface is a FunctionalInterface where its method accepts an instance of java.sql.Connection and return a java.sql.PreparedStatement object, for simplicity, we can use a lambda expression:

String INSERT_MESSAGE_SQL 
  = "insert into sys_message (message) values(?) ";
    
public long insertMessage(String message) {    
    KeyHolder keyHolder = new GeneratedKeyHolder();

    jdbcTemplate.update(connection -> {
        PreparedStatement ps = connection
          .prepareStatement(INSERT_MESSAGE_SQL);
          ps.setString(1, message);
          return ps;
        }, keyHolder);

        return (long) keyHolder.getKey();
    }
}

It’s worth noting that the keyHolder object will contain the auto-generated key return from the JDBCTemplate update() method.

We can retrieve that key by calling keyHolder.getKey().

Besides, we can verify the method:

@Test
public void 
  insertJDBC_whenLoadMessageByKey_thenGetTheSameMessage() {
    long key = messageRepositoryJDBCTemplate.insert(MESSAGE_CONTENT);
    String loadedMessage = messageRepositoryJDBCTemplate
      .getMessageById(key);

    assertEquals(MESSAGE_CONTENT, loadedMessage);
}

3.3. Using SimpleJdbcInsert

In addition to the JDBCTemplate, we also can use SimpleJdbcInsert to achieve the same result.

Hence, we need to initialize an instance of the SimpleJdbcInsert:

@Repository
public class MessageRepositorySimpleJDBCInsert {

    SimpleJdbcInsert simpleJdbcInsert;

    @Autowired
    public MessageRepositorySimpleJDBCInsert(DataSource dataSource) {
        simpleJdbcInsert = new SimpleJdbcInsert(dataSource)
          .withTableName("sys_message").usingGeneratedKeyColumns("id");
    }
    
    //...
}

Consequently, we can call the executeAndReturnKey method of the SimpleJdbcInsert to insert a new record to sys_message table and get back the auto-generated key:

public long insert(String message) {
    Map<String, Object> parameters = new HashMap<>(1);
    parameters.put("message", message);
    Number newId = simpleJdbcInsert.executeAndReturnKey(parameters);
    return (long) newId;
}

Furthermore, we can verify that method quite simply:

@Test
public void 
  insertSimpleInsert_whenLoadMessageKey_thenGetTheSameMessage() {
    long key = messageRepositorySimpleJDBCInsert.insert(MESSAGE_CONTENT);
    String loadedMessage = messageRepositoryJDBCTemplate.getMessageById(key);

    assertEquals(MESSAGE_CONTENT, loadedMessage);
}

4. Conclusion

We’ve explored the possibility of using JDBCTemplate and SimpleJdbcInsert for inserting a new record and getting the auto-generated key back.

As always, we can find the implementation of this article over on Github.