Add support for R2DBC RelationalManagedTypes.

mp911de · mp911de · commit 70eca3a2caf2 · 2022-07-06T15:27:40.000+02:00
See #1279
diff --git a/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/config/AbstractR2dbcConfiguration.java b/spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/config/AbstractR2dbcConfiguration.java
@@ -18,16 +18,22 @@
 import io.r2dbc.spi.ConnectionFactory;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
+import java.util.Set;
 
 import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextAware;
 import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.ClassPathScanningCandidateComponentProvider;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.core.convert.converter.Converter;
+import org.springframework.core.type.filter.AnnotationTypeFilter;
 import org.springframework.data.convert.CustomConversions;
 import org.springframework.data.convert.CustomConversions.StoreConversions;
 import org.springframework.data.r2dbc.convert.MappingR2dbcConverter;
@@ -39,11 +45,15 @@
 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.conversion.BasicRelationalConverter;
 import org.springframework.data.relational.core.mapping.NamingStrategy;
+import org.springframework.data.relational.core.mapping.Table;
 import org.springframework.lang.Nullable;
 import org.springframework.r2dbc.core.DatabaseClient;
 import org.springframework.util.Assert;
+import org.springframework.util.ClassUtils;
+import org.springframework.util.StringUtils;
 
 /**
  * Base class for Spring Data R2DBC configuration containing bean declarations that must be registered for Spring Data
@@ -78,14 +88,43 @@ public void setApplicationContext(ApplicationContext applicationContext) throws
 	 */
 	public abstract ConnectionFactory connectionFactory();
 
