diff --git a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/convert/QueryMapper.java b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/convert/QueryMapper.java index 22edb1308..dd1e4e4b6 100644 --- a/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/convert/QueryMapper.java +++ b/spring-data-cassandra/src/main/java/org/springframework/data/cassandra/core/convert/QueryMapper.java @@ -22,8 +22,8 @@ import java.util.HashSet; import java.util.List; import java.util.Optional; import java.util.Set; -import java.util.stream.Collectors; +import org.springframework.data.cassandra.core.mapping.BasicCassandraPersistentEntity; import org.springframework.data.cassandra.core.mapping.CassandraMappingContext; import org.springframework.data.cassandra.core.mapping.CassandraPersistentEntity; import org.springframework.data.cassandra.core.mapping.CassandraPersistentProperty; @@ -167,8 +167,14 @@ public class QueryMapper { Field field = createPropertyField(entity, column); - columns.getSelector(column).ifPresent(selector -> getCqlIdentifier(column, field) - .ifPresent(cqlIdentifier -> selectors.add(getMappedSelector(selector, cqlIdentifier)))); + columns.getSelector(column).ifPresent(selector -> { + + List mappedColumnNames = getCqlIdentifier(column, field); + + for (CqlIdentifier mappedColumnName : mappedColumnNames) { + selectors.add(getMappedSelector(selector, mappedColumnName)); + } + }); } if (columns.isEmpty()) { @@ -207,17 +213,15 @@ public class QueryMapper { FunctionCall functionCall = (FunctionCall) selector; - List mappedParameters = functionCall.getParameters().stream().map(obj -> { + FunctionCall mappedFunctionCall = FunctionCall.from(functionCall.getExpression(), + functionCall.getParameters().stream().map(obj -> { - if (obj instanceof Selector) { - return getMappedSelector((Selector) obj, cqlIdentifier); - } + if (obj instanceof Selector) { + return getMappedSelector((Selector) obj, cqlIdentifier); + } - return obj; - }) // - .collect(Collectors.toList()); - - FunctionCall mappedFunctionCall = FunctionCall.from(functionCall.getExpression(), mappedParameters.toArray()); + return obj; + }).toArray()); return functionCall.getAlias() // .map(mappedFunctionCall::as) // @@ -252,11 +256,10 @@ public class QueryMapper { for (ColumnName column : columns) { Field field = createPropertyField(entity, column); - field.getProperty().ifPresent(seen::add); columns.getSelector(column).filter(selector -> selector instanceof ColumnSelector) - .ifPresent(columnSelector -> getCqlIdentifier(column, field).ifPresent(columnNames::add)); + .ifPresent(columnSelector -> columnNames.addAll(getCqlIdentifier(column, field))); } if (columns.isEmpty()) { @@ -293,40 +296,49 @@ public class QueryMapper { Field field = createPropertyField(entity, columnName); - Order mappedOrder = getCqlIdentifier(columnName, field) - .map(cqlIdentifier -> new Order(order.getDirection(), cqlIdentifier.toString())).orElse(order); + List mappedColumnNames = getCqlIdentifier(columnName, field); - mappedOrders.add(mappedOrder); + if (mappedColumnNames.isEmpty()) { + mappedOrders.add(order); + } else { + for (CqlIdentifier mappedColumnName : mappedColumnNames) { + mappedOrders.add(new Order(order.getDirection(), mappedColumnName.toString())); + } + } } return Sort.by(mappedOrders); } - private Optional getCqlIdentifier(ColumnName column, Field field) { + private List getCqlIdentifier(ColumnName column, Field field) { + List identifiers = new ArrayList<>(1); try { if (field.getProperty().isPresent()) { - return field.getProperty().map(cassandraPersistentProperty -> { + CassandraPersistentProperty property = field.getProperty().get(); - if (cassandraPersistentProperty.isCompositePrimaryKey()) { - throw new IllegalArgumentException( - "Cannot use composite primary key directly. Reference a property of the composite primary key"); - } + if (property.isCompositePrimaryKey()) { - return cassandraPersistentProperty.getRequiredColumnName(); - }); + BasicCassandraPersistentEntity primaryKeyEntity = mappingContext.getRequiredPersistentEntity(property); + + primaryKeyEntity.forEach(it -> { + identifiers.add(it.getRequiredColumnName()); + }); + } else { + identifiers.add(property.getRequiredColumnName()); + } + } else if (column.getColumnName().isPresent()) { + identifiers.add(CqlIdentifier.fromCql(column.getColumnName().get())); + } else { + column.getCqlIdentifier().ifPresent(identifiers::add); } - if (column.getColumnName().isPresent()) { - return column.getColumnName().map(CqlIdentifier::fromCql); - } - - return column.getCqlIdentifier(); - } catch (IllegalStateException cause) { throw new IllegalArgumentException(cause.getMessage(), cause); } + + return identifiers; } Field createPropertyField(@Nullable CassandraPersistentEntity entity, ColumnName key) { diff --git a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/core/CassandraTemplateIntegrationTests.java b/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/core/CassandraTemplateIntegrationTests.java index c4e5429ef..07865a3b1 100644 --- a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/core/CassandraTemplateIntegrationTests.java +++ b/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/core/CassandraTemplateIntegrationTests.java @@ -197,6 +197,21 @@ class CassandraTemplateIntegrationTests extends AbstractKeyspaceCreatingIntegrat assertThat(loadedEntity.getComment()).isNotNull(); } + @Test // DATACASS-828 + void shouldSelectClosedProjectionWithCompositeKey() { + + CompositeKey key = new CompositeKey("Walter", "White"); + TypeWithCompositeKey user = new TypeWithCompositeKey(key, "comment"); + + template.insert(user); + + WithCompositeKeyProjection loaded = template.query(TypeWithCompositeKey.class).as(WithCompositeKeyProjection.class) + .firstValue(); + assertThat(loaded).isNotNull(); + + assertThat(loaded.getKey()).isEqualTo(key); + } + @Test // DATACASS-343 void shouldSelectOneByQuery() { @@ -733,6 +748,11 @@ class CassandraTemplateIntegrationTests extends AbstractKeyspaceCreatingIntegrat String comment; } + interface WithCompositeKeyProjection { + + CompositeKey getKey(); + } + @Data @PrimaryKeyClass @AllArgsConstructor diff --git a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/core/convert/QueryMapperUnitTests.java b/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/core/convert/QueryMapperUnitTests.java index 530a7e3df..08b21c61e 100644 --- a/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/core/convert/QueryMapperUnitTests.java +++ b/spring-data-cassandra/src/test/java/org/springframework/data/cassandra/core/convert/QueryMapperUnitTests.java @@ -18,6 +18,7 @@ package org.springframework.data.cassandra.core.convert; import static org.assertj.core.api.Assertions.*; import static org.mockito.Matchers.any; import static org.mockito.Mockito.*; +import static org.springframework.data.domain.Sort.Order.*; import lombok.AllArgsConstructor; import lombok.Data; @@ -328,13 +329,15 @@ public class QueryMapperUnitTests { assertThat(mappedObject).contains(new Order(Direction.ASC, "first_name")); } - @Test // DATACASS-343 - void shouldFailMappingSortByCompositePrimaryKeyClass() { + @Test // DATACASS-828 + void allowSortByCompositeKey() { Sort sort = Sort.by("key"); + Query.empty().columns(Columns.from("key")); - assertThatIllegalArgumentException().isThrownBy( - () -> queryMapper.getMappedSort(sort, mappingContext.getRequiredPersistentEntity(TypeWithKeyClass.class))); + Sort mappedSort = queryMapper.getMappedSort(sort, + mappingContext.getRequiredPersistentEntity(TypeWithKeyClass.class)); + assertThat(mappedSort).isEqualTo(Sort.by(asc("first_name"), asc("lastname"))); } @Test // DATACASS-343