聊聊sharding-jdbc的ShardingMasterSlaveRouter

@RequiredArgsConstructor
public final class ShardingMasterSlaveRouter {
    
    private final Collection<MasterSlaveRule> masterSlaveRules;
    
    /**
     * Route Master slave after sharding.
     *
     * @param sqlRouteResult SQL route result
     * @return route result
     */
    public SQLRouteResult route(final SQLRouteResult sqlRouteResult) {
        for (MasterSlaveRule each : masterSlaveRules) {
            route(each, sqlRouteResult);
        }
        return sqlRouteResult;
    }
    
    private void route(final MasterSlaveRule masterSlaveRule, final SQLRouteResult sqlRouteResult) {
        Collection<TableUnit> toBeRemoved = new LinkedList<>();
        Collection<TableUnit> toBeAdded = new LinkedList<>();
        for (TableUnit each : sqlRouteResult.getRoutingResult().getTableUnits().getTableUnits()) {
            if (!masterSlaveRule.getName().equalsIgnoreCase(each.getDataSourceName())) {
                continue;
            }
            toBeRemoved.add(each);
            String actualDataSourceName;
            if (isMasterRoute(sqlRouteResult.getSqlStatement().getType())) {
                MasterVisitedManager.setMasterVisited();
                actualDataSourceName = masterSlaveRule.getMasterDataSourceName();
            } else {
                actualDataSourceName = masterSlaveRule.getLoadBalanceAlgorithm().getDataSource(
                        masterSlaveRule.getName(), masterSlaveRule.getMasterDataSourceName(), new ArrayList<>(masterSlaveRule.getSlaveDataSourceNames()));
            }
            toBeAdded.add(createNewTableUnit(actualDataSourceName, each));
        }
        sqlRouteResult.getRoutingResult().getTableUnits().getTableUnits().removeAll(toBeRemoved);
        sqlRouteResult.getRoutingResult().getTableUnits().getTableUnits().addAll(toBeAdded);
    }
    
    private boolean isMasterRoute(final SQLType sqlType) {
        return SQLType.DQL != sqlType || MasterVisitedManager.isMasterVisited() || HintManager.isMasterRouteOnly();
    }
    
    private TableUnit createNewTableUnit(final String actualDataSourceName, final TableUnit originalTableUnit) {
        TableUnit result = new TableUnit(actualDataSourceName, originalTableUnit.getDataSourceName());
        result.getRoutingTables().addAll(originalTableUnit.getRoutingTables());
        return result;
    }
}

@RequiredArgsConstructor
@Getter
@Setter
public final class SQLRouteResult {
    
    private final SQLStatement sqlStatement;
    
    private final GeneratedKey generatedKey;
    
    // For multiple thread read cached sqlStatement, clone limit on SQLRouteResult, because limit will be modified after cache
    // TODO need more good design here
    private Limit limit;
    
    private RoutingResult routingResult;
    
    private OptimizeResult optimizeResult;
    
    private final Collection<RouteUnit> routeUnits = new LinkedHashSet<>();
    
    public SQLRouteResult(final SQLStatement sqlStatement) {
        this(sqlStatement, null);
    }
}

@Getter
public class RoutingResult {
    
    private final TableUnits tableUnits = new TableUnits();
    
    /**
     * Judge is route for single database and table only or not.
     *
     * @return is route for single database and table only or not
     */
    public boolean isSingleRouting() {
        return 1 == tableUnits.getTableUnits().size();
    }
}

@Getter
@ToString
public final class TableUnits {
    
    private final Collection<TableUnit> tableUnits = new LinkedHashSet<>();
    
    /**
     * Get all data source names.
     *
     * @return all data source names
     */
    public Collection<String> getDataSourceNames() {
        Collection<String> result = new HashSet<>(tableUnits.size(), 1);
        for (TableUnit each : tableUnits) {
            result.add(each.getDataSourceName());
        }
        return result;
    }
    
    /**
     * Get routing table via data source name and actual table name.
     *
     * @param dataSourceName data source name
     * @param actualTableName actual table name
     * @return routing table
     */
    public Optional<RoutingTable> getRoutingTable(final String dataSourceName, final String actualTableName) {
        for (TableUnit each : tableUnits) {
            Optional<RoutingTable> result = each.getRoutingTable(dataSourceName, actualTableName);
            if (result.isPresent()) {
                return result;
            }
        }
        return Optional.absent();
    }
    