+	/**
+	 * Returns the base packages to scan for R2DBC mapped entities at startup. Returns the package name of the
+	 * configuration class' (the concrete class, not this one here) by default. So if you have a
+	 * {@code com.acme.AppConfig} extending {@link AbstractR2dbcConfiguration} the base package will be considered
+	 * {@code com.acme} unless the method is overridden to implement alternate behavior.
+	 *
+	 * @return the base packages to scan for mapped {@link Table} classes or an empty collection to not enable scanning
+	 *         for entities.
+	 * @since 3.0
+	 */
+	protected Collection<String> getMappingBasePackages() {
+
+		Package mappingBasePackage = getClass().getPackage();
+		return Collections.singleton(mappingBasePackage == null ? null : mappingBasePackage.getName());
+	}
+
+	/**
+	 * Returns the a {@link RelationalManagedTypes} object holding the initial entity set.
+	 *
+	 * @return new instance of {@link RelationalManagedTypes}.
+	 * @throws ClassNotFoundException
+	 * @since 3.0
+	 */
+	@Bean
+	public RelationalManagedTypes r2dbcManagedTypes() throws ClassNotFoundException {
+		return RelationalManagedTypes.fromIterable(getInitialEntitySet());
+	}
+
 	/**
 	 * Return a {@link R2dbcDialect} for the given {@link ConnectionFactory}. This method attempts to resolve a
 	 * {@link R2dbcDialect} from {@link io.r2dbc.spi.ConnectionFactoryMetadata}. Override this method to specify a dialect
 	 * instead of attempting to resolve one.
 	 *
 	 * @param connectionFactory the configured {@link ConnectionFactory}.
 	 * @return the resolved {@link R2dbcDialect}.
-	 * @throws org.springframework.data.r2dbc.dialect.DialectResolver.NoDialectException if the {@link R2dbcDialect} cannot be determined.
+	 * @throws org.springframework.data.r2dbc.dialect.DialectResolver.NoDialectException if the {@link R2dbcDialect}
+	 *           cannot be determined.
 	 */
 	public R2dbcDialect getDialect(ConnectionFactory connectionFactory) {
 		return DialectResolver.getDialect(connectionFactory);
@@ -131,17 +170,20 @@ public R2dbcEntityTemplate r2dbcEntityTemplate(DatabaseClient databaseClient,
 	 *
 	 * @param namingStrategy optional {@link NamingStrategy}. Use {@link NamingStrategy#INSTANCE} as fallback.
 	 * @param r2dbcCustomConversions customized R2DBC conversions.
+	 * @param r2dbcManagedTypes R2DBC managed types, typically discovered through {@link #r2dbcManagedTypes() an entity
+	 *          scan}.
 	 * @return must not be {@literal null}.
 	 * @throws IllegalArgumentException if any of the required args is {@literal null}.
 	 */
 	@Bean
 	public R2dbcMappingContext r2dbcMappingContext(Optional<NamingStrategy> namingStrategy,
-			R2dbcCustomConversions r2dbcCustomConversions) {
+			R2dbcCustomConversions r2dbcCustomConversions, RelationalManagedTypes r2dbcManagedTypes) {
 
 		Assert.notNull(namingStrategy, "NamingStrategy must not be null");
 
 		R2dbcMappingContext context = new R2dbcMappingContext(namingStrategy.orElse(NamingStrategy.INSTANCE));
 		context.setSimpleTypeHolder(r2dbcCustomConversions.getSimpleTypeHolder());
+		context.setManagedTypes(r2dbcManagedTypes);
 
 		return context;
 	}
@@ -239,4 +281,56 @@ ConnectionFactory lookupConnectionFactory() {
 
 		return connectionFactory();
 	}
+
+	/**
+	 * Scans the mapping base package for classes annotated with {@link Table}. By default, it scans for entities in all
+	 * packages returned by {@link #getMappingBasePackages()}.
+	 *
+	 * @see #getMappingBasePackages()
+	 * @return
+	 * @throws ClassNotFoundException
+	 * @since 3.0
+	 */
+	protected Set<Class<?>> getInitialEntitySet() throws ClassNotFoundException {
+
+		Set<Class<?>> initialEntitySet = new HashSet<>();
+
+		for (String basePackage : getMappingBasePackages()) {
+			initialEntitySet.addAll(scanForEntities(basePackage));
+		}
+
+		return initialEntitySet;
+	}
+
+	/**
+	 * Scans the given base package for entities, i.e. R2DBC-specific types annotated with {@link Table}.
+	 *
+	 * @param basePackage must not be {@literal null}.
+	 * @return
+	 * @throws ClassNotFoundException
+	 * @since 3.0
+	 */
+	protected Set<Class<?>> scanForEntities(String basePackage) throws ClassNotFoundException {
+
+		if (!StringUtils.hasText(basePackage)) {
+			return Collections.emptySet();
+		}
+
+		Set<Class<?>> initialEntitySet = new HashSet<>();
+
+		if (StringUtils.hasText(basePackage)) {
+
+			ClassPathScanningCandidateComponentProvider componentProvider = new ClassPathScanningCandidateComponentProvider(
+					false);
+			componentProvider.addIncludeFilter(new AnnotationTypeFilter(Table.class));
+
+			for (BeanDefinition candidate : componentProvider.findCandidateComponents(basePackage)) {
+
+				initialEntitySet
+						.add(ClassUtils.forName(candidate.getBeanClassName(), AbstractR2dbcConfiguration.class.getClassLoader()));
+			}
+		}
+
+		return initialEntitySet;
+	}
 }
diff --git a/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/config/R2dbcConfigurationIntegrationTests.java b/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/config/R2dbcConfigurationIntegrationTests.java
@@ -27,7 +27,10 @@
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.data.r2dbc.mapping.R2dbcMappingContext;
+import org.springframework.data.relational.RelationalManagedTypes;
 import org.springframework.r2dbc.core.DatabaseClient;
+import org.springframework.test.util.ReflectionTestUtils;
 
 /**
  * Tests for {@link AbstractR2dbcConfiguration}.
@@ -88,6 +91,22 @@ void shouldRegisterConnectionFactory() {
 		context.stop();
 	}
 
+	@Test // GH-1279
+	void shouldScanForInitialEntities() {
+
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(
+				CustomConnectionFactoryBeanNameConfiguration.class);
+
+		R2dbcMappingContext mappingContext = context.getBean(R2dbcMappingContext.class);
+
+		RelationalManagedTypes managedTypes = (RelationalManagedTypes) ReflectionTestUtils.getField(mappingContext,
+				"managedTypes");
+
+		assertThat(managedTypes.toList()).contains(H2IntegrationTests.LegoSet.class, TopLevelEntity.class);
+
+		context.stop();
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	static class NonBeanConnectionFactoryConfiguration extends AbstractR2dbcConfiguration {
 
diff --git a/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/config/TopLevelEntity.java b/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/config/TopLevelEntity.java
@@ -0,0 +1,26 @@
+/*
+ * Copyright 2022 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.config;
+
+import org.springframework.data.relational.core.mapping.Table;
+
+/**
+ * Empty test entity annotated with {@code @Table}.
+ *
+ * @author Mark Paluch
+ */
+@Table
+class TopLevelEntity {}
diff --git a/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/H2R2dbcRepositoryWithMixedCaseNamesIntegrationTests.java b/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/H2R2dbcRepositoryWithMixedCaseNamesIntegrationTests.java
@@ -31,6 +31,7 @@
 import org.springframework.data.r2dbc.mapping.R2dbcMappingContext;
 import org.springframework.data.r2dbc.repository.config.EnableR2dbcRepositories;
 import org.springframework.data.r2dbc.testing.H2TestSupport;
+import org.springframework.data.relational.RelationalManagedTypes;
 import org.springframework.data.relational.core.mapping.NamingStrategy;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit.jupiter.SpringExtension;
@@ -59,9 +60,10 @@ public ConnectionFactory connectionFactory() {
 
 		@Override
 		public R2dbcMappingContext r2dbcMappingContext(Optional<NamingStrategy> namingStrategy,
-				R2dbcCustomConversions r2dbcCustomConversions) {
+				R2dbcCustomConversions r2dbcCustomConversions, RelationalManagedTypes r2dbcManagedTypes) {
 
-			R2dbcMappingContext r2dbcMappingContext = super.r2dbcMappingContext(namingStrategy, r2dbcCustomConversions);
+			R2dbcMappingContext r2dbcMappingContext = super.r2dbcMappingContext(namingStrategy, r2dbcCustomConversions,
+					r2dbcManagedTypes);
 			r2dbcMappingContext.setForceQuote(true);
 
 			return r2dbcMappingContext;
diff --git a/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/MariaDbR2dbcRepositoryWithMixedCaseNamesIntegrationTests.java b/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/MariaDbR2dbcRepositoryWithMixedCaseNamesIntegrationTests.java
@@ -33,6 +33,7 @@
 import org.springframework.data.r2dbc.repository.config.EnableR2dbcRepositories;
 import org.springframework.data.r2dbc.testing.ExternalDatabase;
 import org.springframework.data.r2dbc.testing.MariaDbTestSupport;
+import org.springframework.data.relational.RelationalManagedTypes;
 import org.springframework.data.relational.core.mapping.NamingStrategy;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit.jupiter.SpringExtension;
@@ -63,9 +64,10 @@ public ConnectionFactory connectionFactory() {
 
 		@Override
 		public R2dbcMappingContext r2dbcMappingContext(Optional<NamingStrategy> namingStrategy,
-				R2dbcCustomConversions r2dbcCustomConversions) {
+				R2dbcCustomConversions r2dbcCustomConversions, RelationalManagedTypes r2dbcManagedTypes) {
 
-			R2dbcMappingContext r2dbcMappingContext = super.r2dbcMappingContext(namingStrategy, r2dbcCustomConversions);
+			R2dbcMappingContext r2dbcMappingContext = super.r2dbcMappingContext(namingStrategy, r2dbcCustomConversions,
+					r2dbcManagedTypes);
 			r2dbcMappingContext.setForceQuote(true);
 
 			return r2dbcMappingContext;
diff --git a/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/OracleR2dbcRepositoryWithMixedCaseNamesIntegrationTests.java b/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/OracleR2dbcRepositoryWithMixedCaseNamesIntegrationTests.java
@@ -35,6 +35,7 @@
 import org.springframework.data.r2dbc.testing.EnabledOnClass;
 import org.springframework.data.r2dbc.testing.ExternalDatabase;
 import org.springframework.data.r2dbc.testing.OracleTestSupport;
+import org.springframework.data.relational.RelationalManagedTypes;
 import org.springframework.data.relational.core.mapping.NamingStrategy;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit.jupiter.SpringExtension;
@@ -66,9 +67,10 @@ public ConnectionFactory connectionFactory() {
 
 		@Override
 		public R2dbcMappingContext r2dbcMappingContext(Optional<NamingStrategy> namingStrategy,
-				R2dbcCustomConversions r2dbcCustomConversions) {
+				R2dbcCustomConversions r2dbcCustomConversions, RelationalManagedTypes r2dbcManagedTypes) {
 
-			R2dbcMappingContext r2dbcMappingContext = super.r2dbcMappingContext(namingStrategy, r2dbcCustomConversions);
+			R2dbcMappingContext r2dbcMappingContext = super.r2dbcMappingContext(namingStrategy, r2dbcCustomConversions,
+					r2dbcManagedTypes);
 			r2dbcMappingContext.setForceQuote(true);
 
 			return r2dbcMappingContext;
diff --git a/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/PostgresR2dbcRepositoryWithMixedCaseNamesIntegrationTests.java b/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/PostgresR2dbcRepositoryWithMixedCaseNamesIntegrationTests.java
@@ -34,6 +34,7 @@
 import org.springframework.data.r2dbc.repository.config.EnableR2dbcRepositories;
 import org.springframework.data.r2dbc.testing.ExternalDatabase;
 import org.springframework.data.r2dbc.testing.PostgresTestSupport;
+import org.springframework.data.relational.RelationalManagedTypes;
 import org.springframework.data.relational.core.mapping.NamingStrategy;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit.jupiter.SpringExtension;
@@ -64,9 +65,10 @@ public ConnectionFactory connectionFactory() {
 
 		@Override
 		public R2dbcMappingContext r2dbcMappingContext(Optional<NamingStrategy> namingStrategy,
-				R2dbcCustomConversions r2dbcCustomConversions) {
+				R2dbcCustomConversions r2dbcCustomConversions, RelationalManagedTypes r2dbcManagedTypes) {
 
-			R2dbcMappingContext r2dbcMappingContext = super.r2dbcMappingContext(namingStrategy, r2dbcCustomConversions);
+			R2dbcMappingContext r2dbcMappingContext = super.r2dbcMappingContext(namingStrategy, r2dbcCustomConversions,
+					r2dbcManagedTypes);
 			r2dbcMappingContext.setForceQuote(true);
 
 			return r2dbcMappingContext;
diff --git a/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/SqlServerR2dbcRepositoryWithMixedCaseNamesIntegrationTests.java b/spring-data-r2dbc/src/test/java/org/springframework/data/r2dbc/repository/SqlServerR2dbcRepositoryWithMixedCaseNamesIntegrationTests.java
@@ -34,6 +34,7 @@
 import org.springframework.data.r2dbc.repository.config.EnableR2dbcRepositories;
 import org.springframework.data.r2dbc.testing.ExternalDatabase;
 import org.springframework.data.r2dbc.testing.SqlServerTestSupport;
+import org.springframework.data.relational.RelationalManagedTypes;
 import org.springframework.data.relational.core.mapping.NamingStrategy;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit.jupiter.SpringExtension;
@@ -64,9 +65,10 @@ public ConnectionFactory connectionFactory() {
 
 		@Override
 		public R2dbcMappingContext r2dbcMappingContext(Optional<NamingStrategy> namingStrategy,
-				R2dbcCustomConversions r2dbcCustomConversions) {
+				R2dbcCustomConversions r2dbcCustomConversions, RelationalManagedTypes r2dbcManagedTypes) {
 
-			R2dbcMappingContext r2dbcMappingContext = super.r2dbcMappingContext(namingStrategy, r2dbcCustomConversions);
+			R2dbcMappingContext r2dbcMappingContext = super.r2dbcMappingContext(namingStrategy, r2dbcCustomConversions,
+					r2dbcManagedTypes);
 			r2dbcMappingContext.setForceQuote(true);
 
 			return r2dbcMappingContext;
