baao baao - 2 months ago 42
Java Question

Change database schema during runtime based on logged in user

I've read many questions and answers about dynamic datasource routing and have implemented a solution using

AbstractRoutingDataSource
and another(see below). That's fine, but requires hardcoded properties for all datasources. As the number of users using the application increases, this isn't a suitable way of routing any more. Also it would require to add an entry to the properties every time a new user registers. The situation is as follows


  • 1 database server

  • many schemas on that server, every user has their own schema.

  • I only need to change the schema name during runtime

  • schema name is retainable by logged in user



I'm using
spring boot 1.4.0
together with
hibernate 5.1
and
spring data jpa


I can't find a way to change the schema completely dynamically. Does someone know how to do it in spring?

EDIT:

Thanks to @Johannes Leimer's answer, I came a further, though it needs a high amount of optimization.

It currently is quite an ugly implementation, here's the code:

User Provider:

@Configuration
public class UserProvider implements Provider<CustomUserDetails> {
private static ThreadLocal<CustomUserDetails> CURRENT;

@Override @Bean
public CustomUserDetails get() {
try {
return CURRENT.get();
} catch (NullPointerException e ) {
e.printStackTrace();
return null;
}
}

public static void setCURRENT(ThreadLocal<CustomUserDetails> CURRENTs) {
System.out.println(CURRENTs);
CURRENT = CURRENTs;
}
}


UserSchemaAwareRoutingDatasource:

public class UserSchemaAwareRoutingDataSource extends AbstractDataSource {
@Autowired
UserProvider user;

@Inject
Environment env;
private LoadingCache<String, DataSource> dataSources = createCache();

@Override
public Connection getConnection() throws SQLException {
try {
return determineTargetDataSource().getConnection();
} catch (ExecutionException e){
e.printStackTrace();

return null;
}
}

// this method is never called.
@Override
public Connection getConnection(String username, String password) throws SQLException {
System.out.println("getConnection" + username);
System.out.println("getConnection2" + password);
try {
return determineTargetDataSource().getConnection(username, password);
} catch (ExecutionException e) {
e.printStackTrace();
return null;
}
}

private DataSource determineTargetDataSource() throws SQLException, ExecutionException {
try {
// if I use the autowired user, it throws a NPE
CustomUserDetails customUserDetails = new UserProvider().get();
String schema = customUserDetails.getUserDatabase();
return dataSources.get(schema);
} catch (NullPointerException e) {
e.printStackTrace();
// gives an error on startup, when no logged in user is found!
return dataSources.get("fooooo");
}

}

private LoadingCache<String, DataSource> createCache() {
return CacheBuilder.newBuilder()
.maximumSize(100)
.expireAfterWrite(10, TimeUnit.MINUTES)
.build(
new CacheLoader<String, DataSource>() {
public DataSource load(String key) throws Exception {
System.out.println("key : " + key);
return buildDataSourceForSchema(key);
}
});
}

private DataSource buildDataSourceForSchema(String schema) {
String url = env.getRequiredProperty("spring.datasource.url") + schema;
return DataSourceBuilder.create()
.driverClassName(env.getRequiredProperty("spring.datasource.driver-class-name"))
.username(env.getRequiredProperty("spring.datasource.username"))
.password(env.getRequiredProperty("spring.datasource.password"))
.url(url)
.build();
}
}


With this code, when I set the current user in every request (Controller method) like this:

CustomUserDetails customUserDetails = getCustomUserDetails();
System.out.println(customUserDetails);

UserProvider.setCURRENT(new ThreadLocal<CustomUserDetails>(){{set(customUserDetails);}});


(where getCustomUserDetails() is):

CustomUserDetails getCustomUserDetails() {
return (CustomUserDetails) SecurityContextHolder.getContext().getAuthentication().getPrincipal();
}


I get the correct result from the correct schema, even for concurrent requests. The set() of the current user could obviously be outsourced in a
RequestInterceptor. However, I didn't manage to implement the suggested Provider, and as commented in the code ->

// if I use the autowired user, it throws a NPE
CustomUserDetails customUserDetails = new UserProvider().get();


