Skip to content

R2DBC @Sequence annotation support #2028

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -56,14 +56,14 @@ public Object onBeforeSave(Object aggregate, MutableAggregateChange<Object> aggr
return aggregate;
}

RelationalPersistentProperty property = entity.getRequiredIdProperty();
RelationalPersistentProperty idProperty = entity.getRequiredIdProperty();
PersistentPropertyAccessor<Object> accessor = entity.getPropertyAccessor(aggregate);

if (!entity.isNew(aggregate) || delegate.hasValue(property, accessor) || !property.hasSequence()) {
if (!entity.isNew(aggregate) || delegate.hasValue(idProperty, accessor) || !idProperty.hasSequence()) {
return aggregate;
}

delegate.generateSequenceValue(property, accessor);
delegate.generateSequenceValue(idProperty, accessor);

return accessor.getBean();
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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}.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -186,11 +186,11 @@ private void writeInternal(Object source, OutboundRow sink, Class<?> userClass)
RelationalPersistentEntity<?> entity = getRequiredPersistentEntity(userClass);
PersistentPropertyAccessor<?> propertyAccessor = entity.getPropertyAccessor(source);

writeProperties(sink, entity, propertyAccessor, entity.isNew(source));
writeProperties(sink, entity, propertyAccessor);
}

private void writeProperties(OutboundRow sink, RelationalPersistentEntity<?> entity,
PersistentPropertyAccessor<?> accessor, boolean isNew) {
PersistentPropertyAccessor<?> accessor) {

for (RelationalPersistentProperty property : entity) {

Expand All @@ -213,24 +213,22 @@ private void writeProperties(OutboundRow sink, RelationalPersistentEntity<?> ent
}

if (getConversions().isSimpleType(value.getClass())) {
writeSimpleInternal(sink, value, isNew, property);
writeSimpleInternal(sink, value, property);
} else {
writePropertyInternal(sink, value, isNew, property);
writePropertyInternal(sink, value, property);
}
}
}

private void writeSimpleInternal(OutboundRow sink, Object value, boolean isNew,
RelationalPersistentProperty property) {
private void writeSimpleInternal(OutboundRow sink, Object value, RelationalPersistentProperty property) {

Object result = getPotentiallyConvertedSimpleWrite(value);

sink.put(property.getColumnName(),
Parameter.fromOrEmpty(result, getPotentiallyConvertedSimpleNullType(property.getType())));
}

private void writePropertyInternal(OutboundRow sink, Object value, boolean isNew,
RelationalPersistentProperty property) {
private void writePropertyInternal(OutboundRow sink, Object value, RelationalPersistentProperty property) {

TypeInformation<?> valueType = TypeInformation.of(value.getClass());

Expand All @@ -239,7 +237,7 @@ private void writePropertyInternal(OutboundRow sink, Object value, boolean isNew
if (valueType.getActualType() != null && valueType.getRequiredActualType().isCollectionLike()) {

// pass-thru nested collections
writeSimpleInternal(sink, value, isNew, property);
writeSimpleInternal(sink, value, property);
return;
}

Expand Down Expand Up @@ -310,7 +308,6 @@ private Class<?> getPotentiallyConvertedSimpleNullType(Class<?> type) {

if (customTarget.isPresent()) {
return customTarget.get();

}

if (type.isEnum()) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -521,15 +521,15 @@ private void potentiallyRemoveId(RelationalPersistentEntity<?> persistentEntity,
return;
}

SqlIdentifier columnName = idProperty.getColumnName();
Parameter parameter = outboundRow.get(columnName);
SqlIdentifier idColumnName = idProperty.getColumnName();
Parameter parameter = outboundRow.get(idColumnName);

if (shouldSkipIdValue(parameter, idProperty)) {
outboundRow.remove(columnName);
if (shouldSkipIdValue(parameter)) {
outboundRow.remove(idColumnName);
}
}

private boolean shouldSkipIdValue(@Nullable Parameter value, RelationalPersistentProperty property) {
private boolean shouldSkipIdValue(@Nullable Parameter value) {

if (value == null || value.getValue() == null) {
return true;
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
/*
* 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.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.mapping.RelationalPersistentProperty;
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);
}

RelationalPersistentProperty property = persistentEntity.getIdProperty();
SqlIdentifier idSequence = property.getSequence();

if (dialect.getIdGeneration().sequencesSupported()) {
return fetchIdFromSeq(entity, row, persistentEntity, idSequence);
} else {
illegalSequenceUsageWarning(entity);
}

return Mono.just(entity);
}

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() //
.map(fetchedId -> { //
row.put( //
persistentEntity.getIdColumn().toSql(dialect.getIdentifierProcessing()), //
Parameter.from(fetchedId) //
);
return 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()));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ public BoundAssignments getMappedObject(BindMarkers markers, Update update, Tabl
* @param entity related {@link RelationalPersistentEntity}, can be {@literal null}.
* @return the mapped {@link BoundAssignments}.
*/
public BoundAssignments getMappedObject(BindMarkers markers, Map<SqlIdentifier, ? extends Object> assignments,
public BoundAssignments getMappedObject(BindMarkers markers, Map<SqlIdentifier, ?> assignments,
Table table, @Nullable RelationalPersistentEntity<?> entity) {

Assert.notNull(markers, "BindMarkers must not be null");
Expand Down
Original file line number Diff line number Diff line change
@@ -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;
}
}
Loading
Loading