R2DBC support for @sequence annotation

Author: mipo256

Diff for: spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/config/AbstractR2dbcConfiguration.java

@@ -39,12 +39,14 @@
 import org.springframework.data.r2dbc.core.DefaultReactiveDataAccessStrategy;
 import org.springframework.data.r2dbc.core.R2dbcEntityTemplate;
 import org.springframework.data.r2dbc.core.ReactiveDataAccessStrategy;
+import org.springframework.data.r2dbc.core.mapping.IdGeneratingBeforeSaveCallback;
 import org.springframework.data.r2dbc.dialect.DialectResolver;
 import org.springframework.data.r2dbc.dialect.R2dbcDialect;
 import org.springframework.data.r2dbc.mapping.R2dbcMappingContext;
 import org.springframework.data.relational.RelationalManagedTypes;
 import org.springframework.data.relational.core.mapping.DefaultNamingStrategy;
 import org.springframework.data.relational.core.mapping.NamingStrategy;
+import org.springframework.data.relational.core.mapping.RelationalMappingContext;
 import org.springframework.data.relational.core.mapping.Table;
 import org.springframework.data.util.TypeScanner;
 import org.springframework.lang.Nullable;
@@ -182,6 +184,18 @@ public R2dbcMappingContext r2dbcMappingContext(Optional<NamingStrategy> namingSt
 		return context;
 	}
 
+	/**
+	 * Register a {@link IdGeneratingBeforeSaveCallback} using
+	 * {@link #r2dbcMappingContext(Optional, R2dbcCustomConversions, RelationalManagedTypes)} and
+	 * {@link #databaseClient()}
+	 */
+	@Bean
+	public IdGeneratingBeforeSaveCallback idGeneratingBeforeSaveCallback(
+			RelationalMappingContext relationalMappingContext, DatabaseClient databaseClient) {
+		return new IdGeneratingBeforeSaveCallback(relationalMappingContext, getDialect(lookupConnectionFactory()),
+				databaseClient);
+	}
+
 	/**
 	 * Creates a {@link ReactiveDataAccessStrategy} using the configured
 	 * {@link #r2dbcConverter(R2dbcMappingContext, R2dbcCustomConversions) R2dbcConverter}.

Diff for: spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/mapping/IdGeneratingBeforeSaveCallback.java

@@ -0,0 +1,103 @@
+/*
+ * Copyright 2020-2025 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.data.r2dbc.core.mapping;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.jetbrains.annotations.NotNull;
+import org.reactivestreams.Publisher;
+import org.springframework.data.r2dbc.dialect.R2dbcDialect;
+import org.springframework.data.r2dbc.mapping.OutboundRow;
+import org.springframework.data.r2dbc.mapping.event.BeforeSaveCallback;
+import org.springframework.data.relational.core.mapping.RelationalMappingContext;
+import org.springframework.data.relational.core.mapping.RelationalPersistentEntity;
+import org.springframework.data.relational.core.sql.SqlIdentifier;
+import org.springframework.r2dbc.core.DatabaseClient;
+import org.springframework.r2dbc.core.Parameter;
+import org.springframework.util.Assert;
+
+import reactor.core.publisher.Mono;
+
+/**
+ * R2DBC Callback for generating ID via the database sequence.
+ *
+ * @author Mikhail Polivakha
+ */
+public class IdGeneratingBeforeSaveCallback implements BeforeSaveCallback<Object> {
+
+	private static final Log LOG = LogFactory.getLog(IdGeneratingBeforeSaveCallback.class);
+
+	private final RelationalMappingContext relationalMappingContext;
+	private final R2dbcDialect dialect;
+
+	private final DatabaseClient databaseClient;
+
+	public IdGeneratingBeforeSaveCallback(RelationalMappingContext relationalMappingContext, R2dbcDialect dialect,
+			DatabaseClient databaseClient) {
+		this.relationalMappingContext = relationalMappingContext;
+		this.dialect = dialect;
+		this.databaseClient = databaseClient;
+	}
+
+	@Override
+	public Publisher<Object> onBeforeSave(Object entity, OutboundRow row, SqlIdentifier table) {
+		Assert.notNull(entity, "The aggregate cannot be null at this point");
+
+		RelationalPersistentEntity<?> persistentEntity = relationalMappingContext.getPersistentEntity(entity.getClass());
+
+		if (!persistentEntity.hasIdProperty() || //
+				!persistentEntity.getIdProperty().hasSequence() || //
+				!persistentEntity.isNew(entity) //
+		) {
+			return Mono.just(entity);
+		}
+
+		SqlIdentifier idSequence = persistentEntity.getIdProperty().getSequence();
+
+		if (dialect.getIdGeneration().sequencesSupported()) {
+			return fetchIdFromSeq(entity, row, persistentEntity, idSequence);
+		} else {
+			illegalSequenceUsageWarning(entity);
+		}
+
+		return Mono.just(entity);
+	}
+
+	@NotNull
+	private Mono<Object> fetchIdFromSeq(Object entity, OutboundRow row, RelationalPersistentEntity<?> persistentEntity,
+			SqlIdentifier idSequence) {
+		String sequenceQuery = dialect.getIdGeneration().createSequenceQuery(idSequence);
+
+		return databaseClient //
+				.sql(sequenceQuery) //
+				.map((r, rowMetadata) -> r.get(0)) //
+				.one() //
+				.doOnNext(fetchedId -> { //
+					row.put(persistentEntity.getIdColumn().toSql(dialect.getIdentifierProcessing()), //
+							Parameter.from(fetchedId) //
+					);
+				}).map(o -> entity);
+	}
+
+	private static void illegalSequenceUsageWarning(Object entity) {
+		LOG.warn("""
+					It seems you're trying to insert an aggregate of type '%s' annotated with @Sequence, but the problem is RDBMS you're
+					working with does not support sequences as such. Falling back to identity columns
+					""".stripIndent().formatted(entity.getClass().getName())
+		);
+	}
+}

