Multitenancy: Switching Databases at runtime with Spring

von

Spring Multitenancy: Switching datasources at runtime

It is very common to have a database layer underneath your web application. However, it is not so common to have multiple clones of that database and switch through them based on, for example, a header value of a web request.

Still, these use cases exist. If you google for “switch databases in Spring” (or Hibernate) it’ll not yield many results. The reason is the term you should be looking for is “multitenancy”. Fortunately, Hibernate supports multitenancy out of the box and so does Spring.

Configuring the databases

The first thing to define: what kind of databases do we have?

A YAML file helps (you can also use other options like properties):

db:
  configurations:
    de:
      url: jdbc:mysql://localhost/db1
      username: db1
      driver: com.mysql.cj.jdbc.Driver
      password: db1
    ch:
      url: jdbc:mysql://localhost/db2
      username: db2
      driver: com.mysql.cj.jdbc.Driver
      password: db2
    at:
      url: jdbc:mysql://localhost/db3
      username: db3
      driver: com.mysql.cj.jdbc.Driver
      password: db3

The idea here is to have a database mapped to each country code.

Of course, it would be pretty nice to have that configuration available as a Java object. Property mapping is pretty decent in Spring. All I have to do is create two classes.

Here is a Java class representing one database configuration:

public class DatabaseConfiguration {
    private String url;
    private String username;
    private String driver;
    private String password;
    // get/set ommitted

    public DataSource createDataSource() {
        DriverManagerDataSource dataSource = new DriverManagerDataSource();
        dataSource.setDriverClassName(driver);
        dataSource.setUrl(url);
        dataSource.setUsername(username);
        dataSource.setPassword(password);
        return dataSource;
    }
}

This entity can be “transformed” in to a data source with “createDataSource()”.

Here is the container holding all configurations:

@ConfigurationProperties(prefix = "db")
public class DatabaseConfigurations {
    private Map<String, DatabaseConfiguration> configurations = new HashMap<>();

    // get/set ommitted

    public Map<Object, Object> createTargetDataSources() {
        Map<Object, Object> result = new HashMap<>();
        configurations.forEach((key, value) ->  result.put(key, value.createDataSource()));
        return result;
    }
}

There are two things which might jump at you.

Firstly, the use of @ConfigurationProperties. This basically tells Spring to map all configuration data found under the key “db” into an object of this class. It would also create nested objects. The object member name “configurations” is also the same as in the YAML file.

Secondly, I created a tiny method to create datasources and put them into a Map. The data sources are created using the factory method above.

Now I can create a Map full of country keys with their corresponding datasources.

Creating an AbstractRoutingDataSource

Here comes the magic. Usually, we’d create some kind of DataSource by either configuration or by adding a @Bean factory method to the applications configuration.

But this time we are going to do something else. Instead of the usual DataSource that we extend from AbstractRoutingDataSource, which is - as the name already suggests - a datasource which can route based on some idea you have in mind.

public class CountryRoutingDataSource extends AbstractRoutingDataSource {
    @Override
    protected Object determineCurrentLookupKey() {
        ServletRequestAttributes attr =
            (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        String pathInfo = attr.getRequest().getRequestURI();
        return pathInfo.substring(1, 3);
    }
}

In this class only determineCurrentLookupKey() is implemented. It is the place where we want to figure out the key from the map we have created earlier. In the case above, I am going to try and get it from the current servlet request. Which means, I expect an URL format of:

http://localhost:8080/de/cars

In this case, “de” would be selected and returned, which would (hopefully) end up in the German database being selected.

This method will be called after every call to “getConnection()”.

This is some minimal configuration. Unfortunately, it may not be enough. It requires you to have a servlet request under all circumstances. This may not always be the case. Think about JUnit tests or when the Spring Container starts. Several connections may be requested when beans are created.

The trick is, to think about defaults or alternative strategies to prevent NullPointers.

For example, you could add something like this to your selection strategy above:

if (attr == null) {
    // when the bean was created, attr is null (startup time)
    return "de";
}

// when running without servlet, this is empty (like in CLI)
if ("".equals(attr.getRequest().getRequestURI())) {
    return "de";
}

You may even come up with something better :)

Application class

Time to connect all this to our application boot class.


@SpringBootApplication
@EnableConfigurationProperties(DatabaseConfigurations.class)
public class ApiApplication implements WebMvcConfigurer {

    @Autowired
    DatabaseConfigurations databaseConfigurations;

    @Bean
    public DataSource dataSource() {
        CustomRoutingDataSource dataSource = new CustomRoutingDataSource();
        dataSource.setTargetDataSources(databaseConfigurations.createTargetDataSources());
        return dataSource;
    }

    public static void main(String[] args) {
        SpringApplication.run(ApiApplication.class, args);
    }
}

No more magic here. This is our custom factory method returning a DataSource, and as indicated before, it returns a CustomRoutingDataSource now. Please don’t forget to add your target datasources, which is the key/value map containing your datasources.

How to do migrations using Flyway?

If you are using Flyway to update your database schemas, we will need to add a few more lines of code.

The first step would be to disable Flyway’s autoconfiguration since we are leaving the “usual path” for 99% of the cases.

Then the application class would need to be modified as:

@SpringBootApplication(
    exclude = {
        FlywayAutoConfiguration.class
    }
)
@EnableConfigurationProperties(DatabaseConfigurations.class)
public class ApiApplication implements WebMvcConfigurer {
   ...
}

Now we have to think of a way to add migrations in some kind of loop at container start time. One way to do this would be to listen to an ApplicationEvent, specifically to the ContextRefreshedEvent which is happening when the context was built. All beans should be ready to use by then.

@Component
public class DatabaseMigration implements ApplicationListener<ContextRefreshedEvent> {
    // use constructor wiring
    @Autowired
    private DatabaseConfigurations databaseConfigurations;

    @Override
    public void onApplicationEvent(ContextRefreshedEvent contextRefreshedEvent) {
        databaseConfigurations.getConfigurations().forEach((key, value) -> {
            Flyway flyway = new Flyway();
            flyway.setDataSource(value.createDataSource());
            flyway.migrate();
        });
    }
}

Whenever the context is read, the onApplicationEvent method is called. We are simply looping through all database configurations and returning their datasources. Using a fresh Flyway object for every database, we can simply set the datasource and call migrate.

Please note that reusing the Flyway object by making it a Spring bean does not work. It caches old behavior and may lead to problems.

Also, remember that this approach might lead to problems when you have multiple servers deployed at once. All of them might run the update scripts at the same time.

With Amazon AWS you can ask for a rolling update (one server at a time). The other option is to have additional code to prevent conflicting parallel updates. In general, once you run more than a few servers you usually use a different approach to roll out your databases changes anyway.

Conclusion

Calling URLs with country codes should result in the selection of the correct database. It should also work with MockMVC and Junit tests in general.

What we haven’t implemented yet is a database connection pool. Depending on your needs, you can either create a pool for any connection or even extend the code in a way it only creates a pool for the most important databases, but sticks to the traditional open/close idea for the rest.

Image Credits

Tags: #Java #Spring #Database #JPA #Login