Handle inheritance hierarchies better.

src/main/java/com/zaxxer/hikari/javassist/HikariClassTransformer.java

@@ -219,10 +219,7 @@ public class HikariClassTransformer implements ClassFileTransformer
         CtMethod[] destMethods = targetClass.getMethods();
         ConstPool constPool = targetClassFile.getConstPool();
 
-        HashSet<CtMethod> srcMethods = new HashSet<CtMethod>();
-        srcMethods.addAll(Arrays.asList(srcClass.getMethods()));
-        srcMethods.addAll(Arrays.asList(srcClass.getDeclaredMethods()));
-        for (CtMethod method : srcMethods)
+        for (CtMethod method : srcClass.getDeclaredMethods())
         {
             if (method.getAnnotation(HikariInject.class) == null)
             {
@@ -267,6 +264,8 @@ public class HikariClassTransformer implements ClassFileTransformer
         {
             CtConstructor copy = CtNewConstructor.copy(srcInitializer, targetClass, null);
             targetClass.addConstructor(copy);
+            CtMethod __static = CtNewMethod.make(Modifier.STATIC, CtClass.voidType, "__static", null, null, "{}", targetClass);
+            targetClass.addMethod(__static);
             LOGGER.debug("Copied static initializer of {} to {}", srcClass.getSimpleName(), targetClass.getSimpleName());
         }
         else
@@ -275,13 +274,16 @@ public class HikariClassTransformer implements ClassFileTransformer
             targetClass.addMethod(method);
             targetClass.removeConstructor(destInitializer);
             LOGGER.debug("Move static initializer of {}", targetClass.getSimpleName());
-            mergeClassInitializers(srcClass, targetClass, targetClassFile);
+            // mergeClassInitializers(srcClass, targetClass, targetClassFile);
+            CtConstructor copy = CtNewConstructor.copy(srcInitializer, targetClass, null);
+            targetClass.addConstructor(copy);
+            LOGGER.debug("Copied static initializer of {} to {}", srcClass.getSimpleName(), targetClass.getSimpleName());
         }
     }
 
     private void injectTryCatch(CtClass targetClass) throws Exception
     {
-        for (CtMethod method : targetClass.getMethods())
+        for (CtMethod method : targetClass.getDeclaredMethods())
         {
             if ((method.getModifiers() & Modifier.PUBLIC) != Modifier.PUBLIC ||  // only public methods
                 method.getAnnotation(HikariInject.class) != null)                // ignore methods we've injected, they already try..catch
@@ -298,6 +300,7 @@ public class HikariClassTransformer implements ClassFileTransformer
             {
                 if ("java.sql.SQLException".equals(exception.getName()))         // only add try..catch to methods throwing SQLException
                 {
+                    LOGGER.debug("Injecting try..catch into {}{}", method.getName(), method.getSignature());
                     method.addCatch("throw checkException($e);", exception);
                     break;
                 }

src/main/java/com/zaxxer/hikari/javassist/HikariInstrumentationAgent.java

@@ -58,9 +58,9 @@ public class HikariInstrumentationAgent
         completionMap = new HashMap<String, Boolean>();
         completionMap.put("java.sql.Connection", false);
         completionMap.put("java.sql.ResultSet", false);
-        completionMap.put("java.sql.Statement", false);
-        completionMap.put("java.sql.CallableStatement", false);
-        completionMap.put("java.sql.PreparedStatement", false);
+//        completionMap.put("java.sql.Statement", false);
+//        completionMap.put("java.sql.CallableStatement", false);
+//        completionMap.put("java.sql.PreparedStatement", false);
     }
 
     private DataSource dataSource;
@@ -112,10 +112,10 @@ public class HikariInstrumentationAgent
         {
             return false;
         }
-        finally
-        {
-            unregisterInstrumenation();
-        }
+//        finally
+//        {
+//            unregisterInstrumenation();
+//        }
     }
 
     /**

src/main/java/com/zaxxer/hikari/proxy/CallableStatementProxy.java

@@ -17,21 +17,167 @@
 package com.zaxxer.hikari.proxy;
 
 import java.sql.CallableStatement;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+
+import com.zaxxer.hikari.javassist.HikariInject;
 
 /**
  *
  * @author Brett Wooldridge
  */
-public class CallableStatementProxy extends PreparedStatementProxy
+public class CallableStatementProxy implements IHikariStatementProxy
 {
+    private static ProxyFactory PROXY_FACTORY;
+
+    @HikariInject private IHikariConnectionProxy _connection;
+    
+    protected Statement delegate;
+
+    static
+    {
+        __static();
+    }
+
     protected CallableStatementProxy(ConnectionProxy connection, CallableStatement statement)
     {
-        super(connection, statement);
+        this._connection = connection;
+        this.delegate = statement;
+    }
+
+    @HikariInject
+    public void setConnectionProxy(IHikariConnectionProxy connection)
+    {
+        this._connection = connection;
+    }
+
+    @HikariInject
+    public SQLException checkException(SQLException e)
+    {
+        return _connection.checkException(e);
     }
 
+
     // **********************************************************************
     //               Overridden java.sql.CallableStatement Methods
     // **********************************************************************
 
-    // TODO implement wrapper
+    @HikariInject
+    public ResultSet executeQuery() throws SQLException
+    {
+        try
+        {
+            IHikariResultSetProxy resultSet = (IHikariResultSetProxy) __executeQuery();
+            if (resultSet == null)
+            {
+                return null;
+            }
+
+            resultSet.setProxyStatement(this);
+            return (ResultSet) resultSet;
+        }
+        catch (SQLException e)
+        {
+            throw checkException(e);
+        }
+    }
+
+    @HikariInject
+    public void close() throws SQLException
+    {
+        _connection.unregisterStatement(this);
+        try
+        {
+            __close();
+        }
+        catch (SQLException e)
+        {
+            throw checkException(e);
+        }
+    }
+
+    @HikariInject
+    public ResultSet executeQuery(String sql) throws SQLException
+    {
+        try
+        {
+            IHikariResultSetProxy resultSet = (IHikariResultSetProxy) __executeQuery(sql);
+            if (resultSet == null)
+            {
+                return null;
+            }
+
+            resultSet.setProxyStatement(this);  
+            return (ResultSet) resultSet;
+        }
+        catch (SQLException e)
+        {
+            throw checkException(e);
+        }
+    }
+
+    @HikariInject
+    public ResultSet getGeneratedKeys() throws SQLException
+    {
+        try
+        {
+            IHikariResultSetProxy resultSet = (IHikariResultSetProxy) __getGeneratedKeys();
+            if (resultSet == null)
+            {
+                return null;
+            }
+
+            resultSet.setProxyStatement(this);  
+            return (ResultSet) resultSet;
+        }
+        catch (SQLException e)
+        {
+            throw checkException(e);
+        }
+    }
+
+    // ***********************************************************************
+    // These methods contain code we do not want injected into the actual
+    // java.sql.Connection implementation class.  These methods are only
+    // used when instrumentation is not available and "conventional" Javassist
+    // delegating proxies are used.
+    // ***********************************************************************
+
+    public ResultSet __executeQuery() throws SQLException
+    {
+        ResultSet resultSet = ((PreparedStatement) delegate).executeQuery();
+        return PROXY_FACTORY.getProxyResultSet(this, resultSet);
+    }
+
+    private static void __static()
+    {
+        if (PROXY_FACTORY == null)
+        {
+            PROXY_FACTORY = JavassistProxyFactoryFactory.getProxyFactory();
+        }
+    }
+    
+    public void __close() throws SQLException
+    {
+        if (delegate.isClosed())
+        {
+            return;
+        }
+
+        delegate.close();        
+    }
+
+    public ResultSet __executeQuery(String sql) throws SQLException
+    {
+        ResultSet resultSet = delegate.executeQuery(sql);
+        return PROXY_FACTORY.getProxyResultSet(this, resultSet);
+    }
+
+    public ResultSet __getGeneratedKeys() throws SQLException
+    {
+        ResultSet generatedKeys = delegate.getGeneratedKeys();
+        return PROXY_FACTORY.getProxyResultSet(this, generatedKeys);
+    }
 }

src/main/java/com/zaxxer/hikari/proxy/PreparedStatementProxy.java

@@ -19,6 +19,7 @@ package com.zaxxer.hikari.proxy;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Statement;
 
 import com.zaxxer.hikari.javassist.HikariInject;
 
@@ -26,11 +27,35 @@ import com.zaxxer.hikari.javassist.HikariInject;
  *
  * @author Brett Wooldridge
  */
-public class PreparedStatementProxy extends StatementProxy
+public class PreparedStatementProxy implements IHikariStatementProxy
 {
+    private static ProxyFactory PROXY_FACTORY;
+
+    @HikariInject private IHikariConnectionProxy _connection;
+    
+    protected Statement delegate;
+
+    static
+    {
+        __static();
+    }
+
     protected PreparedStatementProxy(ConnectionProxy connection, PreparedStatement statement)
     {
-        super(connection, statement);
+        this._connection = connection;
+        this.delegate = statement;
+    }
+
+    @HikariInject
+    public void setConnectionProxy(IHikariConnectionProxy connection)
+    {
+        this._connection = connection;
+    }
+
+    @HikariInject
+    public SQLException checkException(SQLException e)
+    {
+        return _connection.checkException(e);
     }
 
     // **********************************************************************
@@ -57,6 +82,60 @@ public class PreparedStatementProxy extends StatementProxy
     	}
     }
 