    /**
     * Get actual tables group via data source name and logic tables' names.
     * <p>
     * Actual tables in same group are belong one logic name.
     * </p>
     *
     * @param dataSourceName data source name
     * @param logicTableNames logic tables' names
     * @return actual tables group
     */
    public List<Set<String>> getActualTableNameGroups(final String dataSourceName, final Set<String> logicTableNames) {
        List<Set<String>> result = new ArrayList<>();
        for (String logicTableName : logicTableNames) {
            Set<String> actualTableNames = getActualTableNames(dataSourceName, logicTableName);
            if (!actualTableNames.isEmpty()) {
                result.add(actualTableNames);
            }
        }
        return result;
    }
    
    private Set<String> getActualTableNames(final String dataSourceName, final String logicTableName) {
        Set<String> result = new HashSet<>(tableUnits.size(), 1);
        for (TableUnit each : tableUnits) {
            result.addAll(each.getActualTableNames(dataSourceName, logicTableName));
        }
        return result;
    }
    
    /**
     * Get map relationship between data source and logic tables via data sources' names.
     *
     * @param dataSourceNames data sources' names
     * @return  map relationship between data source and logic tables
     */
    public Map<String, Set<String>> getDataSourceLogicTablesMap(final Collection<String> dataSourceNames) {
        Map<String, Set<String>> result = new HashMap<>();
        for (String each : dataSourceNames) {
            Set<String> logicTableNames = getLogicTableNames(each);
            if (!logicTableNames.isEmpty()) {
                result.put(each, logicTableNames);
            }
        }
        return result;
    }
    
    private Set<String> getLogicTableNames(final String dataSourceName) {
        Set<String> result = new HashSet<>(tableUnits.size(), 1);
        for (TableUnit each : tableUnits) {
            if (each.getDataSourceName().equalsIgnoreCase(dataSourceName)) {
                result.addAll(each.getLogicTableNames(dataSourceName));
            }
        }
        return result;
    }
}

@RequiredArgsConstructor
@Getter
@EqualsAndHashCode
@ToString
public final class TableUnit {
    
    private final String dataSourceName;
    
    private final String masterSlaveLogicDataSourceName;
    
    private final List<RoutingTable> routingTables = new LinkedList<>();
    
    public TableUnit(final String dataSourceName) {
        this.dataSourceName = dataSourceName;
        masterSlaveLogicDataSourceName = dataSourceName;
    }
    
    /**
     * Get routing table via data source name and actual table name.
     *
     * @param dataSourceName data source name
     * @param actualTableName actual table name
     * @return routing table
     */
    public Optional<RoutingTable> getRoutingTable(final String dataSourceName, final String actualTableName) {
        for (RoutingTable each : routingTables) {
            if (dataSourceName.equalsIgnoreCase(masterSlaveLogicDataSourceName) && each.getActualTableName().equalsIgnoreCase(actualTableName)) {
                return Optional.of(each);
            }
        }
        return Optional.absent();
    }
    
    /**
     * Get actual tables' names via data source name.
     *
     * @param dataSourceName data source name
     * @param logicTableName logic table name
     * @return  actual tables' names
     */
    public Set<String> getActualTableNames(final String dataSourceName, final String logicTableName) {
        Set<String> result = new HashSet<>(routingTables.size(), 1);
        for (RoutingTable each : routingTables) {
            if (dataSourceName.equalsIgnoreCase(this.dataSourceName) && each.getLogicTableName().equalsIgnoreCase(logicTableName)) {
                result.add(each.getActualTableName());
            }
        }
        return result;
    }
    
    /**
     * Get logic tables' names via data source name.
     *
     * @param dataSourceName data source name
     * @return  logic tables' names
     */
    public Set<String> getLogicTableNames(final String dataSourceName) {
        Set<String> result = new HashSet<>(routingTables.size(), 1);
        for (RoutingTable each : routingTables) {
            if (dataSourceName.equalsIgnoreCase(this.dataSourceName)) {
                result.add(each.getLogicTableName());
            }
        }
        return result;
    }
}

@RequiredArgsConstructor
@Getter
@EqualsAndHashCode
@ToString
public final class RoutingTable {
    
    private final String logicTableName;
    
    private final String actualTableName;
}