Retrieve generated IDs when doing bulk insert using Spring Data R2DBC

Question

I have a scenario where my table has an autogenerated id column and I need to bulk insert items into db and fetch the generated ids. Is there any way I can achieve that?

This is my table:

CREATE TABLE test_table (
  `id` SERIAL NOT NULL,
  `name` VARCHAR(100) NOT NULL,
  `created_date` DATETIME NOT NULL,
  PRIMARY KEY (`id`)
);

To save a list of items into this table, the code I am using:

String initialSql = "INSERT INTO test_table(`name`,`created_date`) VALUES ";

    List<String> values =
        dummyEntities.stream()
            .map(dummyEntity -> "('" + dummyEntity.getName() + "','"
                + dummyEntity.getCreatedDate().atZoneSameInstant(ZoneId.of("UTC")).toLocalDateTime().toString() + "')")
            .collect(Collectors.toList());

    String sqlToExecute =  initialSql + String.join(",", values);
    client.execute(sqlToExecute)
             .//Then what?

The generated SQL statement (from DEBUG Logs):

2020-09-15 18:59:32.613 DEBUG 44801 --- [actor-tcp-nio-1] o.s.d.r2dbc.core.DefaultDatabaseClient   : Executing SQL statement [INSERT INTO test_table(`name`,`created_date`) VALUES ('Abhi57','2020-09-15T13:29:29.951964'),('Abhi92','2020-09-15T13:29:29.952023')]

I even tried using ConnectionFactory, still no clue

    Mono.from(connectionFactory.create())
        .map(Connection::createBatch)
        .map(batch -> {
          dummyEntities.forEach(dummyEntity -> {
            String sql = String.format("INSERT INTO `test_table` (`name`,`created_date`) VALUES ('%s','%s');", dummyEntity.getName(),
                dummyEntity.getCreatedDate().atZoneSameInstant(ZoneId.of("UTC")).toLocalDateTime().toString());
            batch.add(sql);
          });
          return batch;
        })
        .flatMap(batch -> Mono.from(batch.execute()))
        .//Then what?

For reference, dummyEntities variable holds a list of DummyEntity objects. And the DummyEntity class looks like this:

@Table("test_table")
public class DummyEntity implements Persistable<Long> {

  @Id
  @Column("id")
  private Long id;

  @Column("name")
  private String name;

  @Column("created_date")
  private OffsetDateTime createdDate;

  //Getter,Setter
  @Override
  public boolean isNew() {
    return id == null;
  }
}

Dependencies used: 2.3.2.RELEASE

    implementation 'org.springframework.boot:spring-boot-starter-webflux'
    implementation 'org.springframework.boot:spring-boot-starter-data-r2dbc'
    implementation 'dev.miku:r2dbc-mysql:0.8.2.RELEASE'

Answer 1

Using the original ConnectionFacotory is easy to get generated ids.

I have tried to use ConnectionFactory to get the generated ids, worked as expected.

                    .thenMany(
                            Mono.from(conn)
                                    .flatMapMany(
                                            c -> c.createStatement(insertSql)
                                                    .returnGeneratedValues("id")
                                                    .execute()

                                    )
                    )
                    .flatMap(data -> Flux.from(data.map((row, rowMetadata) -> row.get("id"))))
                    .doOnNext(id -> log.info("generated id: {}", id))

The complete code example is here.

It prints logs in the console like this.

2020-09-19 10:43:30,815 INFO [main] com.example.demo.H2Tests$Sql:89 generated id: 1
2020-09-19 10:43:30,815 INFO [main] com.example.demo.H2Tests$Sql:89 generated id: 2

And I think the new DatabaseClient in Spring framework 5.3 is just a thin wrapper of the connectionfactories,and use a filter to get generated ids.

databaseClient.sql("INSERT INTO  posts (title, content, metadata) VALUES (:title, :content, :metadata)")
.filter((statement, executeFunction) -> statement.returnGeneratedValues("id").execute())

Check the complete example codes(but this example only retrieve a single id).

Answer 2

Late response, but since i faced the same issue on R2DBC PostgreSQL, here is a possible solution:

// -- generate 100 Product entities and mark them as CREATE
// -- Product is implementing Persistable<UUID> and isNew() will return TRUE for CREATE and FALSE for UPDATE 
List<Product> list = Stream.generate(() -> Product.builder()
    .articleNumber(UUID.randomUUID().toString().substring(0, 10))
    .name(UUID.randomUUID().toString())
    .images(Arrays.asList("1", "2"))
    .saveAction(SaveAction.CREATE)
    .build()).limit(100).collect(Collectors.toList());

// -- create the PostgresqlBatch and return the Result flux
Flux<? extends Result> flux = databaseClient.inConnectionMany(connection -> {
  Batch batch = connection.createBatch();

  list.forEach(p -> {
    batch.add("INSERT INTO product(\"id\",\"article_number\",\"name\") VALUES ('" + p.getId() + "','" + p.getArticleNumber() + "','" + p
        .getName() + "') RETURNING id, name, article_number");
  });

  return Flux.from(batch.execute());
});

// -- transform the Result flux into a Product flux
return flux.flatMap(result -> result.map((row, rowMetadata) -> Product.builder()
    .id(row.get("id", UUID.class))
    .name(row.get("name", String.class))
    .articleNumber(row.get("article_number", String.class))
    .build()));

The secret is RETURNING column1, column2, column3 added to each insert statement.