Support Composite IDs in R2DBC

pom.xml

@@ -5,7 +5,7 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-relational-parent</artifactId>
-	<version>4.0.0-SNAPSHOT</version>
+	<version>4.0.0-2011-embedded-SNAPSHOT</version>
 	<packaging>pom</packaging>
 
 	<name>Spring Data Relational Parent</name>

spring-data-jdbc-distribution/pom.xml

@@ -14,7 +14,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-relational-parent</artifactId>
-		<version>4.0.0-SNAPSHOT</version>
+		<version>4.0.0-2011-embedded-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-jdbc/pom.xml

@@ -6,7 +6,7 @@
 	<modelVersion>4.0.0</modelVersion>
 
 	<artifactId>spring-data-jdbc</artifactId>
-	<version>4.0.0-SNAPSHOT</version>
+	<version>4.0.0-2011-embedded-SNAPSHOT</version>
 
 	<name>Spring Data JDBC</name>
 	<description>Spring Data module for JDBC repositories.</description>
@@ -15,7 +15,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-relational-parent</artifactId>
-		<version>4.0.0-SNAPSHOT</version>
+		<version>4.0.0-2011-embedded-SNAPSHOT</version>
 	</parent>
 
 	<properties>

spring-data-r2dbc/pom.xml

@@ -6,7 +6,7 @@
 	<modelVersion>4.0.0</modelVersion>
 
 	<artifactId>spring-data-r2dbc</artifactId>
-	<version>4.0.0-SNAPSHOT</version>
+	<version>4.0.0-2011-embedded-SNAPSHOT</version>
 
 	<name>Spring Data R2DBC</name>
 	<description>Spring Data module for R2DBC</description>
@@ -15,7 +15,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-relational-parent</artifactId>
-		<version>4.0.0-SNAPSHOT</version>
+		<version>4.0.0-2011-embedded-SNAPSHOT</version>
 	</parent>
 
 	<properties>

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/convert/MappingR2dbcConverter.java

@@ -57,6 +57,7 @@
  *
  * @author Mark Paluch
  * @author Oliver Drotbohm
+ * @author Jens Schauder
  */
 public class MappingR2dbcConverter extends MappingRelationalConverter implements R2dbcConverter {
 
@@ -189,8 +190,17 @@ private void writeInternal(Object source, OutboundRow sink, Class<?> userClass)
 		writeProperties(sink, entity, propertyAccessor);
 	}
 
