java - Set schema name at run time using REST API input in spring boot based application

Question

In a postgres database, the schema name is set as the buyer id(so, if buyer id is buyer_2213 then schema name will be buyer_2213). These schemas have multiple tables and these tables have common structure for all schemas.

Now, I am writing REST API using spring boot to get buyer data from these schema, but since schema name is buyer id dependent, I have to write manual queries and could not use JPA features for that.

Is there any way to set schema name to entity using REST API request parameters. So, in below entity, can we set schema using buyerId passed at API call defined in BuyerController:

  @Entity
  @Table(name="buyer_table", schema="set_it_using_API_input")
  public class BuyerTable{
    ...
  }


  @RestController
  @RequestMapping("/buyer")
  public class BuyerController{
    
    @GetMapping("/{buyerId}")
    public void getBuyerData(@PathVariable(required = true) String buyerId){
      ...
    }
  }

Also, the buyer Id is not same as logged in user(consider this situation as admin user who is trying to get buyer's details) and will be provided as API request param only(or any other way as API input). Hence I could not find relevance to this

Answer 1

I've finally found a working solution for this. This solution uses configurations mostly from here but is more specific to the my question requirements.

The idea is obviously to use AbstractDataSource and data source configurations are pretty much similar to what shown in here with just the schema name will be set using a setter which can be called from inside API logic.
First, we need to write an implementation of AbstractDataSource which will pretty much look like this:

public class BuyerSchemaDataSource extends AbstractDataSource {

    private LoadingCache<String, DataSource> dataSources = createCache();

    public void setSchemaName(String schemaName){
        this.schemaName = schemaName;
    }

    @Override public Connection getConnection() throws SQLException {
        try {
            return determineTargetDataSource().getConnection();
        } catch (ExecutionException e) {
            //print exception
            return null;
        }
    }

    @Override public Connection getConnection(String username, String password)
        throws SQLException {
        try {
            return determineTargetDataSource().getConnection(username,password);
        } catch (ExecutionException e) {
            //print exception
            return null;
        }
    }

    private DataSource determineTargetDataSource() throws ExecutionException {
        if(!utils.isNullOrEmpty(schemaName)){
            return dataSources.get(schemaName);
        }
        return buildDataSourceFromSchema(null);
    }

    private LoadingCache<String, DataSource> createCache(){
        return CacheBuilder.newBuilder()
            .maximumSize(100)
            .expireAfterWrite(10, TimeUnit.MINUTES)
            .build(new CacheLoader<String, DataSource>() {
                @Override public DataSource load(String key) throws Exception {
                    return buildDataSourceFromSchema(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, just like any other data source this can be used in a spring configuration file like this:

@Configuration
@EnableTransactionManagement
@EnableJpaRepositories(entityManagerFactoryRef = "schemaSpecificEntityManagerFactory",
    transactionManagerRef = "schemaSpecificTransactionManager")
public class SchemaSpecificConfig {

    @Bean(name = "schemaSpecificDataSource")
    public DataSource schemaSpecificDataSource(){
        return new BuyerSchemaDataSource();
    }

    @Bean(name = "schemaSpecificEntityManagerFactory")
    public LocalContainerEntityManagerFactoryBean schemaSpecificEntityManagerFactory(
        EntityManagerFactoryBuilder builder, @Qualifier("schemaSpecificDataSource") DataSource dataSource) {

        HashMap<String, Object> properties = new HashMap<>();
        properties.put("hibernate.hbm2ddl.auto", "update");
        properties.put("hibernate.dialect", "org.hibernate.dialect.PostgreSQLDialect");
        return builder.dataSource(dataSource).properties(properties)
            .persistenceUnit("SchemaSpecific").build();
    }

    @Bean(name = "schemaSpecificTransactionManager")
    public PlatformTransactionManager schemaSpecificTransactionManager(
        @Qualifier("schemaSpecificEntityManagerFactory") EntityManagerFactory schemaSpecificEntityManagerFactory) {
        return new JpaTransactionManager(schemaSpecificEntityManagerFactory);
    }
}

Now, the setSchema() method defined in BuyerSchemaDataSource can be called from inside API logic in controller.

This looks like a bad way to solve it but I have not find anything better than this and all suggestions/edits are appreciated.