the autowiring fails (understandably).

Can you suggest how to make it work properly?

Answer

Assumptions

Because I don't have the reputation yet to post a comment below your question, my answer is based on the following assumtions:

  • The current schema name to be used for the current user is accessible through a Spring JSR-330 Provider like private javax.inject.Provider<User> user; String schema = user.get().getSchema();. This is ideally a ThreadLocal-based proxy.

  • To build a DataSource which is fully configured in a way you need it requires the same properties. Every time. The only thing which is different is the schema name. (It would easily possible to obtain other different parameters as well, but this would be too much for this answer)

  • Each schema is already set up with the needed DDL, so there is no need for hibernate to create tables or something else

  • Each database schema looks completely the same except for its name

  • You need to reuse a DataSource every time the corresponding user makes a request to your application. But you don't want to have every DataSource of every user permanently in the memory.

My solution idea

Use a combination of ThreadLocal proxys to get the schema name and a Singleton-DataSource which behaves different on every user request. This solution is inspired by your hint to AbstractRoutingDataSource, Meherzad's comments and own experience.

A dynamic DataSource

I suggest to facilitate the AbstractDataSource of Spring and implement it like the AbstractRoutingDataSource. Instead of a static Map-like approach we use a Guava Cache to get an easy to use cache.

public class UserSchemaAwareRoutingDataSource extends AbstractDataSource {
    private @Inject javax.inject.Provider<User> user;
    private @Inject Environment env;
    private LoadingCache<String, DataSource> dataSources = createCache();

    @Override
    public Connection getConnection() throws SQLException {
        return determineTargetDataSource().getConnection();
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        return determineTargetDataSource().getConnection(username, password);
    }

    private DataSource determineTargetDataSource() {
        String schema = user.get().getSchema();
        return dataSources.get(schema);
    }

    private LoadingCache<String, DataSource> createCache() {
        return CacheBuilder.newBuilder()
           .maximumSize(100)
           .expireAfterWrite(10, TimeUnit.MINUTES)
           .build(
               new CacheLoader<String, DataSource>() {
                 public DataSource load(String key) throws AnyException {
                   return buildDataSourceForSchema(key);
                 }
               });
    }

    private DataSource buildDataSourceForSchema(String schema) {
        // e.g. of property: "jdbc:postgresql://localhost:5432/mydatabase?currentSchema="
        String url = env.getRequiredProperty("spring.datasource.url") + schema;
        return DataSourceBuilder.create()
            .driverClassName(env.getRequiredProperty("spring.datasource.driverClassName"))
            [...]
            .url(url)
            .build();
    }
}

Now you have a `DataSource┬┤ which acts different for every user. Once a DataSource is created it's gonna be cached for 10 minutes. That's it.

Make the application aware of our dynamic DataSource

The place to integrate our newly created DataSource is the DataSource singleton known to the spring context and used in all beans e.g. the EntityManagerFactory

So we need an equivalent to this:

@Primary
@Bean(name = "dataSource")
@ConfigurationProperties(prefix="spring.datasource")
public DataSource dataSource() {
    return DataSourceBuilder.create().build();
}

but it has to be more dynamic, than a plain property based DataSourceBuilder:

@Primary
@Bean(name = "dataSource")
public UserSchemaAwareRoutingDataSource dataSource() {
    return new UserSchemaAwareRoutingDataSource();
}

Conclusion

We have a transparent dynamic DataSource which uses the correct DataSource everytime.

Open questions

  • What to do, when no user is logged in? Is there no database access allowed?
  • Who sets up the schemes?

Disclaimer

I haven't tested this code!

EDIT: To implement a Provider<CustomUserDetails> with Spring you need to define this as prototype. You can utilize Springs support of JSR-330 and Spring Securitys SecurityContextHolder:

@Bean @Scope("prototype")
public CustomUserDetails customUserDetails() {
    return return (CustomUserDetails) SecurityContextHolder.getContext().getAuthentication().getPrincipal();
}

You don't need a RequestInterceptor, the UserProvider or the controller code to update the user anymore.

Does this help?