+	/**
+	 * write the values of the properties of an {@link RelationalPersistentEntity} to an {@link OutboundRow}.
+	 * 
+	 * @param sink must not be {@literal null}.
+	 * @param entity must not be {@literal null}.
+	 * @param accessor used for accessing the property values of {@literal entity}. May be {@literal null}. A
+	 *          {@literal null} value is used when this is an embedded {@literal null} entity, resulting in all its
+	 *          property values to be {@literal null} as well.
+	 */
 	private void writeProperties(OutboundRow sink, RelationalPersistentEntity<?> entity,
-			PersistentPropertyAccessor<?> accessor) {
+			@Nullable PersistentPropertyAccessor<?> accessor) {
 
 		for (RelationalPersistentProperty property : entity) {
 
@@ -200,11 +210,27 @@ private void writeProperties(OutboundRow sink, RelationalPersistentEntity<?> ent
 
 			Object value;
 
-			if (property.isIdProperty()) {
-				IdentifierAccessor identifierAccessor = entity.getIdentifierAccessor(accessor.getBean());
-				value = identifierAccessor.getIdentifier();
+			if (accessor == null) {
+				value = null;
 			} else {
-				value = accessor.getProperty(property);
+				if (property.isIdProperty()) {
+					IdentifierAccessor identifierAccessor = entity.getIdentifierAccessor(accessor.getBean());
+					value = identifierAccessor.getIdentifier();
+				} else {
+					value = accessor.getProperty(property);
+				}
+			}
+
+			if (property.isEmbedded()) {
+
+				RelationalPersistentEntity<?> embeddedEntity = getMappingContext().getRequiredPersistentEntity(property);
+				PersistentPropertyAccessor<Object> embeddedAccessor = null;
+				if (value != null) {
+					embeddedAccessor = embeddedEntity.getPropertyAccessor(value);
+				}
+				writeProperties(sink, embeddedEntity, embeddedAccessor);
+
+				continue;
 			}
 
 			if (value == null) {

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/DefaultReactiveDataAccessStrategy.java

@@ -43,6 +43,7 @@
 import org.springframework.data.relational.core.dialect.ArrayColumns;
 import org.springframework.data.relational.core.dialect.Dialect;
 import org.springframework.data.relational.core.dialect.RenderContextFactory;
+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;
@@ -66,7 +67,7 @@ public class DefaultReactiveDataAccessStrategy implements ReactiveDataAccessStra
 	private final R2dbcDialect dialect;
 	private final R2dbcConverter converter;
 	private final UpdateMapper updateMapper;
-	private final MappingContext<RelationalPersistentEntity<?>, ? extends RelationalPersistentProperty> mappingContext;
+	private final RelationalMappingContext mappingContext;
 	private final StatementMapper statementMapper;
 	private final NamedParameterExpander expander = new NamedParameterExpander();
 
@@ -119,16 +120,14 @@ public static R2dbcConverter createConverter(R2dbcDialect dialect, Collection<?>
 	 * @param dialect the {@link R2dbcDialect} to use.
 	 * @param converter must not be {@literal null}.
 	 */
-	@SuppressWarnings("unchecked")
 	public DefaultReactiveDataAccessStrategy(R2dbcDialect dialect, R2dbcConverter converter) {
 
 		Assert.notNull(dialect, "Dialect must not be null");
 		Assert.notNull(converter, "RelationalConverter must not be null");
 
 		this.converter = converter;
 		this.updateMapper = new UpdateMapper(dialect, converter);
-		this.mappingContext = (MappingContext<RelationalPersistentEntity<?>, ? extends RelationalPersistentProperty>) this.converter
-				.getMappingContext();
+		this.mappingContext = (RelationalMappingContext) this.converter.getMappingContext();
 		this.dialect = dialect;
 
 		RenderContextFactory factory = new RenderContextFactory(dialect);
@@ -141,13 +140,22 @@ public List<SqlIdentifier> getAllColumns(Class<?> entityType) {
 
 		RelationalPersistentEntity<?> persistentEntity = getPersistentEntity(entityType);
 
+		return getAllColumns(persistentEntity);
+	}
+
+	private List<SqlIdentifier> getAllColumns(@Nullable RelationalPersistentEntity<?> persistentEntity) {
+
 		if (persistentEntity == null) {
 			return Collections.singletonList(SqlIdentifier.unquoted("*"));
 		}
 
 		List<SqlIdentifier> columnNames = new ArrayList<>();
 		for (RelationalPersistentProperty property : persistentEntity) {
-			columnNames.add(property.getColumnName());
+			if (property.isEmbedded()) {
+				columnNames.addAll(getAllColumns(mappingContext.getRequiredPersistentEntity(property)));
+			} else {
+				columnNames.add(property.getColumnName());
+			}
 		}
 
 		return columnNames;
@@ -159,12 +167,8 @@ public List<SqlIdentifier> getIdentifierColumns(Class<?> entityType) {
 		RelationalPersistentEntity<?> persistentEntity = getRequiredPersistentEntity(entityType);
 
 		List<SqlIdentifier> columnNames = new ArrayList<>();
-		for (RelationalPersistentProperty property : persistentEntity) {
-
-			if (property.isIdProperty()) {
-				columnNames.add(property.getColumnName());
-			}
-		}
+		mappingContext.getAggregatePath(persistentEntity).getTableInfo().idColumnInfos()
+				.forEach((__, ci) -> columnNames.add(ci.name()));
 
 		return columnNames;
 	}

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/DefaultStatementMapper.java

@@ -143,7 +143,7 @@ protected List<Expression> getSelectList(SelectSpec selectSpec, @Nullable Relati
 		List<Expression> mapped = new ArrayList<>(selectList.size());
 
 		for (Expression expression : selectList) {
-			mapped.add(updateMapper.getMappedObject(expression, entity));
+			mapped.addAll(updateMapper.getMappedObjects(expression, entity));
 		}
 
 		return mapped;

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/R2dbcEntityTemplate.java

@@ -23,6 +23,7 @@
 import reactor.core.publisher.Mono;
 
 import java.util.Collections;
+import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -96,6 +97,7 @@
  * @author Robert Heim
  * @author Sebastian Wieland
  * @author Mikhail Polivakha
+ * @author Jens Schauder
  * @since 1.1
  */
 public class R2dbcEntityTemplate implements R2dbcEntityOperations, BeanFactoryAware, ApplicationContextAware {
@@ -350,8 +352,8 @@ <T, P extends Publisher<T>> P doSelect(Query query, Class<?> entityClass, SqlIde
 		return (P) ((Flux<?>) result).concatMap(it -> maybeCallAfterConvert(it, tableName));
 	}
 
-	private <T> RowsFetchSpec<T> doSelect(Query query, Class<?> entityType, SqlIdentifier tableName,
-			Class<T> returnType, Function<? super Statement, ? extends Statement> filterFunction) {
+	private <T> RowsFetchSpec<T> doSelect(Query query, Class<?> entityType, SqlIdentifier tableName, Class<T> returnType,
+			Function<? super Statement, ? extends Statement> filterFunction) {
 
 		StatementMapper statementMapper = dataAccessStrategy.getStatementMapper().forType(entityType);
 
@@ -378,11 +380,8 @@ private <T> RowsFetchSpec<T> doSelect(Query query, Class<?> entityType, SqlIdent
 
 		PreparedOperation<?> operation = statementMapper.getMappedObject(selectSpec);
 
-		return getRowsFetchSpec(
-				databaseClient.sql(operation).filter(statementFilterFunction.andThen(filterFunction)),
-			entityType,
-			returnType
-		);
+		return getRowsFetchSpec(databaseClient.sql(operation).filter(statementFilterFunction.andThen(filterFunction)),
+				entityType, returnType);
 	}
 
 	@Override
@@ -622,16 +621,28 @@ private <T> Mono<T> doUpdate(T entity, SqlIdentifier tableName) {
 			return maybeCallBeforeSave(entityToUse, outboundRow, tableName) //
 					.flatMap(onBeforeSave -> {
 
-						SqlIdentifier idColumn = persistentEntity.getRequiredIdProperty().getColumnName();
-						Parameter id = outboundRow.remove(idColumn);
+						Map<SqlIdentifier, Object> idValues = new LinkedHashMap<>();
+						List<SqlIdentifier> identifierColumns = dataAccessStrategy.getIdentifierColumns(persistentEntity.getType());
+						Assert.state(!identifierColumns.isEmpty(), entityToUse + " has no Identifier. Update is not possible.");
+
+						identifierColumns.forEach(sqlIdentifier -> {
+							idValues.put(sqlIdentifier, outboundRow.remove(sqlIdentifier));
+						});
 
 						persistentEntity.forEach(p -> {
 							if (p.isInsertOnly()) {
 								outboundRow.remove(p.getColumnName());
 							}
 						});
 
-						Criteria criteria = Criteria.where(dataAccessStrategy.toSql(idColumn)).is(id);
+						Criteria criteria = null;
+						for (Map.Entry<SqlIdentifier, Object> idAndValue : idValues.entrySet()) {
+							if (criteria == null) {
+								criteria = Criteria.where(dataAccessStrategy.toSql(idAndValue.getKey())).is(idAndValue.getValue());
+							} else {
+								criteria = criteria.and(dataAccessStrategy.toSql(idAndValue.getKey())).is(idAndValue.getValue());
+							}
+						}
 
 						if (matchingVersionCriteria != null) {
 							criteria = criteria.and(matchingVersionCriteria);