How to provide a custom DataSource

[omarmalave]

Previous to version 2.0.0, I extended the class DefaultFlywayDataSourceContext and the reload method to provide a custom DataSource; how can I do the same thing in the new release? Thanks.

[tomix26]

Can you please describe in more detail the reasons why you need the custom data source? Because the more I know, the better I can advise you.

[omarmalave]

The embedded Datasource runs as user 'Postgres'. For other applications, this would work fine, but as we are using Postgres row-level security, we need to be in control of the user we use to connect.

Reason: the object owner always skips row-level security; only the application user (created by the non-flyway setup scripts) will be subject to row-level security.

[tomix26]

Ok, just to clarify. At the beginning of the test suite you execute some sql logic to create a new user with some specific permissions. And then each time the DefaultFlywayDataSourceContext context is reloaded, you create a copy of the embedded data source and change the user. Do I understand it correctly?

[omarmalave]

Correct

[tomix26]

Ok, so in that case, I would recommend extending a database provider instead of the database context. The principle should be very similar to your previous solution, see the example below. The only downside is that if you were using multiple types of databases/providers, you would have to extend all the associated providers. But I guess that shouldn't be your case. In the future I will try to find a different and simpler way to manage multiple users.

@JdbcTest
@AutoConfigureEmbeddedDatabase
public class SpringBootMultipleUsersIntegrationTest {

    // the super class may differ according to the provider you are currently using - docker, zonky, ...
    public static class CustomZonkyPostgresDatabaseProvider extends ZonkyPostgresDatabaseProvider {

        public CustomZonkyPostgresDatabaseProvider(Environment environment, ObjectProvider<List<Consumer<EmbeddedPostgres.Builder>>> databaseCustomizers) {
            super(environment, databaseCustomizers);
        }

        @Override
        public EmbeddedDatabase createDatabase(DatabaseRequest request) throws ProviderException {
            EmbeddedDatabase database = super.createDatabase(request);

            JdbcTemplate jdbcTemplate = new JdbcTemplate(database);
            // some initialization logic - users are shared resources, that's the reason for handling the exception below
            jdbcTemplate.execute("DO $$\n" +
                    "    BEGIN\n" +
                    "        CREATE USER customuser WITH PASSWORD 'custompass';\n" +
                    "        EXCEPTION WHEN duplicate_object THEN RAISE NOTICE '%, skipping', SQLERRM USING ERRCODE = SQLSTATE;\n" +
                    "    END\n" +
                    "$$;");
            jdbcTemplate.execute("GRANT USAGE ON SCHEMA test TO customuser");
            jdbcTemplate.execute("GRANT SELECT ON ALL TABLES IN SCHEMA test TO customuser");

            try {
                PGSimpleDataSource dataSource = database.unwrap(PGSimpleDataSource.class);
                dataSource.setUser("customuser");
                dataSource.setPassword("custompass");
            } catch (SQLException e) {
                throw new ProviderException("Unexpected error when creating a database", e);
            }

            return database;
        }
    }

    @Configuration
    static class Config {

        @Bean
        @Provider(type = "zonky", database = "postgres") // the provider type may differ according to the provider you are overriding - check EmbeddedDatabaseAutoConfiguration for more details
        public DatabaseProvider zonkyPostgresDatabaseProvider(DatabaseProviderFactory postgresDatabaseProviderFactory) {
            return postgresDatabaseProviderFactory.createProvider(CustomZonkyPostgresDatabaseProvider.class);
        }
    }

    @Autowired
    private DataSource dataSource;

    @Test
    @FlywayTest
    public void testJdbcTemplate() {
        JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource);

        List<Map<String, Object>> persons = jdbcTemplate.queryForList("select * from test.person");
        assertThat(persons).isNotNull().hasSize(1);
    }
}

[tomix26]

You can move the inner classes somewhere outside of the test class, of coarse. This is just an example.

[omarmalave]

That worked @tomix26. Thanks a lot!

I have another question; how can I provide a custom preparer to be used by the parent class createDatabase method?

[tomix26]

@omarmalave I'm really sorry, but I don't know what you mean. Could you be more specific, please?

[omarmalave]

Sure, the thing is that in my previous implementation, I was also running flyway migrations before returning the final DataSource; it was easy because the reload method provides the flyway bean. Now, If I extend a Provider class, such as in your example, I won't have access to that bean, and I can't inject it because of the circular dependency it creates. So as far as I understand, the ZonkyCustomProvider.createDatabase runs the preparers that come in the request param, so I thought that maybe I could provide a custom preparer, perhaps extending FlywayDatabasePreparer, and do the migrations there

[omarmalave]

@tomix26 Forget that last comment; I solved it using a different approach. Thanks again.

[tomix26]

All flyway migrations should already be executed/applied before returning the final data source. The preparers that come in the request param are created with respect to the current setup of the Flyway bean. So if you want to change the behavior of the preparers, just change the configuration of the corresponding Flyway bean. There should be no reason to create a custom flyway preparer, or I don't know about any at least.

However, if you really think you need it, you can use the AopProxyUtils.getDatabaseContext method to get access to an underlying DatabaseContext API and via this context add or replace any of the existing preparers. But keep in mind that this is an internal API, so it may change over time, which may lead to slightly more complicated upgrades in the future.

[omarmalave]

Yeah, didn't need it at the end.