From bcc6722b50220038a920216aed3ec3969396bc2b Mon Sep 17 00:00:00 2001
From: wenshao <szujobs@qq.com>
Date: Fri, 9 Dec 2011 13:46:59 +0800
Subject: [PATCH] ha

---
 .../alibaba/druid/pool/ha/HADataSource.java   |  2 +-
 .../pool/ha/MultiDataSourceConnection.java    | 20 ++++
 .../druid/bvt/pool/ha/HADataSourceTest2.java  | 92 ++++++++++++-------
 3 files changed, 82 insertions(+), 32 deletions(-)

diff --git a/src/main/java/com/alibaba/druid/pool/ha/HADataSource.java b/src/main/java/com/alibaba/druid/pool/ha/HADataSource.java
index 20c992b9f2..63ef3071a9 100644
--- a/src/main/java/com/alibaba/druid/pool/ha/HADataSource.java
+++ b/src/main/java/com/alibaba/druid/pool/ha/HADataSource.java
@@ -99,7 +99,7 @@ public class HADataSource extends MultiDataSource implements HADataSourceMBean,
             masterConnectCount.incrementAndGet();
         }
 
-        if (slave.isEnable()) {
+        if (conn == null && slave.isEnable()) {
             conn = slave.getConnection();
             slaveConnectCount.incrementAndGet();
         }
diff --git a/src/main/java/com/alibaba/druid/pool/ha/MultiDataSourceConnection.java b/src/main/java/com/alibaba/druid/pool/ha/MultiDataSourceConnection.java
index 6a6210dae6..7902d437e8 100644
--- a/src/main/java/com/alibaba/druid/pool/ha/MultiDataSourceConnection.java
+++ b/src/main/java/com/alibaba/druid/pool/ha/MultiDataSourceConnection.java
@@ -579,6 +579,26 @@ public class MultiDataSourceConnection extends WrapperAdapter implements Connect
     public Date getConnectedTime() {
         return connectedTime;
     }
+    
+    public boolean isWrapperFor(Class<?> iface) throws SQLException {
+    	boolean result = super.isWrapperFor(iface);
+    	
+    	if (result == false && conn != null) {
+    		result = conn.isWrapperFor(iface);
+    	}
+    	
+    	return result;
+    }
+    
+    public <T> T unwrap(Class<T> iface) throws SQLException {
+    	T object = super.unwrap(iface);
+    	
+    	if (object == null && conn != null) {
+    		object = conn.unwrap(iface);
+    	}
+    	
+    	return object;
+    }
 
     public void setSchema(String schema) throws SQLException {
         throw new SQLFeatureNotSupportedException();
diff --git a/src/test/java/com/alibaba/druid/bvt/pool/ha/HADataSourceTest2.java b/src/test/java/com/alibaba/druid/bvt/pool/ha/HADataSourceTest2.java
index 5a6efc7bcc..b69bd75ad6 100644
--- a/src/test/java/com/alibaba/druid/bvt/pool/ha/HADataSourceTest2.java
+++ b/src/test/java/com/alibaba/druid/bvt/pool/ha/HADataSourceTest2.java
@@ -1,6 +1,7 @@
 package com.alibaba.druid.bvt.pool.ha;
 
 import java.sql.Connection;
+import java.sql.Statement;
 
 import junit.framework.Assert;
 import junit.framework.TestCase;
@@ -12,41 +13,70 @@ import com.alibaba.druid.stat.DruidDataSourceStatManager;
 
 public class HADataSourceTest2 extends TestCase {
 
-    private DruidDataSource dataSourceA;
-    private DruidDataSource dataSourceB;
+	private DruidDataSource dataSourceA;
+	private DruidDataSource dataSourceB;
 
-    private HADataSource    dataSourceHA;
+	private HADataSource dataSourceHA;
 
-    protected void setUp() throws Exception {
-        Assert.assertEquals(0, DruidDataSourceStatManager.getInstance().getDataSourceList().size());
-        
-        dataSourceA = new DruidDataSource();
-        dataSourceA.setUrl("jdbc:mock:ha1");
-        dataSourceA.setFilters("trace");
+	private String MASTER_URL = "jdbc:mock:master";
+	private String SLAVE_URL = "jdbc:mock:slave";
 
-        dataSourceB = new DruidDataSource();
-        dataSourceB.setUrl("jdbc:mock:ha2");
-        dataSourceB.setFilters("stat");
+	protected void setUp() throws Exception {
+		Assert.assertEquals(0, DruidDataSourceStatManager.getInstance()
+				.getDataSourceList().size());
 
-        dataSourceHA = new HADataSource();
-        dataSourceHA.setMaster(dataSourceA);
-        dataSourceHA.setSlave(dataSourceB);
-    }
+		dataSourceA = new DruidDataSource();
+		dataSourceA.setUrl(MASTER_URL);
+		dataSourceA.setFilters("trace");
 
-    protected void tearDown() throws Exception {
-        dataSourceHA.close();
-        
-        for (DruidDataSource dataSource : DruidDataSourceStatManager.getDruidDataSourceInstances()) {
-            dataSource.close();
-        }
-        
-        Assert.assertEquals(0, DruidDataSourceStatManager.getInstance().getDataSourceList().size());
-    }
+		dataSourceB = new DruidDataSource();
+		dataSourceB.setUrl(SLAVE_URL);
+		dataSourceB.setFilters("stat");
 
-   public void test_switch() throws Exception {
-       Connection conn = dataSourceHA.getConnection();
-       
-       MockConnection mockConn = conn.unwrap(MockConnection.class);
-       conn.close();
-   }
+		dataSourceHA = new HADataSource();
+		dataSourceHA.setMaster(dataSourceA);
+		dataSourceHA.setSlave(dataSourceB);
+	}
+
+	protected void tearDown() throws Exception {
+		dataSourceHA.close();
+
+		for (DruidDataSource dataSource : DruidDataSourceStatManager
+				.getDruidDataSourceInstances()) {
+			dataSource.close();
+		}
+
+		Assert.assertEquals(0, DruidDataSourceStatManager.getInstance()
+				.getDataSourceList().size());
+	}
+
+	public void test_switch() throws Exception {
+		{
+			Connection conn = dataSourceHA.getConnection();
+
+			Statement stmt = conn.createStatement();
+			stmt.execute("select 1");
+
+			MockConnection mockConn = conn.unwrap(MockConnection.class);
+			Assert.assertNotNull(mockConn);
+
+			Assert.assertEquals(dataSourceA.getUrl(), mockConn.getUrl());
+
+			conn.close();
+		}
+		dataSourceHA.setMasterEnable(false);
+		{
+			Connection conn = dataSourceHA.getConnection();
+			
+			Statement stmt = conn.createStatement();
+			stmt.execute("select 1");
+			
+			MockConnection mockConn = conn.unwrap(MockConnection.class);
+			Assert.assertNotNull(mockConn);
+			
+			Assert.assertEquals(dataSourceB.getUrl(), mockConn.getUrl());
+			
+			conn.close();
+		}
+	}
 }