how can i switch the schema after login?

Your setting seams the classical situation for two different DataSources. Here is a Baeldung-Blog-Post how to configure Spring Data JPA.

First thing to notice, they are using @Primary. This is helping and standing in your way at the same time. You can only have ONE primary bean of a certain type. This is causing trouble for some people, since they try to “override” a spring bean by making their testing spring beans primary. Which results in having two primary beans with the same type. So be careful, when setting up your tests.

But it also eases things up, if you are mostly referring to one DataSource and only in a few cases to the other. This seams to be your case, so lets adopt it.

Your DataSource configuration could look like

@Configuration
public class DataSourceConfiguration {
    @Bean(name="loginDataSource")
    public DataSource loginDataSource(Environment env) {
        String url = env.getRequiredProperty("spring.logindatasource.url");
        return DataSourceBuilder.create()
            .driverClassName(env.getRequiredProperty("spring.logindatasource.driverClassName"))
            [...]
            .url(url)
            .build();
    }
    
    @Bean(name="companyDependentDataSource")
    @Primary // use with caution, I'd recommend to use name based autowiring. See @Qualifier
    public DataSource companyDependentDataSource(Environment env) {
        return new UserSchemaAwareRoutingDataSource(); // Autowiring is done afterwards by Spring
    }
}

These two DataSources can now be used in your repositories/DAOs or how ever you structure your program

@Autowired // This references the primary datasource, because no qualifier is given. UserSchemaAwareRoutingDataSource is its implementation
// @Qualifier("companyDependentDataSource") if @Primary is omitted
private DataSource companyDependentDataSource;

@Autowired
@Qualifier(name="loginDataSource") // reference by bean name
private DataSource loginDataSource

Here is an example how to configure Spring Data JPA with a DataSource referenced by name:

@Configuration
@EnableJpaRepositories(
    basePackages = "<your entity package>", 
    entityManagerFactoryRef = "companyEntityManagerFactory", 
    transactionManagerRef = "companyTransactionManager"
)
public class CompanyPersistenceConfiguration {

    @Autowired
    @Qualifier("companyDependentDataSource")
    private DataSource companyDependentDataSource;
    
    @Bean(name="companyEntityManagerFactory")
    public LocalContainerEntityManagerFactoryBean companyEntityManagerFactory() {
        LocalContainerEntityManagerFactoryBean emf = new LocalContainerEntityManagerFactoryBean();
        emf.setDataSource(companyDependentDataSource);
        // ... see Baeldung Blog Post
        return emf;
    }

    @Bean(name="companyTransactionManager")
    public PlatformTransactionManager companyTransactionManager() {
        JpaTransactionManager tm = new JpaTransactionManager();
        tm.setEntityManagerFactory(companyEntityManagerFactory().getObject());
        return tm;
    }
}

As described in my SO-answer you referred to there is an important assumption

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.

This is the trick which makes the UserSchemaAwareRoutingDataSource implementation possible. Spring beans are mostly singletons and therefore stateless. This also applies to the normal usage of DataSources. They are treated as stateless singletons and the references to them are passed over in the whole program. So we need to find a way to provide a single instance of the companyDependentDataSource which is behaving different on user basis regardless. To get that behavior I suggest to use a request-scoped bean.

In a web application, you can use @Scope(REQUEST_SCOPE) to create such objects. There is also a Bealdung Post talking about that topic. As usual, @Bean annotated methods reside in @Confiugration annotated classes.

@Configuration
public class UsuarioConfiguration {
    @Bean
    @Scope(value = WebApplicationContext.SCOPE_REQUEST,
     proxyMode = ScopedProxyMode.TARGET_CLASS) // or just @RequestScope
    public Usuario usario() {
        // based on your edit2 
        return (Usuario) SecurityContextHolder.getContext().getAuthentication().getPrincipal();
    }
}

Now you can use this request scoped object with a provider inside your singleton DataSource to behave different according to the logged in user:

@Autowired
private Usario usario; // this is now a request-scoped proxy which will create the corresponding bean (see UsuarioConfiguration.usario()

private DataSource determineTargetDataSource() {
    try {
        String db_schema = this.usuario.getTunnel().getDb_schema();
        return dataSources.get(db_schema);
    } catch (Exception ex) {
        ex.printStackTrace();
    }
    return null;
}

I hope this helps you understand the request scope concept of Spring.

So your login process would look something like

  1. User input username and password
  2. A normal spring bean, referencing the userDataSource by name, is checking the login and is putting the user information into the session/securitycontext/cookie/….
  3. When successful, during the next request the companyDependentDataSource is capable of retrieving a properly setup Usario object
  4. You can use this datasource now to do user specific stuff.

To verify your DataSource is properly working you could create a small Spring MVC endpoint

@RestController
public class DataSourceVerificationController {
    @Autowired
    private Usario usario;
    
    @Autowired
    @Qualifier("companyDependentDataSource") // omit this annotation if you use @Primary
    private DataSource companyDependentDataSource;

    @GetRequest("/test")
    public String test() throws Exception {
        String schema = usario.getTunnel().getDb_schema()
    
        Connection con = companyDependentDataSource.getConnection();
        Statement stmt = con.createStatement();
        ResultSet rs = stmt.executeQuery("select name from Employee"); // just a random guess
        rs.next();
        String name = rs.getString("name")
        rs.close();
        stmt.close();
        con.close();
        
        return "name = '" + name + "', schema = '" + schema + "'";
    }
}

Take your favorite browser go to your login page, do a valid login and call http://localhost:8080/test afterwards

CLICK HERE to find out more related problems solutions.

Leave a Comment

Your email address will not be published.

Scroll to Top