/** | |
* | |
* Licensed to the Apache Software Foundation (ASF) under one or more | |
* contributor license agreements. See the NOTICE file distributed with | |
* this work for additional information regarding copyright ownership. | |
* The ASF licenses this file to You under the Apache License, Version 2.0 | |
* (the "License"); you may not use this file except in compliance with | |
* the License. You may obtain a copy of the License at | |
* | |
* http://www.apache.org/licenses/LICENSE-2.0 | |
* | |
* Unless required by applicable law or agreed to in writing, software | |
* distributed under the License is distributed on an "AS IS" BASIS, | |
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
* See the License for the specific language governing permissions and | |
* limitations under the License. | |
*/ | |
package org.apache.openejb.resource.jdbc; | |
import static org.junit.Assert.assertEquals; | |
import java.sql.*; | |
import java.util.ArrayList; | |
import java.util.List; | |
import java.util.Map; | |
import java.util.Properties; | |
import java.util.concurrent.ConcurrentHashMap; | |
import javax.ejb.Local; | |
import javax.ejb.Stateless; | |
import javax.ejb.TransactionAttribute; | |
import javax.ejb.TransactionAttributeType; | |
import javax.naming.Context; | |
import javax.naming.InitialContext; | |
import javax.naming.NamingException; | |
import javax.persistence.Entity; | |
import javax.persistence.EntityManager; | |
import javax.persistence.Id; | |
import javax.persistence.PersistenceContext; | |
import javax.sql.DataSource; | |
import org.apache.openejb.assembler.classic.Assembler; | |
import org.apache.openejb.assembler.classic.ProxyFactoryInfo; | |
import org.apache.openejb.assembler.classic.ResourceInfo; | |
import org.apache.openejb.assembler.classic.SecurityServiceInfo; | |
import org.apache.openejb.assembler.classic.StatelessSessionContainerInfo; | |
import org.apache.openejb.assembler.classic.TransactionServiceInfo; | |
import org.apache.openejb.core.LocalInitialContextFactory; | |
import org.apache.openejb.config.AppModule; | |
import org.apache.openejb.config.ConfigurationFactory; | |
import org.apache.openejb.config.EjbModule; | |
import org.apache.openejb.config.PersistenceModule; | |
import org.apache.openejb.config.sys.Resource; | |
import org.apache.openejb.jee.EjbJar; | |
import org.apache.openejb.jee.StatelessBean; | |
import org.apache.openejb.jee.jpa.unit.Persistence; | |
import org.apache.openejb.jee.jpa.unit.PersistenceUnit; | |
import org.apache.openejb.jee.jpa.unit.TransactionType; | |
import org.apache.openejb.loader.SystemInstance; | |
import org.apache.openejb.spi.ContainerSystem; | |
import org.junit.Test; | |
public class DynamicDataSourceTest { | |
@Test | |
public void route() throws Exception { | |
System.setProperty(javax.naming.Context.INITIAL_CONTEXT_FACTORY, LocalInitialContextFactory.class.getName()); | |
ConfigurationFactory config = new ConfigurationFactory(); | |
Assembler assembler = new Assembler(); | |
assembler.createProxyFactory(config.configureService(ProxyFactoryInfo.class)); | |
assembler.createTransactionManager(config.configureService(TransactionServiceInfo.class)); | |
assembler.createSecurityService(config.configureService(SecurityServiceInfo.class)); | |
// resources | |
for (int i = 1; i <= 3; i++) { | |
String dbName = "database" + i; | |
Resource resourceDs = new Resource(dbName, "DataSource"); | |
Properties p = resourceDs.getProperties(); | |
p.put("JdbcDriver", "org.hsqldb.jdbcDriver"); | |
p.put("JdbcUrl", "jdbc:hsqldb:mem:db" + i); | |
p.put("UserName", "sa"); | |
p.put("Password", ""); | |
p.put("JtaManaged", "true"); | |
assembler.createResource(config.configureService(resourceDs, ResourceInfo.class)); | |
} | |
Resource resourceRouter = new Resource("My Router", "org.apache.openejb.router.test.DynamicDataSourceTest$DeterminedRouter", "org.router:DeterminedRouter"); | |
resourceRouter.getProperties().setProperty("DatasourceNames", "database1 database2 database3"); | |
resourceRouter.getProperties().setProperty("DefaultDataSourceName", "database1"); | |
assembler.createResource(config.configureService(resourceRouter, ResourceInfo.class)); | |
Resource resourceRoutedDs = new Resource("Routed Datasource", "org.apache.openejb.resource.jdbc.Router", "RoutedDataSource"); | |
resourceRoutedDs.getProperties().setProperty("Router", "My Router"); | |
assembler.createResource(config.configureService(resourceRoutedDs, ResourceInfo.class)); | |
// containers | |
StatelessSessionContainerInfo statelessContainerInfo = config.configureService(StatelessSessionContainerInfo.class); | |
assembler.createContainer(statelessContainerInfo); | |
EjbJar ejbJar = new EjbJar(); | |
ejbJar.addEnterpriseBean(new StatelessBean(RoutedEJBBean.class)); | |
ejbJar.addEnterpriseBean(new StatelessBean(UtilityBean.class)); | |
EjbModule ejbModule = new EjbModule(ejbJar); | |
// Create an "ear" | |
AppModule appModule = new AppModule(ejbModule.getClassLoader(), "test-dynamic-data-source"); | |
appModule.getEjbModules().add(ejbModule); | |
// Create a persistence-units | |
PersistenceUnit unit = new PersistenceUnit("router"); | |
unit.addClass(Person.class); | |
unit.getProperties().put("openjpa.jdbc.SynchronizeMappings", "buildSchema"); | |
unit.setTransactionType(TransactionType.JTA); | |
unit.setJtaDataSource("Routed Datasource"); | |
appModule.getPersistenceModules().add(new PersistenceModule("root", new Persistence(unit))); | |
for (int i = 1; i <= 3; i++) { | |
PersistenceUnit u = new PersistenceUnit("db" + i); | |
u.addClass(Person.class); | |
u.getProperties().put("openjpa.jdbc.SynchronizeMappings", "buildSchema"); | |
u.setTransactionType(TransactionType.JTA); | |
u.setJtaDataSource("database" + i); | |
appModule.getPersistenceModules().add(new PersistenceModule("root", new Persistence(u))); | |
} | |
assembler.createApplication(config.configureApplication(appModule)); | |
// context | |
Context ctx = new InitialContext(); | |
// running persist on all "routed" databases | |
final List<String> databases = new ArrayList<String>(); | |
databases.add("database1"); | |
databases.add("database2"); | |
databases.add("database3"); | |
// convinient bean to create tables for each persistence unit | |
Utility utility = (Utility) ctx.lookup("UtilityBeanLocal"); | |
utility.initDatabase(); | |
RoutedEJB ejb = (RoutedEJB) ctx.lookup("RoutedEJBBeanLocal"); | |
for (int i = 0; i < 18; i++) { | |
String name = "record " + i; | |
String db = databases.get(i % 3); | |
ejb.persist(i, name, db); | |
} | |
// assert database records number using jdbc | |
for (int i = 1; i <= 3; i++) { | |
Connection connection = DriverManager.getConnection("jdbc:hsqldb:mem:db" + i, "sa", ""); | |
Statement st = connection.createStatement(); | |
ResultSet rs = st.executeQuery("select count(*) from \"DynamicDataSourceTest$Person\""); | |
rs.next(); | |
assertEquals(6, rs.getInt(1)); | |
st.close(); | |
connection.close(); | |
} | |
} | |
@Stateless | |
@Local(RoutedEJB.class) | |
public static class RoutedEJBBean implements RoutedEJB { | |
@PersistenceContext(unitName = "router") | |
private EntityManager em; | |
@javax.annotation.Resource(name = "My Router", type = DeterminedRouter.class) | |
private DeterminedRouter router; | |
public void persist(int id, String name, String ds) { | |
router.setDataSource(ds); | |
em.persist(new Person(id, name)); | |
} | |
} | |
@Stateless | |
@Local(Utility.class) | |
public static class UtilityBean implements Utility { | |
@PersistenceContext(unitName = "db1") | |
private EntityManager em1; | |
@PersistenceContext(unitName = "db2") | |
private EntityManager em2; | |
@PersistenceContext(unitName = "db3") | |
private EntityManager em3; | |
@TransactionAttribute(TransactionAttributeType.SUPPORTS) | |
public void initDatabase() { | |
em1.find(Person.class, 0); | |
em2.find(Person.class, 0); | |
em3.find(Person.class, 0); | |
} | |
} | |
public static interface RoutedEJB { | |
void persist(int id, String name, String ds); | |
} | |
public static interface Utility { | |
void initDatabase(); | |
} | |
@Entity | |
public static class Person { | |
@Id | |
private long id; | |
private String name; | |
public Person() { | |
// no-op | |
} | |
public Person(int i, String n) { | |
id = i; | |
name = n; | |
} | |
public long getId() { | |
return id; | |
} | |
public void setId(long id) { | |
this.id = id; | |
} | |
public String getName() { | |
return name; | |
} | |
public void setName(String name) { | |
this.name = name; | |
} | |
} | |
public static class DeterminedRouter implements Router { | |
private String dataSourceNames; | |
private String defaultDataSourceName; | |
private Map<String, DataSource> dataSources = null; | |
private ThreadLocal<DataSource> currentDataSource = new ThreadLocal<DataSource>(); | |
/** | |
* @param datasourceList datasource resource name, separator is a space | |
*/ | |
public void setDataSourceNames(String datasourceList) { | |
dataSourceNames = datasourceList; | |
} | |
/** | |
* lookup datasource in openejb resources | |
*/ | |
private void init() { | |
dataSources = new ConcurrentHashMap<String, DataSource>(); | |
for (String ds : dataSourceNames.split(" ")) { | |
ContainerSystem containerSystem = SystemInstance.get().getComponent(ContainerSystem.class); | |
Object o = null; | |
Context ctx = containerSystem.getJNDIContext(); | |
try { | |
o = ctx.lookup("openejb:Resource/" + ds); | |
if (o instanceof DataSource) { | |
dataSources.put(ds, (DataSource) o); | |
} | |
} catch (NamingException e) { | |
} | |
} | |
} | |
/** | |
* @return the user selected data source if it is set | |
* or the default one | |
* @throws IllegalArgumentException if the data source is not found | |
*/ | |
public DataSource getDataSource() { | |
// lazy init of routed datasources | |
if (dataSources == null) { | |
init(); | |
} | |
// if no datasource is selected use the default one | |
if (currentDataSource.get() == null) { | |
if (dataSources.containsKey(defaultDataSourceName)) { | |
return dataSources.get(defaultDataSourceName); | |
} else { | |
throw new IllegalArgumentException("you have to specify at least one datasource"); | |
} | |
} | |
// the developper set the datasource to use | |
return currentDataSource.get(); | |
} | |
/** | |
* | |
* @param datasourceName data source name | |
*/ | |
public void setDataSource(String datasourceName) { | |
if (dataSources == null) { | |
init(); | |
} | |
if (!dataSources.containsKey(datasourceName)) { | |
throw new IllegalArgumentException("data source called " + datasourceName + " can't be found."); | |
} | |
DataSource ds = dataSources.get(datasourceName); | |
currentDataSource.set(ds); | |
} | |
/** | |
* reset the data source | |
*/ | |
public void clear() { | |
currentDataSource.remove(); | |
} | |
public void setDefaultDataSourceName(String name) { | |
this.defaultDataSourceName = name; | |
} | |
} | |
} |