org.springframework.jdbc.datasource.LazyConnectionDataSourceProxy Java Examples

public ReplicationRoutingDataSource(DataSource master, List<DataSource> slaves) {

		Assert.notNull(master, "master datasource can not be null.");

		Map<Object, Object> targetDataSources = new HashMap<>();
		targetDataSources.put(MASTER_KEY, new LazyConnectionDataSourceProxy(master));
		if (!CollectionUtils.isEmpty(slaves)) {
			slaveSize = slaves.size();
			for (int i = 0; i < slaveSize; i++) {
				targetDataSources.put(SLAVE_PREFIX + i,
						new LazyConnectionDataSourceProxy(slaves.get(i)));
			}
		}
		else {
			this.onlyMaster = true;
		}
		setTargetDataSources(targetDataSources);
		setDefaultTargetDataSource(targetDataSources.get(MASTER_KEY));
	}
 

@Override
public void afterPropertiesSet() throws Exception {
	if (this.shardsDataSources == null) {
		throw new IllegalArgumentException("Property 'shardsDataSources' is required");
	}
	this.resolvedShardsDataSources = new HashMap<Object, DataSource>(this.shardsDataSources.size());
	for (Map.Entry<Object, Object> entry : this.shardsDataSources.entrySet()) {
		DataSource dataSource = resolveSpecifiedDataSource(entry.getValue());
		LazyConnectionDataSourceProxy lazyDataSourceProxy = new LazyConnectionDataSourceProxy();
		lazyDataSourceProxy.setTargetDataSource(dataSource);
		this.resolvedShardsDataSources.put(entry.getKey(), lazyDataSourceProxy);
	}
	if (this.defaultDataSource == null) {
		throw new IllegalArgumentException("Property 'defaultDataSource' is required");
	}
	if(this.defaultDataSource != null){
		resolvedDefaultDataSource = this.resolveSpecifiedDataSource(defaultDataSource);
	}
	
}
 

@Override
protected void destroyInstance(DataSource instance) throws Exception {
	if (instance instanceof LazyConnectionDataSourceProxy) {
		DataSource targetDataSource = ((LazyConnectionDataSourceProxy) instance)
				.getTargetDataSource();
		if (targetDataSource instanceof ReadOnlyRoutingDataSource) {
			List<Object> dataSources = ((ReadOnlyRoutingDataSource) targetDataSource)
					.getDataSources();
			for (Object candidate : dataSources) {
				if (candidate instanceof DataSource) {
					super.destroyInstance(instance);
				}
			}
		}
	}
}
 

@Test
void getConnection_NoReadReplicaAvailableNoTransactionActive_returnsDefaultDataSource()
		throws Exception {

	// Arrange
	DataSource defaultDataSource = mock(DataSource.class);
	Connection connection = mock(Connection.class);

	when(defaultDataSource.getConnection()).thenReturn(connection);

	ReadOnlyRoutingDataSource readOnlyRoutingDataSource = new ReadOnlyRoutingDataSource();
	readOnlyRoutingDataSource.setTargetDataSources(Collections.emptyMap());
	readOnlyRoutingDataSource.setDefaultTargetDataSource(defaultDataSource);
	readOnlyRoutingDataSource.afterPropertiesSet();

	LazyConnectionDataSourceProxy dataSource = new LazyConnectionDataSourceProxy(
			readOnlyRoutingDataSource);

	// Act
	Connection connectionReturned = dataSource.getConnection();

	// Assert
	assertThat(((ConnectionProxy) connectionReturned).getTargetConnection())
			.isSameAs(connection);
}
 

/**
 * Constructs a
 * {@link org.springframework.cloud.aws.jdbc.datasource.ReadOnlyRoutingDataSource}
 * data source that contains the regular data source as a default, and all
 * read-replicas as additional data source. The
 * {@link org.springframework.cloud.aws.jdbc.datasource.ReadOnlyRoutingDataSource} is
 * additionally wrapped with a
 * {@link org.springframework.jdbc.datasource.LazyConnectionDataSourceProxy}, because
 * the read-only flag is only available after the transactional context has been
 * established. This is only the case if the physical connection is requested after
 * the transaction start and not while starting a transaction.
 * @return a ReadOnlyRoutingDataSource that is wrapped with a
 * LazyConnectionDataSourceProxy
 * @throws Exception if the underlying data source setup throws any exception
 */