+    @HikariInject
+    public void close() throws SQLException
+    {
+        _connection.unregisterStatement(this);
+        try
+        {
+            __close();
+        }
+        catch (SQLException e)
+        {
+            throw checkException(e);
+        }
+    }
+
+    @HikariInject
+    public ResultSet executeQuery(String sql) throws SQLException
+    {
+        try
+        {
+            IHikariResultSetProxy resultSet = (IHikariResultSetProxy) __executeQuery(sql);
+            if (resultSet == null)
+            {
+                return null;
+            }
+
+            resultSet.setProxyStatement(this);  
+            return (ResultSet) resultSet;
+        }
+        catch (SQLException e)
+        {
+            throw checkException(e);
+        }
+    }
+
+    @HikariInject
+    public ResultSet getGeneratedKeys() throws SQLException
+    {
+        try
+        {
+            IHikariResultSetProxy resultSet = (IHikariResultSetProxy) __getGeneratedKeys();
+            if (resultSet == null)
+            {
+                return null;
+            }
+
+            resultSet.setProxyStatement(this);  
+            return (ResultSet) resultSet;
+        }
+        catch (SQLException e)
+        {
+            throw checkException(e);
+        }
+    }
+
     // ***********************************************************************
     // These methods contain code we do not want injected into the actual
     // java.sql.Connection implementation class.  These methods are only