Diff for: spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/core/mapping/IdGeneratingBeforeSaveCallbackTest.java

@@ -0,0 +1,124 @@
+/*
+ * Copyright 2020-2025 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.data.r2dbc.core.mapping;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.util.function.BiFunction;
+
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+import org.reactivestreams.Publisher;
+import org.springframework.data.annotation.Id;
+import org.springframework.data.r2dbc.dialect.MySqlDialect;
+import org.springframework.data.r2dbc.dialect.PostgresDialect;
+import org.springframework.data.r2dbc.mapping.OutboundRow;
+import org.springframework.data.r2dbc.mapping.R2dbcMappingContext;
+import org.springframework.data.relational.core.mapping.Sequence;
+import org.springframework.data.relational.core.sql.SqlIdentifier;
+import org.springframework.r2dbc.core.DatabaseClient;
+import org.springframework.r2dbc.core.Parameter;
+
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+/**
+ * Unit tests for {@link IdGeneratingBeforeSaveCallback}.
+ *
+ * @author Mikhail Polivakha
+ */
+class IdGeneratingBeforeSaveCallbackTest {
+
+	@Test
+	void testIdGenerationIsNotSupported() {
+		R2dbcMappingContext r2dbcMappingContext = new R2dbcMappingContext();
+		r2dbcMappingContext.getPersistentEntity(SimpleEntity.class);
+		MySqlDialect dialect = MySqlDialect.INSTANCE;
+		DatabaseClient databaseClient = mock(DatabaseClient.class);
+
+		IdGeneratingBeforeSaveCallback callback = new IdGeneratingBeforeSaveCallback(r2dbcMappingContext, dialect,
+				databaseClient);
+
+		OutboundRow row = new OutboundRow("name", Parameter.from("my_name"));
+		SimpleEntity entity = new SimpleEntity();
+		Publisher<Object> publisher = callback.onBeforeSave(entity, row, SqlIdentifier.unquoted("simple_entity"));
+
+		StepVerifier.create(publisher).expectNext(entity).expectComplete().verify();
+		assertThat(row).hasSize(1); // id is not added
+	}
+
+	@Test
+	void testEntityIsNotAnnotatedWithSequence() {
+		R2dbcMappingContext r2dbcMappingContext = new R2dbcMappingContext();
+		r2dbcMappingContext.getPersistentEntity(SimpleEntity.class);
+		PostgresDialect dialect = PostgresDialect.INSTANCE;
+		DatabaseClient databaseClient = mock(DatabaseClient.class);
+
+		IdGeneratingBeforeSaveCallback callback = new IdGeneratingBeforeSaveCallback(r2dbcMappingContext, dialect,
+				databaseClient);
+
+		OutboundRow row = new OutboundRow("name", Parameter.from("my_name"));
+		SimpleEntity entity = new SimpleEntity();
+		Publisher<Object> publisher = callback.onBeforeSave(entity, row, SqlIdentifier.unquoted("simple_entity"));
+
+		StepVerifier.create(publisher).expectNext(entity).expectComplete().verify();
+		assertThat(row).hasSize(1); // id is not added
+	}
+
+	@Test
+	void testIdGeneratedFromSequenceHappyPath() {
+		R2dbcMappingContext r2dbcMappingContext = new R2dbcMappingContext();
+		r2dbcMappingContext.getPersistentEntity(WithSequence.class);
+		PostgresDialect dialect = PostgresDialect.INSTANCE;
+		DatabaseClient databaseClient = mock(DatabaseClient.class, RETURNS_DEEP_STUBS);
+		long generatedId = 1L;
+
+		when(databaseClient.sql(Mockito.anyString()).map(Mockito.any(BiFunction.class)).one()).thenReturn(
+				Mono.just(generatedId));
+
+		IdGeneratingBeforeSaveCallback callback = new IdGeneratingBeforeSaveCallback(r2dbcMappingContext, dialect,
+				databaseClient);
+
+		OutboundRow row = new OutboundRow("name", Parameter.from("my_name"));
+		WithSequence entity = new WithSequence();
+		Publisher<Object> publisher = callback.onBeforeSave(entity, row, SqlIdentifier.unquoted("simple_entity"));
+
+		StepVerifier.create(publisher).expectNext(entity).expectComplete().verify();
+		assertThat(row).hasSize(2)
+				.containsEntry(SqlIdentifier.unquoted("id"), Parameter.from(generatedId));
+	}
+
+	static class SimpleEntity {
+
+		@Id
+		private Long id;
+
+		private String name;
+	}
+
+	static class WithSequence {
+
+		@Id
+		@Sequence(sequence = "seq_name")
+		private Long id;
+
+		private String name;
+	}
+}

