-
Notifications
You must be signed in to change notification settings - Fork 32
/
MultiTenantManager.java
119 lines (98 loc) · 3.86 KB
/
MultiTenantManager.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
package io.github.cepr0.demo.multitenant;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.datasource.DriverManagerDataSource;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Collection;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import static java.lang.String.format;
@Slf4j
@Configuration
public class MultiTenantManager {
private final ThreadLocal<String> currentTenant = new ThreadLocal<>();
private final Map<Object, Object> tenantDataSources = new ConcurrentHashMap<>();
private final DataSourceProperties properties;
private Function<String, DataSourceProperties> tenantResolver;
private AbstractRoutingDataSource multiTenantDataSource;
public MultiTenantManager(DataSourceProperties properties) {
this.properties = properties;
}
@Bean
public DataSource dataSource() {
multiTenantDataSource = new AbstractRoutingDataSource() {
@Override
protected Object determineCurrentLookupKey() {
return currentTenant.get();
}
};
multiTenantDataSource.setTargetDataSources(tenantDataSources);
multiTenantDataSource.setDefaultTargetDataSource(defaultDataSource());
multiTenantDataSource.afterPropertiesSet();
return multiTenantDataSource;
}
public void setTenantResolver(Function<String, DataSourceProperties> tenantResolver) {
this.tenantResolver = tenantResolver;
}
public void setCurrentTenant(String tenantId) throws SQLException, TenantNotFoundException, TenantResolvingException {
if (tenantIsAbsent(tenantId)) {
if (tenantResolver != null) {
DataSourceProperties properties;
try {
properties = tenantResolver.apply(tenantId);
log.debug("[d] Datasource properties resolved for tenant ID '{}'", tenantId);
} catch (Exception e) {
throw new TenantResolvingException(e, "Could not resolve the tenant!");
}
String url = properties.getUrl();
String username = properties.getUsername();
String password = properties.getPassword();
addTenant(tenantId, url, username, password);
} else {
throw new TenantNotFoundException(format("Tenant %s not found!", tenantId));
}
}
currentTenant.set(tenantId);
log.debug("[d] Tenant '{}' set as current.", tenantId);
}
public void addTenant(String tenantId, String url, String username, String password) throws SQLException {
DataSource dataSource = DataSourceBuilder.create()
.driverClassName(properties.getDriverClassName())
.url(url)
.username(username)
.password(password)
.build();
// Check that new connection is 'live'. If not - throw exception
try(Connection c = dataSource.getConnection()) {
tenantDataSources.put(tenantId, dataSource);
multiTenantDataSource.afterPropertiesSet();
log.debug("[d] Tenant '{}' added.", tenantId);
}
}
public DataSource removeTenant(String tenantId) {
Object removedDataSource = tenantDataSources.remove(tenantId);
multiTenantDataSource.afterPropertiesSet();
return (DataSource) removedDataSource;
}
public boolean tenantIsAbsent(String tenantId) {
return !tenantDataSources.containsKey(tenantId);
}
public Collection<Object> getTenantList() {
return tenantDataSources.keySet();
}
private DriverManagerDataSource defaultDataSource() {
DriverManagerDataSource defaultDataSource = new DriverManagerDataSource();
defaultDataSource.setDriverClassName("org.h2.Driver");
defaultDataSource.setUrl("jdbc:h2:mem:default");
defaultDataSource.setUsername("default");
defaultDataSource.setPassword("default");
return defaultDataSource;
}
}