MasterSlaveDataSource.java

/*
 * Copyright 1999-2015 dangdang.com.
 * <p>
 * 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.
 * </p>
 */

package com.dangdang.ddframe.rdb.sharding.jdbc;

import com.dangdang.ddframe.rdb.sharding.api.strategy.slave.RoundRobinSlaveLoadBalanceStrategy;
import com.dangdang.ddframe.rdb.sharding.api.strategy.slave.SlaveLoadBalanceStrategy;
import com.dangdang.ddframe.rdb.sharding.hint.HintManagerHolder;
import com.dangdang.ddframe.rdb.sharding.jdbc.adapter.AbstractDataSourceAdapter;
import com.dangdang.ddframe.rdb.sharding.parser.result.router.SQLStatementType;
import com.google.common.base.Preconditions;
import lombok.RequiredArgsConstructor;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.List;

/**
 * 支持读写分离的数据源.
 *
 * @author zhangliang
 */
@RequiredArgsConstructor
public final class MasterSlaveDataSource extends AbstractDataSourceAdapter {
    
    private static final ThreadLocal<Boolean> DML_FLAG = new ThreadLocal<Boolean>() {
        
        @Override
        protected Boolean initialValue() {
            return false;
        }
    };
    
    private final String name;
    
    private final DataSource masterDataSource;
    
    private final List<DataSource> slaveDataSources;
    
    private final SlaveLoadBalanceStrategy slaveLoadBalanceStrategy = new RoundRobinSlaveLoadBalanceStrategy();
    
    
    static boolean isDML(final SQLStatementType sqlStatementType) {
        return SQLStatementType.SELECT != sqlStatementType || DML_FLAG.get() || HintManagerHolder.isMasterRouteOnly();
    }
    
    /**
     * 获取主或从节点的数据源名称.
     *
     * @param sqlStatementType SQL类型
     * @return 主或从节点的数据源
     */
    public DataSource getDataSource(final SQLStatementType sqlStatementType) {
        if (isDML(sqlStatementType)) {
            DML_FLAG.set(true);
            return masterDataSource;
        }
        return slaveLoadBalanceStrategy.getDataSource(name, slaveDataSources);
    }
    
    String getDatabaseProductName() throws SQLException {
        String result;
        try (Connection masterConnection = masterDataSource.getConnection()) {
            result = masterConnection.getMetaData().getDatabaseProductName();
        }
        for (DataSource each : slaveDataSources) {
            String slaveDatabaseProductName;
            try (Connection slaveConnection = each.getConnection()) {
                slaveDatabaseProductName = slaveConnection.getMetaData().getDatabaseProductName();    
            }
            Preconditions.checkState(result.equals(slaveDatabaseProductName), String.format("Database type inconsistent with '%s' and '%s'", result, slaveDatabaseProductName));
        }
        return result;
    }
    
    @Override
    public Connection getConnection() throws SQLException {
        throw new UnsupportedOperationException("Master slave data source cannot support get connection directly.");
    }
    
    /**
     * 重置更新标记.
     */
    public static void resetDMLFlag() {
        DML_FLAG.remove();
    }
}