Diff for: spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/PostgresR2dbcRepositoryIntegrationTests.java

@@ -15,8 +15,13 @@
  */
 package org.springframework.data.r2dbc.repository;
 
-import io.r2dbc.postgresql.codec.Json;
-import io.r2dbc.spi.ConnectionFactory;
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.Collections;
+import java.util.Map;
+
+import javax.sql.DataSource;
+
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.api.extension.RegisterExtension;
@@ -32,22 +37,20 @@
 import org.springframework.data.r2dbc.repository.support.R2dbcRepositoryFactory;
 import org.springframework.data.r2dbc.testing.ExternalDatabase;
 import org.springframework.data.r2dbc.testing.PostgresTestSupport;
+import org.springframework.data.relational.core.mapping.Sequence;
 import org.springframework.data.relational.core.mapping.Table;
 import org.springframework.data.repository.reactive.ReactiveCrudRepository;
 import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.r2dbc.core.DatabaseClient;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit.jupiter.SpringExtension;
+
+import io.r2dbc.postgresql.codec.Json;
+import io.r2dbc.spi.ConnectionFactory;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
-import javax.sql.DataSource;
-import java.util.Collections;
-import java.util.Map;
-
-import static org.assertj.core.api.Assertions.*;
-
 /**
  * Integration tests for {@link LegoSetRepository} using {@link R2dbcRepositoryFactory} against Postgres.
  *
@@ -62,12 +65,14 @@ public class PostgresR2dbcRepositoryIntegrationTests extends AbstractR2dbcReposi
 
 	@Autowired WithJsonRepository withJsonRepository;
 
+	@Autowired WithIdFromSequenceRepository withIdFromSequenceRepository;
+
 	@Autowired WithHStoreRepository hstoreRepositoryWith;
 
 	@Configuration
 	@EnableR2dbcRepositories(considerNestedRepositories = true,
 			includeFilters = @Filter(
-					classes = { PostgresLegoSetRepository.class, WithJsonRepository.class, WithHStoreRepository.class },
+					classes = { PostgresLegoSetRepository.class, WithJsonRepository.class, WithHStoreRepository.class, WithIdFromSequenceRepository.class },
 					type = FilterType.ASSIGNABLE_TYPE))
 	static class IntegrationTestConfiguration extends AbstractR2dbcConfiguration {
 
@@ -151,6 +156,51 @@ void shouldSaveAndLoadJson() {
 		}).verifyComplete();
 	}
 
+	@Test
+	void shouldInsertWithAutoGeneratedId() {
+
+		JdbcTemplate template = new JdbcTemplate(createDataSource());
+
+		template.execute("DROP TABLE IF EXISTS with_id_from_sequence");
+		template.execute("CREATE SEQUENCE IF NOT EXISTS target_sequence START WITH 15");
+		template.execute("CREATE TABLE with_id_from_sequence(\n" //
+				+ "    id    BIGINT PRIMARY KEY,\n" //
+				+ "    name  TEXT NOT NULL" //
+				+ ");");
+
+		WithIdFromSequence entity = new WithIdFromSequence(null, "Jordane");
+		withIdFromSequenceRepository.save(entity).as(StepVerifier::create).expectNextCount(1).verifyComplete();
+
+		withIdFromSequenceRepository.findAll().as(StepVerifier::create).consumeNextWith(actual -> {
+
+			assertThat(actual.id).isNotNull().isEqualTo(15);
+			assertThat(actual.name).isEqualTo("Jordane");
+		}).verifyComplete();
+	}
+
+	@Test
+	void shouldUpdateNoIdGenerationHappens() {
+
+		JdbcTemplate template = new JdbcTemplate(createDataSource());
+
+		template.execute("DROP TABLE IF EXISTS with_id_from_sequence");
+		template.execute("CREATE SEQUENCE IF NOT EXISTS target_sequence");
+		template.execute("CREATE TABLE with_id_from_sequence(\n" //
+				+ "    id    BIGINT PRIMARY KEY,\n" //
+				+ "    name  TEXT NOT NULL" //
+				+ ");");
+		template.execute("INSERT INTO with_id_from_sequence VALUES(4, 'Alex');");
+
+		WithIdFromSequence entity = new WithIdFromSequence(4L, "NewName");
+		withIdFromSequenceRepository.save(entity).as(StepVerifier::create).expectNextCount(1).verifyComplete();
+
+		withJsonRepository.findAll().as(StepVerifier::create).consumeNextWith(actual -> {
+
+			assertThat(actual.jsonValue).isNotNull().isEqualTo(4);
+			assertThat(actual.jsonValue.asString()).isEqualTo("NewName");
+		}).verifyComplete();
+	}
+
 	@Test // gh-492
 	void shouldSaveAndLoadHStore() {
 
@@ -188,6 +238,24 @@ interface WithJsonRepository extends ReactiveCrudRepository<WithJson, Long> {
 
 	}
 
+	static class WithIdFromSequence {
+
+		@Id
+		@Sequence(sequence = "target_sequence")
+		Long id;
+
+		String name;
+
+		public WithIdFromSequence(Long id, String name) {
+			this.id = id;
+			this.name = name;
+		}
+	}
+
+	interface WithIdFromSequenceRepository extends ReactiveCrudRepository<WithIdFromSequence, Long> {
+
+	}
+
 	@Table("with_hstore")
 	static class WithHStore {