@@ -70,5 +149,33 @@ public class PreparedStatementProxy extends StatementProxy
         return PROXY_FACTORY.getProxyResultSet(this, resultSet);
     }
 
-    // TODO: fix wrapper
+    private static void __static()
+    {
+        if (PROXY_FACTORY == null)
+        {
+            PROXY_FACTORY = JavassistProxyFactoryFactory.getProxyFactory();
+        }
+    }
+    
+    public void __close() throws SQLException
+    {
+        if (delegate.isClosed())
+        {
+            return;
+        }
+
+        delegate.close();        
+    }
+
+    public ResultSet __executeQuery(String sql) throws SQLException
+    {
+        ResultSet resultSet = delegate.executeQuery(sql);
+        return PROXY_FACTORY.getProxyResultSet(this, resultSet);
+    }
+
+    public ResultSet __getGeneratedKeys() throws SQLException
+    {
+        ResultSet generatedKeys = delegate.getGeneratedKeys();
+        return PROXY_FACTORY.getProxyResultSet(this, generatedKeys);
+    }
 }

src/main/java/com/zaxxer/hikari/proxy/StatementProxy.java

@@ -27,9 +27,9 @@ import com.zaxxer.hikari.javassist.HikariInject;
  */
 public class StatementProxy implements IHikariStatementProxy
 {
-    protected static ProxyFactory PROXY_FACTORY;
+    private static ProxyFactory PROXY_FACTORY;
 
-    @HikariInject protected IHikariConnectionProxy _connection;
+    @HikariInject private IHikariConnectionProxy _connection;
     
     protected Statement delegate;
 

src/test/java/com/zaxxer/hikari/mocks/StubStatement.java

@@ -21,7 +21,6 @@ import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.SQLWarning;
 import java.sql.Statement;
-import java.util.ArrayList;
 
 /**
  *
@@ -30,7 +29,6 @@ import java.util.ArrayList;
 public class StubStatement implements Statement
 {
     private boolean closed;
-    private ArrayList<ResultSet> resultSets = new ArrayList<ResultSet>();
 
     /** {@inheritDoc} */
     public <T> T unwrap(Class<T> iface) throws SQLException
@@ -48,7 +46,6 @@ public class StubStatement implements Statement
     public ResultSet executeQuery(String sql) throws SQLException
     {
         StubResultSet resultSet = new StubResultSet();
-        resultSets.add(resultSet);
         return resultSet;
     }
 
@@ -61,11 +58,6 @@ public class StubStatement implements Statement
     /** {@inheritDoc} */
     public void close() throws SQLException
     {
-        for (ResultSet resultSet : resultSets)
-        {
-            resultSet.close();
-        }
-
         closed = true;
     }