How to use user defined database proxy in @DataJpaTest

Question

We need to track database metrics so we are using datasource-proxy to track this to integrate the same in spring boot project we have created custom datasource as below

@Component
@Slf4j
@ConfigurationProperties(prefix = "spring.datasource")
public class DataSourceBeanConfig
{

    public DataSource actualDataSource()
    {
        EmbeddedDatabaseBuilder databaseBuilder = new EmbeddedDatabaseBuilder();
        return databaseBuilder.setType(EmbeddedDatabaseType.H2).build();
    }

    @Bean
    @Primary
    public DataSource dataSource() {
        // use pretty formatted query with multiline enabled
        PrettyQueryEntryCreator creator = new PrettyQueryEntryCreator();
        creator.setMultiline(true);

        log.info("Inside Proxy Creation");

        SystemOutQueryLoggingListener listener = new SystemOutQueryLoggingListener();
        listener.setQueryLogEntryCreator(creator);

        return ProxyDataSourceBuilder
                .create(actualDataSource())
                .countQuery()
                .name("MyDS")
                .listener(listener)
                .build();
    }
}

When we run main application datasource-proxy is picked up but when we use @DataJpaTest it is not picking up. How to enable datasource-proxy in JUNIT test cases?

Edit::

Using Spring BeanPostProcessor to configure Proxy DataSource

@Slf4j
@Configuration
public class DataSourceBeanConfig implements BeanPostProcessor
{
    public Object postProcessAfterInitialization(Object bean, String beanName)
            throws BeansException
    {

        if (bean instanceof DataSource)
        {
            System.out.println("AfterInitialization : " + beanName);

            // use pretty formatted query with multiline enabled
            PrettyQueryEntryCreator creator = new PrettyQueryEntryCreator();
            creator.setMultiline(true);

            log.info("Inside Proxy Creation");

            SystemOutQueryLoggingListener listener = new SystemOutQueryLoggingListener();
            listener.setQueryLogEntryCreator(creator);

            return ProxyDataSourceBuilder.create((DataSource) bean).countQuery()
                    .name("MyDS").listener(listener).build();

        }
        return bean; // you can return any other object as well
    }
}

For starters it should be a `@Configration` and not a `@Component` and `@ConfigurationProperties` doesn't really do much here. However instead of wrapping it like this, I suggest to use a `BeanPostProcessor` which wraps and replaces the `DataSource`. — M. Deinum, May 29 '17 at 08:47
Even after wrapping main program is using proxy but @DataJpaTest is not using — rajadilipkolli, May 29 '17 at 09:29
After wrapping what? Have you applied my comment? I don't understand what you are trying to explain with your comment. — M. Deinum, May 29 '17 at 09:55
I have updated the question which implements your suggestion from comments, Even now my JUNIT test is not using proxy — rajadilipkolli, May 29 '17 at 10:04
I gave 2 suggestion and you cobbled them together into 1. It should be a regular class and registered as a `static` `@Bean` in one of your configuration files. I'm still not sure if that would detect it as `@DataJpaTest` will do an educated guess on what to bootstrap. — M. Deinum, May 29 '17 at 13:47

rajadilipkolli · Accepted Answer · 2023-07-09T00:51:59.923

Here is the solution we need to create TestConfiguration to use in @DataJpaTest

@RunWith(SpringRunner.class)
@DataJpaTest
public class DataTestJPA
{

    @TestConfiguration
    static class ProxyDataSourceConfig implements BeanPostProcessor
    {
        public Object postProcessAfterInitialization(Object bean, String beanName)
                throws BeansException
        {

            if (bean instanceof DataSource source && !(bean instanceof ProxyDataSource)) {
            {
               return ProxyDataSourceBuilder
                            .create(source)
                            .countQuery()
                            .name("MyDS")
                            .build();
                // @formatter:on

            }
            return bean; // you can return any other object as well
        }
    }
}

How to use user defined database proxy in @DataJpaTest

1 Answers1