Configuration class to register and retrieve multiple data sources from Spring Boot application properties

MultiDataSourceConfig.java

package org.example;

import lombok.RequiredArgsConstructor;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.FactoryBean;
import org.springframework.beans.factory.NoSuchBeanDefinitionException;
import org.springframework.beans.factory.NoUniqueBeanDefinitionException;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.BeanDefinitionRegistryPostProcessor;
import org.springframework.boot.context.properties.bind.Bindable;
import org.springframework.boot.context.properties.bind.Binder;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.ResolvableType;
import org.springframework.core.env.Environment;

import javax.sql.DataSource;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * Configuration class for adding {@link DataSource} beans to Spring context from the configuration.
 * A {@link DataSourceProvider} bean is provided to obtain the data sources from the bean definitions
 * generated by this configuration.
 * <br/>
 * <strong>Properties</strong><br/>
 * The expected properties are the same as the ones defined using {@code spring.datasource}, as the property binding
 * mechanism is similar in both cases.  For example, the following properties will result in generation of two data
 * source beans named {@code foo} and {@code bar} respectively.
 * <br/>
 * The configuration prefix is defined in the public constant {@link MultiDataSourceConfig#CONFIG_PREFIX}.
 * Due to constrains in the way {@code BeanDefinitionRegistryPostProcessor} are instanced, this constant is static.
 * <pre>
 * config.prefix.foo.jdbc-url=jdbc:postgresql://localhost:5644/postgres
 * config.prefix.foo.username=postgres
 * config.prefix.foo.password=postgres
 * config.prefix.foo.dialect=org.hibernate.dialect.PostgreSQLDialect
 * config.prefix.foo.primary=true
 *
 * config.prefix.bar.jdbc-url=jdbc:postgresql://localhost:4546/postgres
 * config.prefix.bar.username=postgres
 * config.prefix.bar.password=postgres
 * config.prefix.bar.dialect=org.hibernate.dialect.PostgreSQLDialect
 * </pre>
 * <br/>
 * To obtain each data source, use {@link DataSourceProvider#getDataSource} and provide the data source name.
 */
@Configuration
@RequiredArgsConstructor
public class MultiDataSourceConfig {
    /**
     * The configuration prefix for data sources properties picked up by this configuration class.
     */
    public static final String CONFIG_PREFIX = "config.prefix";
    private static final String DATA_SOURCE_FACTORY_BEAN_NAME = "multiConfig_DataSourceFactoryBean";

    @Bean
    public static MultiDataSourceRegistrar multiDataSourceRegistrar(Environment environment) {
        return new MultiDataSourceRegistrar(environment);
    }

    /**
     * {@code BeanDefinitionRegistryPostProcessor} that registers new bean definitions for the data sources defined
     * in the application properties. These beans are defined as singletons and their lifecycle is managed by the IoC
     * container.
     *
     * @see BeanDefinitionRegistryPostProcessor
     */
    @RequiredArgsConstructor
    private static class MultiDataSourceRegistrar implements BeanDefinitionRegistryPostProcessor {
        private final Environment environment;

        @Override
        public void postProcessBeanDefinitionRegistry(BeanDefinitionRegistry registry) throws BeansException {
            Map<String, Map<String, Object>> properties = Binder.get(environment).bindOrCreate(CONFIG_PREFIX,
                    Bindable.of(ResolvableType.forClassWithGenerics(Map.class,
                            ResolvableType.forClass(String.class),
                            ResolvableType.forClassWithGenerics(Map.class, String.class, Object.class))));

            for (var entry : properties.entrySet()) {
                String name = entry.getKey();
                Map<String, Object> config = entry.getValue();
                BeanDefinitionBuilder builder = BeanDefinitionBuilder.genericBeanDefinition(DataSource.class,
                                () -> createBoundDataSource(name))
                        .setLazyInit(true);
                if (Boolean.parseBoolean(String.valueOf(config.get("primary")))) {
                    builder.setPrimary(true);
                }
                String beanDefName = buildBeanDefinitionName(name);
                registry.registerBeanDefinition(beanDefName, builder.getBeanDefinition());
            }
        }

        private DataSource createBoundDataSource(String name) {
            var properties = CONFIG_PREFIX + "." + name;
            return Binder.get(environment).bind(properties, Bindable.of(DataSource.class)
                            .withSuppliedValue(() -> DataSourceBuilder.create().build()))
                    .orElseThrow(() -> new BeansException("No data source was bound to properties: " + properties) {
                    });
        }

        @Override
        public void postProcessBeanFactory(ConfigurableListableBeanFactory beanFactory) throws BeansException {
            /* noop */
        }
    }

    private static String buildBeanDefinitionName(String base) {
        return MultiDataSourceRegistrar.class.getName() + "_" + base;
    }

    /**
     * Bean definition of a data source factory that returns a data source from the bean definitions in the
     * primary bean factory lazily.
     * <br/>
     * This bean is registered very early and prevents Spring Boot autoconfiguration from creating new data source bean.
     * Bean definitions registered using a {@code BeanDefinitionRegistryPostProcessor} are not taken into account
     * for {@link org.springframework.context.annotation.Condition} resolving.
     */
    @Bean(DATA_SOURCE_FACTORY_BEAN_NAME)
    public FactoryBean<DataSource> dataSourceFactoryBean(ConfigurableListableBeanFactory beanFactory) {
        return new FactoryBean<>() {
            @Override
            public DataSource getObject() {
                List<String> names = Arrays.stream(beanFactory.getBeanNamesForType(DataSource.class))
                        .filter(name -> !DATA_SOURCE_FACTORY_BEAN_NAME.equals(name)).toList();
                if (names.size() == 1) {
                    return beanFactory.getBean(names.get(0), DataSource.class);
                }

                List<String> primary = names.stream().filter(this::isPrimaryBean).toList();
                if (primary.size() == 1) {
                    return beanFactory.getBean(primary.get(0), DataSource.class);
                } else if (primary.size() > 1) {
                    throw new NoUniqueBeanDefinitionException(DataSource.class, primary);
                }

                throw new NoUniqueBeanDefinitionException(DataSource.class, names);
            }

            private boolean isPrimaryBean(String name) {
                BeanDefinition beanDef = beanFactory.getBeanDefinition(name);
                return beanDef.isPrimary();
            }

            @Override
            public Class<DataSource> getObjectType() {
                return DataSource.class;
            }
        };
    }

    @Bean
    public DataSourceProvider dataSourceProvider(BeanFactory beanFactory) {
        return new DefaultDataSourceProvider(beanFactory);
    }

    /**
     * A provider for the datasource beans registered using properties.
     */
    public interface DataSourceProvider {
        /**
         * Obtain a data source bean from the application context.
         *
         * @param name the name of the data source
         * @return the data source bean
         */
        DataSource getDataSource(String name);
    }

    @RequiredArgsConstructor
    private static class DefaultDataSourceProvider implements DataSourceProvider {
        private final BeanFactory beanFactory;

        @Override
        public DataSource getDataSource(String name) {
            String beanDefName = buildBeanDefinitionName(name);
            try {
                return beanFactory.getBean(beanDefName, DataSource.class);
            } catch (NoSuchBeanDefinitionException ex) {
                throw new BeansException("No data source with name \"" + name + "\". " +
                        "Verify the application properties contain configuration prefixed with \"" + CONFIG_PREFIX + "." + name + "\"", ex) {
                };
            } catch (BeansException ex) {
                throw new BeansException("Unable to provide data source with name \"" + name + "\"", ex) {
                };
            }
        }
    }
}