@Override
protected DataSource createInstance() throws Exception {
	DBInstance dbInstance = getDbInstance(getDbInstanceIdentifier());

	// If there is no read replica available, delegate to super class
	if (dbInstance.getReadReplicaDBInstanceIdentifiers().isEmpty()) {
		return super.createInstance();
	}

	HashMap<Object, Object> replicaMap = new HashMap<>(
			dbInstance.getReadReplicaDBInstanceIdentifiers().size());

	for (String replicaName : dbInstance.getReadReplicaDBInstanceIdentifiers()) {
		replicaMap.put(replicaName, createDataSourceInstance(replicaName));
	}

	// Create the data source
	ReadOnlyRoutingDataSource dataSource = new ReadOnlyRoutingDataSource();
	dataSource.setTargetDataSources(replicaMap);
	dataSource.setDefaultTargetDataSource(
			createDataSourceInstance(getDbInstanceIdentifier()));

	// Initialize the class
	dataSource.afterPropertiesSet();

	return new LazyConnectionDataSourceProxy(dataSource);
}
 

@Test
void getConnection_NoReadReplicaAvailableReadOnlyTransactionActive_returnsDefaultDataSource()
		throws Exception {

	// Arrange
	DataSource defaultDataSource = mock(DataSource.class);
	Connection connection = mock(Connection.class);

	when(defaultDataSource.getConnection()).thenReturn(connection);

	ReadOnlyRoutingDataSource readOnlyRoutingDataSource = new ReadOnlyRoutingDataSource();
	readOnlyRoutingDataSource.setTargetDataSources(Collections.emptyMap());
	readOnlyRoutingDataSource.setDefaultTargetDataSource(defaultDataSource);
	readOnlyRoutingDataSource.afterPropertiesSet();

	LazyConnectionDataSourceProxy dataSource = new LazyConnectionDataSourceProxy(
			readOnlyRoutingDataSource);

	DefaultTransactionDefinition transactionDefinition = new DefaultTransactionDefinition();
	transactionDefinition.setReadOnly(true);

	TransactionTemplate transactionTemplate = new TransactionTemplate(
			new DataSourceTransactionManager(dataSource), transactionDefinition);

	// Act
	Connection connectionReturned = transactionTemplate.execute(status -> {
		try {
			return ((ConnectionProxy) dataSource.getConnection())
					.getTargetConnection();
		}
		catch (SQLException e) {
			fail(e.getMessage());
		}
		return null;
	});

	// Assert
	assertThat(connectionReturned).isSameAs(connection);
}
 

public void afterPropertiesSet() throws Exception {
    if (getHaDataSourceCreator() == null) {
        setHaDataSourceCreator(new NonHADataSourceCreator());
    }
    if (CollectionUtils.isEmpty(dataSourceDescriptors)) {
        return;
    }

    for (CobarDataSourceDescriptor descriptor : getDataSourceDescriptors()) {
        Validate.notEmpty(descriptor.getIdentity());
        Validate.notNull(descriptor.getTargetDataSource());

        DataSource dataSourceToUse = descriptor.getTargetDataSource();

        if (descriptor.getStandbyDataSource() != null) {
            dataSourceToUse = getHaDataSourceCreator().createHADataSource(descriptor);
            if (CollectionUtils.isNotEmpty(dataSourcePostProcessor)) {
                for (IDataSourcePostProcessor pp : dataSourcePostProcessor) {
                    dataSourceToUse = pp.postProcess(dataSourceToUse);
                }
            }
        }

        dataSources.put(descriptor.getIdentity(), new LazyConnectionDataSourceProxy(
                dataSourceToUse));
    }
}
 

