Allow keyspace customization for Cassandra tables and UDTs.

We now allow setting the keyspace on `@Table` and `@UserDefinedType` to interact with tables and types residing in a specific keyspace. @table(keyspace = "some_ks") class Person { … } Keyspace values can be either string literals or value expressions (#{…}, ${…}) to use the context or config properties to compute the actual keyspace. Special honors go to @tomekl007 and @mipo256 for their contributions. See #1400 See #279 Closes #921
spring-projects · Jul 25, 2024 · e154ef6 · e154ef6
1 parent 0816576
commit e154ef6
Show file tree

Hide file tree

Showing 60 changed files with 1,007 additions and 131 deletions.
diff --git a/...ssandra/src/main/java/org/springframework/data/cassandra/core/AsyncCassandraTemplate.java b/...ssandra/src/main/java/org/springframework/data/cassandra/core/AsyncCassandraTemplate.java
@@ -722,8 +722,10 @@ public CompletableFuture<Void> truncate(Class<?> entityClass) {
 
 		Assert.notNull(entityClass, "Entity type must not be null");
 
-		CqlIdentifier tableName = getTableName(entityClass);
-		Truncate truncate = QueryBuilder.truncate(tableName);
+		CassandraPersistentEntity<?> entity = getRequiredPersistentEntity(entityClass);
+		CqlIdentifier tableName = entity.getTableName();
+
+		Truncate truncate = QueryBuilder.truncate(entity.getKeyspace(), tableName);
 		SimpleStatement statement = truncate.build();
 
 		maybeEmitEvent(() -> new BeforeDeleteEvent<>(statement, entityClass, tableName));

diff --git a/.../java/org/springframework/data/cassandra/core/CassandraPersistentEntitySchemaDropper.java b/.../java/org/springframework/data/cassandra/core/CassandraPersistentEntitySchemaDropper.java
@@ -32,6 +32,7 @@
 import org.springframework.util.MultiValueMap;
 
 import com.datastax.oss.driver.api.core.CqlIdentifier;
+import com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata;
 import com.datastax.oss.driver.api.core.metadata.schema.RelationMetadata;
 import com.datastax.oss.driver.api.core.type.DataType;
 import com.datastax.oss.driver.api.core.type.ListType;
@@ -80,12 +81,17 @@ public CassandraPersistentEntitySchemaDropper(CassandraMappingContext mappingCon
 	 */
 	public void dropTables(boolean dropUnused) {
 
+		KeyspaceMetadata keyspaceMetadata = this.cassandraAdminOperations.getKeyspaceMetadata();
+		Set<CqlIdentifier> canRecreate = this.mappingContext.getTableEntities().stream()
+				.filter(it -> isInKeyspace(it, keyspaceMetadata)).map(CassandraPersistentEntity::getTableName)
+				.collect(Collectors.toSet());
+
 		this.cassandraAdminOperations.getKeyspaceMetadata() //
 				.getTables() //
 				.values() //
 				.stream() //
 				.map(RelationMetadata::getName) //
-				.filter(table -> dropUnused || this.mappingContext.usesTable(table)) //
+				.filter(table -> canRecreate.contains(table) || (dropUnused && !mappingContext.usesTable(table))) //
 				.forEach(this.cassandraAdminOperations::dropTable);
 	}
 
@@ -98,18 +104,28 @@ public void dropTables(boolean dropUnused) {
 	 */
 	public void dropUserTypes(boolean dropUnused) {
 
+		KeyspaceMetadata keyspaceMetadata = this.cassandraAdminOperations.getKeyspaceMetadata();
 		Set<CqlIdentifier> canRecreate = this.mappingContext.getUserDefinedTypeEntities().stream()
-				.map(CassandraPersistentEntity::getTableName).collect(Collectors.toSet());
+				.filter(it -> isInKeyspace(it, keyspaceMetadata)).map(CassandraPersistentEntity::getTableName)
+				.collect(Collectors.toSet());
 
-		Collection<UserDefinedType> userTypes = this.cassandraAdminOperations.getKeyspaceMetadata().getUserDefinedTypes()
-				.values();
+		Collection<UserDefinedType> userTypes = keyspaceMetadata.getUserDefinedTypes().values();
 
 		getUserTypesToDrop(userTypes) //
 				.stream() //
 				.filter(it -> canRecreate.contains(it) || (dropUnused && !mappingContext.usesUserType(it))) //
 				.forEach(this.cassandraAdminOperations::dropUserType);
 	}
 
+	private static boolean isInKeyspace(CassandraPersistentEntity<?> entity, KeyspaceMetadata keyspaceMetadata) {
+
+		if (entity.hasKeyspace() && keyspaceMetadata.getName().equals(entity.getRequiredKeyspace())) {
+			return true;
+		}
+
+		return !entity.hasKeyspace();
+	}
+
 	/**
 	 * Create {@link List} of {@link CqlIdentifier} with User-Defined type names to drop considering dependencies between
 	 * UDTs.

diff --git a/...ta-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraTemplate.java b/...ta-cassandra/src/main/java/org/springframework/data/cassandra/core/CassandraTemplate.java
@@ -748,8 +748,10 @@ public void truncate(Class<?> entityClass) {
 
 		Assert.notNull(entityClass, "Entity type must not be null");
 
-		CqlIdentifier tableName = getTableName(entityClass);
-		Truncate truncate = QueryBuilder.truncate(tableName);
+		CassandraPersistentEntity<?> entity = getRequiredPersistentEntity(entityClass);
+		CqlIdentifier tableName = entity.getTableName();
+
+		Truncate truncate = QueryBuilder.truncate(entity.getKeyspace(), tableName);
 		SimpleStatement statement = truncate.build();
 
 		maybeEmitEvent(() -> new BeforeDeleteEvent<>(statement, entityClass, tableName));

diff --git a/...ndra/src/main/java/org/springframework/data/cassandra/core/ReactiveCassandraTemplate.java b/...ndra/src/main/java/org/springframework/data/cassandra/core/ReactiveCassandraTemplate.java
@@ -728,8 +728,10 @@ public Mono<Void> truncate(Class<?> entityClass) {
 
 		Assert.notNull(entityClass, "Entity type must not be null");
 
-		CqlIdentifier tableName = getTableName(entityClass);
-		Truncate truncate = QueryBuilder.truncate(tableName);
+		CassandraPersistentEntity<?> entity = getRequiredPersistentEntity(entityClass);
+		CqlIdentifier tableName = entity.getTableName();
+
+		Truncate truncate = QueryBuilder.truncate(entity.getKeyspace(), tableName);
 		SimpleStatement statement = truncate.build();
 
 		Mono<Boolean> result = doExecute(statement, ReactiveResultSet::wasApplied)

diff --git a/...ata-cassandra/src/main/java/org/springframework/data/cassandra/core/StatementFactory.java b/...ata-cassandra/src/main/java/org/springframework/data/cassandra/core/StatementFactory.java
@@ -110,7 +110,7 @@ public class StatementFactory {
 
 	private final UpdateMapper updateMapper;
 
-	private KeyspaceProvider keyspaceProvider = KeyspaceProviders.EMPTY_KEYSPACE;
+	private KeyspaceProvider keyspaceProvider = KeyspaceProviders.ENTITY_KEYSPACE;
 
 	/**
 	 * Create {@link StatementFactory} given {@link CassandraConverter}.
@@ -1252,7 +1252,25 @@ public interface KeyspaceProvider {
 
 	}
 
+	/**
+	 * Implementations of {@link KeyspaceProvider}.
+	 */
 	enum KeyspaceProviders implements KeyspaceProvider {
+
+		/**
+		 * Derive the keyspace from the given {@link CassandraPersistentEntity}.
+		 */
+		ENTITY_KEYSPACE {
+			@Nullable
+			@Override
+			public CqlIdentifier getKeyspace(CassandraPersistentEntity<?> entity, CqlIdentifier tableName) {
+				return entity.getKeyspace();
+			}
+		},
+
+		/**
+		 * Use the session's keyspace.
+		 */
 		EMPTY_KEYSPACE {
 			@Nullable
 			@Override

diff --git a/.../main/java/org/springframework/data/cassandra/core/convert/DefaultColumnTypeResolver.java b/.../main/java/org/springframework/data/cassandra/core/convert/DefaultColumnTypeResolver.java
@@ -486,10 +486,19 @@ private CodecRegistry getCodecRegistry() {
 	}
 
 	private UserDefinedType getUserType(CassandraPersistentEntity<?> persistentEntity, boolean frozen) {
-		return getUserType(persistentEntity.getTableName()).copy(frozen);
+		return (persistentEntity.hasKeyspace()
+				? getUserType(persistentEntity.getRequiredKeyspace(), persistentEntity.getTableName())
+				: getUserType(persistentEntity.getTableName())).copy(frozen);
 	}
 
 	private UserDefinedType getUserType(String userTypeName) {
+
+		if (userTypeName.contains(".") && !userTypeName.contains("\"")) {
+
+			String[] split = userTypeName.split("\\.");
+			return getUserType(CqlIdentifier.fromCql(split[0]), CqlIdentifier.fromCql(split[1]));
+		}
+
 		return getUserType(CqlIdentifier.fromCql(userTypeName));
 	}
 
@@ -504,6 +513,17 @@ private UserDefinedType getUserType(CqlIdentifier userTypeName) {
 		return type;
 	}
 
+	private UserDefinedType getUserType(CqlIdentifier keyspace, CqlIdentifier userTypeName) {
+
+		UserDefinedType type = userTypeResolver.resolveType(keyspace, userTypeName);
+
+		if (type == null) {
+			throw new MappingException(String.format("User type [%s] in keyspace [%s] not found", userTypeName, keyspace));
+		}
+
+		return type;
+	}
+
 	private static void assertTypeArguments(int args, int expected) {
 
 		if (args != expected) {

diff --git a/.../main/java/org/springframework/data/cassandra/core/convert/IndexSpecificationFactory.java b/.../main/java/org/springframework/data/cassandra/core/convert/IndexSpecificationFactory.java
@@ -34,9 +34,12 @@
 import org.springframework.data.cassandra.core.mapping.SASI.Normalization;
 import org.springframework.data.cassandra.core.mapping.SASI.StandardAnalyzed;
 import org.springframework.data.mapping.MappingException;
+import org.springframework.lang.Nullable;
 import org.springframework.util.ObjectUtils;
 import org.springframework.util.StringUtils;
 
+import com.datastax.oss.driver.api.core.CqlIdentifier;
+
 /**
  * Factory to create {@link org.springframework.data.cassandra.core.cql.keyspace.CreateIndexSpecification} based on
  * index-annotated {@link CassandraPersistentProperty properties}.
@@ -70,13 +73,15 @@ class IndexSpecificationFactory {
 	 * @param property must not be {@literal null}.
 	 * @return {@link List} of {@link CreateIndexSpecification}.
 	 */
-	static List<CreateIndexSpecification> createIndexSpecifications(CassandraPersistentProperty property) {
+	static List<CreateIndexSpecification> createIndexSpecifications(@Nullable CqlIdentifier keyspace,
+			CassandraPersistentProperty property) {
 
 		List<CreateIndexSpecification> indexes = new ArrayList<>();
 
 		if (property.isAnnotationPresent(Indexed.class)) {
 
-			CreateIndexSpecification index = createIndexSpecification(property.findAnnotation(Indexed.class), property);
+			CreateIndexSpecification index = createIndexSpecification(keyspace, property.findAnnotation(Indexed.class),
+					property);
 
 			if (property.isMapLike()) {
 				index.entries();
@@ -86,7 +91,7 @@ static List<CreateIndexSpecification> createIndexSpecifications(CassandraPersist
 		}
 
 		if (property.isAnnotationPresent(SASI.class)) {
-			indexes.add(createIndexSpecification(property.findAnnotation(SASI.class), property));
+			indexes.add(createIndexSpecification(keyspace, property.findAnnotation(SASI.class), property));
 		}
 
 		if (property.isMapLike()) {
@@ -113,41 +118,41 @@ static List<CreateIndexSpecification> createIndexSpecifications(CassandraPersist
 				}
 
 				if (keyIndex != null) {
-					indexes.add(createIndexSpecification(keyIndex, property).keys());
+					indexes.add(createIndexSpecification(keyspace, keyIndex, property).keys());
 				}
 
 				if (valueIndex != null) {
-					indexes.add(createIndexSpecification(valueIndex, property).values());
+					indexes.add(createIndexSpecification(keyspace, valueIndex, property).values());
 				}
 			}
 		}
 
 		return indexes;
 	}
 
-	static CreateIndexSpecification createIndexSpecification(Indexed annotation,
+	static CreateIndexSpecification createIndexSpecification(@Nullable CqlIdentifier keyspace, Indexed annotation,
 			CassandraPersistentProperty property) {
 
 		CreateIndexSpecification index;
 
 		if (StringUtils.hasText(annotation.value())) {
-			index = CreateIndexSpecification.createIndex(annotation.value());
+			index = CreateIndexSpecification.createIndex(keyspace, CqlIdentifier.fromCql(annotation.value()));
 		} else {
-			index = CreateIndexSpecification.createIndex();
+			index = CreateIndexSpecification.createIndex(keyspace, null);
 		}
 
 		return index.columnName(property.getRequiredColumnName());
 	}
 
-	private static CreateIndexSpecification createIndexSpecification(SASI annotation,
+	private static CreateIndexSpecification createIndexSpecification(@Nullable CqlIdentifier keyspace, SASI annotation,
 			CassandraPersistentProperty property) {
 
 		CreateIndexSpecification index;
 
 		if (StringUtils.hasText(annotation.value())) {
-			index = CreateIndexSpecification.createIndex(annotation.value());
+			index = CreateIndexSpecification.createIndex(keyspace, CqlIdentifier.fromCql(annotation.value()));
 		} else {
-			index = CreateIndexSpecification.createIndex();
+			index = CreateIndexSpecification.createIndex(keyspace, null);
 		}
 
 		index.using("org.apache.cassandra.index.sasi.SASIIndex") //

diff --git a/.../main/java/org/springframework/data/cassandra/core/convert/MappingCassandraConverter.java b/.../main/java/org/springframework/data/cassandra/core/convert/MappingCassandraConverter.java
@@ -141,7 +141,19 @@ public MappingCassandraConverter(CassandraMappingContext mappingContext) {
 
 		Assert.notNull(mappingContext, "CassandraMappingContext must not be null");
 
-		UserTypeResolver userTypeResolver = userTypeName -> getUserTypeResolver().resolveType(userTypeName);
+		UserTypeResolver userTypeResolver = new UserTypeResolver() {
+			@Nullable
+			@Override
+			public UserDefinedType resolveType(CqlIdentifier typeName) {
+				return MappingCassandraConverter.this.getUserTypeResolver().resolveType(typeName);
+			}
+
+			@Nullable
+			@Override
+			public UserDefinedType resolveType(CqlIdentifier keyspace, CqlIdentifier typeName) {
+				return MappingCassandraConverter.this.getUserTypeResolver().resolveType(keyspace, typeName);
+			}
+		};
 
 		this.mappingContext = mappingContext;
 		this.setCodecRegistry(mappingContext.getCodecRegistry());