Add support for RelationalManagedTypes.

See #1269

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/repository/config/AbstractJdbcConfiguration.java

@@ -16,33 +16,43 @@
 package org.springframework.data.jdbc.repository.config;
 
 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.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.NoSuchBeanDefinitionException;
+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.context.annotation.Lazy;
 import org.springframework.core.convert.converter.Converter;
+import org.springframework.core.type.filter.AnnotationTypeFilter;
 import org.springframework.data.convert.CustomConversions;
 import org.springframework.data.jdbc.core.JdbcAggregateOperations;
 import org.springframework.data.jdbc.core.JdbcAggregateTemplate;
 import org.springframework.data.jdbc.core.convert.*;
-import org.springframework.data.jdbc.core.convert.JdbcArrayColumns;
 import org.springframework.data.jdbc.core.dialect.JdbcDialect;
 import org.springframework.data.jdbc.core.mapping.JdbcMappingContext;
 import org.springframework.data.jdbc.core.mapping.JdbcSimpleTypes;
 import org.springframework.data.mapping.model.SimpleTypeHolder;
+import org.springframework.data.relational.RelationalManagedTypes;
 import org.springframework.data.relational.core.conversion.RelationalConverter;
 import org.springframework.data.relational.core.dialect.Dialect;
 import org.springframework.data.relational.core.mapping.NamingStrategy;
+import org.springframework.data.relational.core.mapping.Table;
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcOperations;
+import org.springframework.util.ClassUtils;
+import org.springframework.util.StringUtils;
 
 /**
  * Beans that must be registered for Spring Data JDBC to work.
@@ -63,19 +73,50 @@ public class AbstractJdbcConfiguration implements ApplicationContextAware {
 
 	private ApplicationContext applicationContext;
 
+	/**
+	 * Returns the base packages to scan for JDBC 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 AbstractJdbcConfiguration} 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 jdbcManagedTypes() throws ClassNotFoundException {
+		return RelationalManagedTypes.fromIterable(getInitialEntitySet());
+	}
+
 	/**
 	 * Register a {@link JdbcMappingContext} and apply an optional {@link NamingStrategy}.
 	 *
 	 * @param namingStrategy optional {@link NamingStrategy}. Use {@link NamingStrategy#INSTANCE} as fallback.
 	 * @param customConversions see {@link #jdbcCustomConversions()}.
+	 * @param jdbcManagedTypes JDBC managed types, typically discovered through {@link #jdbcManagedTypes() an entity
+	 *          scan}.
 	 * @return must not be {@literal null}.
 	 */
 	@Bean
 	public JdbcMappingContext jdbcMappingContext(Optional<NamingStrategy> namingStrategy,
-			JdbcCustomConversions customConversions) {
+			JdbcCustomConversions customConversions, RelationalManagedTypes jdbcManagedTypes) {
 
 		JdbcMappingContext mappingContext = new JdbcMappingContext(namingStrategy.orElse(NamingStrategy.INSTANCE));
 		mappingContext.setSimpleTypeHolder(customConversions.getSimpleTypeHolder());
+		mappingContext.setManagedTypes(jdbcManagedTypes);
 
 		return mappingContext;
 	}
@@ -190,4 +231,56 @@ public class AbstractJdbcConfiguration implements ApplicationContextAware {
 	public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
 		this.applicationContext = applicationContext;
 	}
+
+	/**
+	 * 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. JDBC-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(), AbstractJdbcConfiguration.class.getClassLoader()));
+			}
+		}
+
+		return initialEntitySet;
+	}
 }

spring-data-jdbc/src/test/java/org/springframework/data/jdbc/repository/config/AbstractJdbcConfigurationIntegrationTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019-2021 the original author or authors.
+ * Copyright 2019-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.
@@ -37,6 +37,7 @@ import org.springframework.data.jdbc.core.convert.DataAccessStrategy;
 import org.springframework.data.jdbc.core.convert.JdbcConverter;
 import org.springframework.data.jdbc.core.convert.JdbcCustomConversions;
 import org.springframework.data.jdbc.core.mapping.JdbcMappingContext;
+import org.springframework.data.relational.RelationalManagedTypes;
 import org.springframework.data.relational.core.dialect.Dialect;
 import org.springframework.data.relational.core.dialect.LimitClause;
 import org.springframework.data.relational.core.dialect.LockClause;
@@ -44,13 +45,15 @@ import org.springframework.data.relational.core.sql.render.SelectRenderContext;
 import org.springframework.jdbc.core.JdbcOperations;
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcOperations;
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
+import org.springframework.test.util.ReflectionTestUtils;
 
 /**
  * Integration tests for {@link AbstractJdbcConfiguration}.
  *
  * @author Oliver Drotbohm
+ * @author Mark Paluch
  */