@Test
public void testTransactionCommitWithNonExistingDatabaseAndLazyConnection() throws Exception {
	DriverManagerDataSource dsTarget = new DriverManagerDataSource();
	final LazyConnectionDataSourceProxy ds = new LazyConnectionDataSourceProxy();
	ds.setTargetDataSource(dsTarget);
	ds.setDefaultAutoCommit(true);
	ds.setDefaultTransactionIsolation(Connection.TRANSACTION_READ_COMMITTED);
	//ds.setDefaultTransactionIsolationName("TRANSACTION_READ_COMMITTED");

	LocalSessionFactoryBean lsfb = new LocalSessionFactoryBean();
	lsfb.setDataSource(ds);
	Properties props = new Properties();
	props.setProperty("hibernate.dialect", HSQLDialect.class.getName());
	props.setProperty("hibernate.cache.provider_class", NoCacheProvider.class.getName());
	props.setProperty("hibernate.temp.use_jdbc_metadata_defaults", "false");
	lsfb.setHibernateProperties(props);
	lsfb.afterPropertiesSet();
	final SessionFactory sf = lsfb.getObject();

	HibernateTransactionManager tm = new HibernateTransactionManager();
	tm.setSessionFactory(sf);
	tm.afterPropertiesSet();
	TransactionTemplate tt = new TransactionTemplate(tm);
	tt.setIsolationLevel(TransactionDefinition.ISOLATION_SERIALIZABLE);
	tt.setTimeout(10);
	assertTrue("Hasn't thread session", !TransactionSynchronizationManager.hasResource(sf));
	assertTrue("Hasn't thread connection", !TransactionSynchronizationManager.hasResource(ds));
	assertTrue("JTA synchronizations not active", !TransactionSynchronizationManager.isSynchronizationActive());

	tt.execute(new TransactionCallback() {
		@Override
		public Object doInTransaction(TransactionStatus status) {
			assertTrue("Has thread session", TransactionSynchronizationManager.hasResource(sf));
			assertTrue("Has thread connection", TransactionSynchronizationManager.hasResource(ds));
			HibernateTemplate ht = new HibernateTemplate(sf);
			return ht.find("from java.lang.Object");
		}
	});

	assertTrue("Hasn't thread session", !TransactionSynchronizationManager.hasResource(sf));
	assertTrue("Hasn't thread connection", !TransactionSynchronizationManager.hasResource(ds));
	assertTrue("JTA synchronizations not active", !TransactionSynchronizationManager.isSynchronizationActive());
}
 

@Test
public void testTransactionCommitWithNonExistingDatabaseAndLazyConnection() throws Exception {
	DriverManagerDataSource dsTarget = new DriverManagerDataSource();
	final LazyConnectionDataSourceProxy ds = new LazyConnectionDataSourceProxy();
	ds.setTargetDataSource(dsTarget);
	ds.setDefaultAutoCommit(true);
	ds.setDefaultTransactionIsolation(Connection.TRANSACTION_READ_COMMITTED);
	//ds.setDefaultTransactionIsolationName("TRANSACTION_READ_COMMITTED");

	LocalSessionFactoryBean lsfb = new LocalSessionFactoryBean();
	lsfb.setDataSource(ds);
	Properties props = new Properties();
	props.setProperty("hibernate.dialect", HSQLDialect.class.getName());
	props.setProperty("hibernate.temp.use_jdbc_metadata_defaults", "false");
	lsfb.setHibernateProperties(props);
	lsfb.afterPropertiesSet();
	final SessionFactory sf = lsfb.getObject();

	HibernateTransactionManager tm = new HibernateTransactionManager();
	tm.setSessionFactory(sf);
	tm.afterPropertiesSet();
	TransactionTemplate tt = new TransactionTemplate(tm);
	tt.setIsolationLevel(TransactionDefinition.ISOLATION_SERIALIZABLE);
	tt.setTimeout(10);
	assertTrue("Hasn't thread session", !TransactionSynchronizationManager.hasResource(sf));
	assertTrue("Hasn't thread connection", !TransactionSynchronizationManager.hasResource(ds));
	assertTrue("JTA synchronizations not active", !TransactionSynchronizationManager.isSynchronizationActive());

	tt.execute(new TransactionCallback() {
		@Override
		public Object doInTransaction(TransactionStatus status) {
			assertTrue("Has thread session", TransactionSynchronizationManager.hasResource(sf));
			assertTrue("Has thread connection", TransactionSynchronizationManager.hasResource(ds));
			Session session = ((SessionHolder) TransactionSynchronizationManager.getResource(sf)).getSession();
			return session.createQuery("from java.lang.Object").list();
		}
	});

	assertTrue("Hasn't thread session", !TransactionSynchronizationManager.hasResource(sf));
	assertTrue("Hasn't thread connection", !TransactionSynchronizationManager.hasResource(ds));
	assertTrue("JTA synchronizations not active", !TransactionSynchronizationManager.isSynchronizationActive());
}
 

