com.nesscomputing.jdbc.wrappers.ConnectionWrapper Maven / Gradle / Ivy

Go to download
/**
 * Copyright (C) 2012 Ness Computing, Inc.
 *
 * Licensed 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 com.nesscomputing.jdbc.wrappers;

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Map;
import java.util.Set;

import javax.annotation.Nullable;
import javax.sql.DataSource;

import com.google.common.base.Function;
import com.google.common.base.Functions;
import com.google.common.collect.Maps;
import com.google.inject.Binding;
import com.google.inject.Inject;
import com.google.inject.Injector;
import com.google.inject.Key;
import com.google.inject.Singleton;
import com.google.inject.TypeLiteral;
import com.nesscomputing.jdbc.wrappers.AbstractProxyInvocationHandler.MethodInterceptor;


/**
 * Adds wrappers to all handed out connections. These in turn can be wrapped into other interceptors.
 */
@Singleton
public final class ConnectionWrapper implements Function
{
    private final Map> interceptors = Maps.newHashMap();
    private final Annotation annotation;

    private Function connectionWrapper = Functions.identity();

    public ConnectionWrapper(final Annotation annotation)
    {
        this.annotation = annotation;

        // Intercept all attempts to hand out a connection.
        final MethodInterceptor connectionInterceptor = new ConnectionInterceptor();
        try {
            // @{link DataSource#getConnection()}
            interceptors.put(DataSource.class.getMethod("getConnection", new Class[0]), connectionInterceptor);
            // @{link DataSource#getConnection(String username, String password)}
            interceptors.put(DataSource.class.getMethod("getConnection", new Class[] {String.class, String.class}), connectionInterceptor);
        }
        catch (NoSuchMethodException nsme) {
            throw new ExceptionInInitializerError(nsme);
        }
    }

    @Inject
    void getWrappers(final Injector injector)
    {
        final Binding>> connectionBindings = injector.getExistingBinding(Key.get(new TypeLiteral>> () { }, annotation));
        if (connectionBindings != null) {
            for (Function fn : connectionBindings.getProvider().get()) {
                connectionWrapper = Functions.compose(connectionWrapper, fn);
            }
        }
    }

    /**
     * Returns a proxy for a data source. All methods on the datasource object are passed through the proxy. In addition, the
     * DataSource also implements {@link java.io.Closeable}.
     */
    @Override
    public DataSource apply(@Nullable final DataSource dataSource)
    {
        if (dataSource == null) {
            return null;
        }
        else {
            return (DataSource) Proxy.newProxyInstance(dataSource.getClass().getClassLoader(),
                                                       dataSource.getClass().getInterfaces(),
                                                       new ConnectionWrapperInvocationHandler(dataSource));
        }
    }

    class ConnectionWrapperInvocationHandler extends AbstractProxyInvocationHandler
    {
        ConnectionWrapperInvocationHandler(final DataSource dataSource)
        {
            super(dataSource);
        }

        @Override
        protected Object intercept(final Method method, final Object object) throws Throwable
        {
            final MethodInterceptor interceptor = interceptors.get(method);
            return (interceptor == null) ? object : interceptor.intercept(Connection.class.cast(object));
        }
    }

    /**
     * Wrap a connection object by applying all the wrappers.
     */
    class ConnectionInterceptor implements MethodInterceptor
    {
        @Override
        public Connection intercept(final Connection parent) throws SQLException
        {
            final Connection c = (Connection) parent;

            if (c == null) {
                throw new IllegalStateException("null target in method interceptor");
            }

            return connectionWrapper.apply(c);
        }
    }
}