-public class AbstractJdbcConfigurationIntegrationTests {
+class AbstractJdbcConfigurationIntegrationTests {
 
 	@Test // DATAJDBC-395
 	void configuresInfrastructureComponents() {
@@ -97,7 +100,22 @@ public class AbstractJdbcConfigurationIntegrationTests {
 		}, AbstractJdbcConfigurationUnderTest.class, Infrastructure.class);
 	}
 
-	protected static void assertApplicationContext(Consumer<ConfigurableApplicationContext> verification,
+	@Test // GH-1269
+	void detectsInitialEntities() {
+
+		assertApplicationContext(context -> {
+
+			JdbcMappingContext mappingContext = context.getBean(JdbcMappingContext.class);
+			RelationalManagedTypes managedTypes = (RelationalManagedTypes) ReflectionTestUtils.getField(mappingContext,
+					"managedTypes");
+
+			assertThat(managedTypes.toList()).contains(JdbcRepositoryConfigExtensionUnitTests.Sample.class,
+					TopLevelEntity.class);
+
+		}, AbstractJdbcConfigurationUnderTest.class, Infrastructure.class);
+	}
+
+	static void assertApplicationContext(Consumer<ConfigurableApplicationContext> verification,
 			Class<?>... configurationClasses) {
 
 		try (AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext()) {

spring-data-jdbc/src/test/java/org/springframework/data/jdbc/repository/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.jdbc.repository.config;
+
+import org.springframework.data.relational.core.mapping.Table;
+
+/**
+ * Empty test entity annotated with {@code @Table}.
+ *
+ * @author Mark Paluch
+ */
+@Table
+class TopLevelEntity {}

spring-data-relational/src/main/java/org/springframework/data/relational/RelationalManagedTypes.java

@@ -0,0 +1,83 @@
+/*
+ * 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.relational;
+
+import java.util.Arrays;
+import java.util.function.Consumer;
+
+import org.springframework.data.domain.ManagedTypes;
+
+/**
+ * Relational-specific extension to {@link ManagedTypes}.
+ *
+ * @author Mark Paluch
+ * @since 3.0
+ */
+public final class RelationalManagedTypes implements ManagedTypes {
+
+	private final ManagedTypes delegate;
+
+	private RelationalManagedTypes(ManagedTypes types) {
+		this.delegate = types;
+	}
+
+	/**
+	 * Wraps an existing {@link ManagedTypes} object with {@link RelationalManagedTypes}.
+	 *
+	 * @param managedTypes
+	 * @return
+	 */
+	public static RelationalManagedTypes from(ManagedTypes managedTypes) {
+		return new RelationalManagedTypes(managedTypes);
+	}
+
+	/**
+	 * Factory method used to construct {@link RelationalManagedTypes} from the given array of {@link Class types}.
+	 *
+	 * @param types array of {@link Class types} used to initialize the {@link ManagedTypes}; must not be {@literal null}.
+	 * @return new instance of {@link RelationalManagedTypes} initialized from {@link Class types}.
+	 */
+	public static RelationalManagedTypes from(Class<?>... types) {
+		return fromIterable(Arrays.asList(types));
+	}
+
+	/**
+	 * Factory method used to construct {@link RelationalManagedTypes} from the given, required {@link Iterable} of
+	 * {@link Class types}.
+	 *
+	 * @param types {@link Iterable} of {@link Class types} used to initialize the {@link ManagedTypes}; must not be
+	 *          {@literal null}.
+	 * @return new instance of {@link RelationalManagedTypes} initialized the given, required {@link Iterable} of
+	 *         {@link Class types}.
+	 */
+	public static RelationalManagedTypes fromIterable(Iterable<? extends Class<?>> types) {
+		return from(ManagedTypes.fromIterable(types));
+	}
+
+	/**
+	 * Factory method to return an empty {@link RelationalManagedTypes} object.
+	 *
+	 * @return an empty {@link RelationalManagedTypes} object.
+	 */
+	public static RelationalManagedTypes empty() {
+		return from(ManagedTypes.empty());
+	}
+
+	@Override
+	public void forEach(Consumer<Class<?>> action) {
+		delegate.forEach(action);
+	}
+}

spring-data-relational/src/main/java/org/springframework/data/relational/aot/RelationalManagedTypesBeanRegistrationAotProcessor.java

@@ -0,0 +1,38 @@
+/*
+ * 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.relational.aot;
+
+import org.springframework.data.aot.ManagedTypesBeanRegistrationAotProcessor;
+import org.springframework.data.relational.RelationalManagedTypes;
+import org.springframework.lang.Nullable;
+import org.springframework.util.ClassUtils;
+
+/**
+ * Relational-specific extension to {@link ManagedTypesBeanRegistrationAotProcessor}.
+ *
+ * @author Mark Paluch
+ * @since 3.0
+ */
+class RelationalManagedTypesBeanRegistrationAotProcessor extends ManagedTypesBeanRegistrationAotProcessor {
+
+	protected boolean isMatch(@Nullable Class<?> beanType, @Nullable String beanName) {
+		return this.matchesByType(beanType);
+	}
+
+	protected boolean matchesByType(@Nullable Class<?> beanType) {
+		return beanType != null && ClassUtils.isAssignable(RelationalManagedTypes.class, beanType);
+	}
+}

spring-data-relational/src/main/java/org/springframework/data/relational/aot/package-info.java

@@ -0,0 +1,7 @@
+/**
+ * Ahead of Time processing utilities for Spring Data Relational.
+ */
+@NonNullApi
+package org.springframework.data.relational.aot;
+
+import org.springframework.lang.NonNullApi;

spring-data-relational/src/main/resources/META-INF/spring/aot.factories

@@ -0,0 +1,2 @@
+org.springframework.beans.factory.aot.BeanRegistrationAotProcessor=\
+	org.springframework.data.relational.aot.RelationalManagedTypesBeanRegistrationAotProcessor