@Test
void getConnection_ReadReplicaAvailableReadOnlyTransactionActive_returnsReadReplicaDataSource()
		throws Exception {

	// Arrange
	DataSource defaultDataSource = mock(DataSource.class);
	Connection connection = mock(Connection.class);

	DataSource readOnlyDataSource = mock(DataSource.class);
	Connection readOnlyConnection = mock(Connection.class);

	when(readOnlyDataSource.getConnection()).thenReturn(readOnlyConnection);
	when(defaultDataSource.getConnection()).thenReturn(connection);

	ReadOnlyRoutingDataSource readOnlyRoutingDataSource = new ReadOnlyRoutingDataSource();
	readOnlyRoutingDataSource.setTargetDataSources(
			Collections.singletonMap("read1", readOnlyDataSource));
	readOnlyRoutingDataSource.setDefaultTargetDataSource(defaultDataSource);
	readOnlyRoutingDataSource.afterPropertiesSet();

	LazyConnectionDataSourceProxy dataSource = new LazyConnectionDataSourceProxy(
			readOnlyRoutingDataSource);

	DefaultTransactionDefinition transactionDefinition = new DefaultTransactionDefinition();
	transactionDefinition.setReadOnly(true);

	TransactionTemplate transactionTemplate = new TransactionTemplate(
			new DataSourceTransactionManager(dataSource), transactionDefinition);

	// Act
	Connection connectionReturned = transactionTemplate.execute(status -> {
		try {
			return ((ConnectionProxy) dataSource.getConnection())
					.getTargetConnection();
		}
		catch (SQLException e) {
			fail(e.getMessage());
		}
		return null;
	});

	// Assert
	assertThat(connectionReturned).isSameAs(readOnlyConnection);
}
 

@Test
void getConnection_ReadReplicaAvailableWriteTransactionActive_returnsDefaultDataSource()
		throws Exception {

	// Arrange
	DataSource defaultDataSource = mock(DataSource.class);
	Connection connection = mock(Connection.class);

	DataSource readOnlyDataSource = mock(DataSource.class);
	Connection readOnlyConnection = mock(Connection.class);

	when(readOnlyDataSource.getConnection()).thenReturn(readOnlyConnection);
	when(defaultDataSource.getConnection()).thenReturn(connection);

	ReadOnlyRoutingDataSource readOnlyRoutingDataSource = new ReadOnlyRoutingDataSource();
	readOnlyRoutingDataSource.setTargetDataSources(
			Collections.singletonMap("read1", readOnlyDataSource));
	readOnlyRoutingDataSource.setDefaultTargetDataSource(defaultDataSource);
	readOnlyRoutingDataSource.afterPropertiesSet();

	LazyConnectionDataSourceProxy dataSource = new LazyConnectionDataSourceProxy(
			readOnlyRoutingDataSource);

	DefaultTransactionDefinition transactionDefinition = new DefaultTransactionDefinition();
	transactionDefinition.setReadOnly(false);

	TransactionTemplate transactionTemplate = new TransactionTemplate(
			new DataSourceTransactionManager(dataSource), transactionDefinition);

	// Act
	Connection connectionReturned = transactionTemplate.execute(status -> {
		try {
			return ((ConnectionProxy) dataSource.getConnection())
					.getTargetConnection();
		}
		catch (SQLException e) {
			fail(e.getMessage());
		}
		return null;
	});

	// Assert
	assertThat(connectionReturned).isSameAs(connection);
}
 

@Test
void afterPropertiesSet_instanceWithReadReplica_createsDataSourceRouter()
		throws Exception {
	// Arrange
	AmazonRDS amazonRDS = mock(AmazonRDS.class);
	DataSourceFactory dataSourceFactory = mock(DataSourceFactory.class);

	when(amazonRDS.describeDBInstances(
			new DescribeDBInstancesRequest().withDBInstanceIdentifier("test")))
					.thenReturn(new DescribeDBInstancesResult().withDBInstances(
							new DBInstance().withDBInstanceStatus("available")
									.withDBName("test")
									.withDBInstanceIdentifier("test")
									.withEngine("mysql").withMasterUsername("admin")
									.withEndpoint(new Endpoint()
											.withAddress("localhost").withPort(3306))
									.withReadReplicaDBInstanceIdentifiers("read1",
											"read2")));

	when(amazonRDS.describeDBInstances(new DescribeDBInstancesRequest()
			.withDBInstanceIdentifier("read1"))).thenReturn(
					new DescribeDBInstancesResult().withDBInstances(new DBInstance()
							.withDBInstanceStatus("available").withDBName("read1")
							.withDBInstanceIdentifier("read1").withEngine("mysql")
							.withMasterUsername("admin").withEndpoint(new Endpoint()
									.withAddress("localhost").withPort(3306))));

	when(amazonRDS.describeDBInstances(new DescribeDBInstancesRequest()
			.withDBInstanceIdentifier("read2"))).thenReturn(
					new DescribeDBInstancesResult().withDBInstances(new DBInstance()
							.withDBInstanceStatus("available").withDBName("read2")
							.withDBInstanceIdentifier("read2").withEngine("mysql")
							.withMasterUsername("admin").withEndpoint(new Endpoint()
									.withAddress("localhost").withPort(3306))));

	DataSource createdDataSource = mock(DataSource.class);
	Connection connection = mock(Connection.class);

	when(dataSourceFactory.createDataSource(new DataSourceInformation(
			DatabaseType.MYSQL, "localhost", 3306, "test", "admin", "secret")))
					.thenReturn(createdDataSource);
	when(dataSourceFactory.createDataSource(new DataSourceInformation(
			DatabaseType.MYSQL, "localhost", 3306, "read1", "admin", "secret")))
					.thenReturn(createdDataSource);
	when(dataSourceFactory.createDataSource(new DataSourceInformation(
			DatabaseType.MYSQL, "localhost", 3306, "read2", "admin", "secret")))
					.thenReturn(createdDataSource);
	when(createdDataSource.getConnection()).thenReturn(connection);

	AmazonRdsReadReplicaAwareDataSourceFactoryBean factoryBean = new AmazonRdsReadReplicaAwareDataSourceFactoryBean(
			amazonRDS, "test", "secret");
	factoryBean.setDataSourceFactory(dataSourceFactory);

	// Act
	factoryBean.afterPropertiesSet();

	// Assert
	DataSource datasource = factoryBean.getObject();
	assertThat(datasource).isNotNull();
	assertThat(datasource instanceof LazyConnectionDataSourceProxy).isTrue();

	ReadOnlyRoutingDataSource source = (ReadOnlyRoutingDataSource) ((LazyConnectionDataSourceProxy) datasource)
			.getTargetDataSource();
	assertThat(source.getDataSources().size()).isEqualTo(3);
}
 

/**
 * {@link org.springframework.jdbc.datasource.LazyConnectionDataSourceProxy}로 감싸서
 * 트랜잭션 동기화가 이루어진 뒤에 실제 커넥션을 확보하도록 해준다.
 *
 * @param routingDataSource
 * @return
 */
@Bean
public DataSource dataSource(@Qualifier("routingDataSource") DataSource routingDataSource) {
    return new LazyConnectionDataSourceProxy(routingDataSource);
}
 

/**
 * {@link org.springframework.jdbc.datasource.LazyConnectionDataSourceProxy}로 감싸서
 * 트랜잭션 동기화가 이루어진 뒤에 실제 커넥션을 확보하도록 해준다.
 *
 * @param routingDataSource
 * @return
 */
@Bean
public DataSource dataSource(@Qualifier("routingDataSource") DataSource routingDataSource) {
    return new LazyConnectionDataSourceProxy(routingDataSource);
}