Skip to content

Commit

Permalink
Add ShadowTableHintDataSourceMappingsRetriever
Browse files Browse the repository at this point in the history
  • Loading branch information
terrymanu committed Nov 5, 2024
1 parent 479d4e7 commit 090fd86
Show file tree
Hide file tree
Showing 2 changed files with 86 additions and 42 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -25,20 +25,17 @@
import org.apache.shardingsphere.shadow.condition.ShadowColumnCondition;
import org.apache.shardingsphere.shadow.condition.ShadowCondition;
import org.apache.shardingsphere.shadow.route.determiner.ColumnShadowAlgorithmDeterminer;
import org.apache.shardingsphere.shadow.route.determiner.HintShadowAlgorithmDeterminer;
import org.apache.shardingsphere.shadow.route.retriever.ShadowDataSourceMappingsRetriever;
import org.apache.shardingsphere.shadow.route.retriever.hint.ShadowTableHintDataSourceMappingsRetriever;
import org.apache.shardingsphere.shadow.rule.ShadowRule;
import org.apache.shardingsphere.shadow.spi.ShadowAlgorithm;
import org.apache.shardingsphere.shadow.spi.ShadowOperationType;
import org.apache.shardingsphere.shadow.spi.column.ColumnShadowAlgorithm;
import org.apache.shardingsphere.shadow.spi.hint.HintShadowAlgorithm;
import org.apache.shardingsphere.sql.parser.statement.core.segment.generic.table.SimpleTableSegment;

import java.util.Collection;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Optional;

/**
* Abstract shadow DML statement data source mappings retriever.
Expand All @@ -48,15 +45,15 @@ public abstract class AbstractShadowDMLStatementDataSourceMappingsRetriever impl

private final ShadowOperationType operationType;

private final boolean isShadow;

@Getter
private final Map<String, String> tableAliasAndNameMappings;

private final ShadowTableHintDataSourceMappingsRetriever tableHintDataSourceMappingsRetriever;

protected AbstractShadowDMLStatementDataSourceMappingsRetriever(final SQLStatementContext sqlStatementContext, final HintValueContext hintValueContext, final ShadowOperationType operationType) {
this.operationType = operationType;
isShadow = hintValueContext.isShadow();
tableAliasAndNameMappings = getTableAliasAndNameMappings(((TableAvailable) sqlStatementContext).getTablesContext().getSimpleTables());
tableHintDataSourceMappingsRetriever = new ShadowTableHintDataSourceMappingsRetriever(operationType, hintValueContext.isShadow(), tableAliasAndNameMappings);
}

private Map<String, String> getTableAliasAndNameMappings(final Collection<SimpleTableSegment> tableSegments) {
Expand All @@ -71,41 +68,8 @@ private Map<String, String> getTableAliasAndNameMappings(final Collection<Simple

@Override
public Map<String, String> retrieve(final ShadowRule rule) {
Collection<String> shadowTables = rule.filterShadowTables(tableAliasAndNameMappings.values());
if (shadowTables.isEmpty() && isMatchDefaultAlgorithm(rule)) {
return rule.getAllShadowDataSourceMappings();
}
Map<String, String> result = findBySQLHints(rule, shadowTables);
return result.isEmpty() ? findByShadowColumn(rule, shadowTables) : result;
}

@SuppressWarnings("unchecked")
private boolean isMatchDefaultAlgorithm(final ShadowRule rule) {
Optional<ShadowAlgorithm> defaultAlgorithm = rule.getDefaultShadowAlgorithm();
if (defaultAlgorithm.isPresent() && defaultAlgorithm.get() instanceof HintShadowAlgorithm<?>) {
return HintShadowAlgorithmDeterminer.isShadow((HintShadowAlgorithm<Comparable<?>>) defaultAlgorithm.get(), new ShadowCondition(), rule, isShadow);
}
return false;
}

private Map<String, String> findBySQLHints(final ShadowRule rule, final Collection<String> shadowTables) {
Map<String, String> result = new LinkedHashMap<>();
for (String each : shadowTables) {
if (isContainsShadowInSQLHints(rule, each, new ShadowCondition(each, operationType))) {
result.putAll(rule.getShadowDataSourceMappings(each));
return result;
}
}
return result;
}

private boolean isContainsShadowInSQLHints(final ShadowRule rule, final String tableName, final ShadowCondition shadowCondition) {
for (HintShadowAlgorithm<Comparable<?>> each : rule.getHintShadowAlgorithms(tableName)) {
if (HintShadowAlgorithmDeterminer.isShadow(each, shadowCondition, rule, isShadow)) {
return true;
}
}
return false;
Map<String, String> result = tableHintDataSourceMappingsRetriever.retrieve(rule);
return result.isEmpty() ? findByShadowColumn(rule, rule.filterShadowTables(tableAliasAndNameMappings.values())) : result;
}

private Map<String, String> findByShadowColumn(final ShadowRule rule, final Collection<String> shadowTables) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You 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 org.apache.shardingsphere.shadow.route.retriever.hint;

import lombok.RequiredArgsConstructor;
import org.apache.shardingsphere.infra.annotation.HighFrequencyInvocation;
import org.apache.shardingsphere.shadow.condition.ShadowCondition;
import org.apache.shardingsphere.shadow.route.determiner.HintShadowAlgorithmDeterminer;
import org.apache.shardingsphere.shadow.route.retriever.ShadowDataSourceMappingsRetriever;
import org.apache.shardingsphere.shadow.rule.ShadowRule;
import org.apache.shardingsphere.shadow.spi.ShadowAlgorithm;
import org.apache.shardingsphere.shadow.spi.ShadowOperationType;
import org.apache.shardingsphere.shadow.spi.hint.HintShadowAlgorithm;

import java.util.Collection;
import java.util.Collections;
import java.util.Map;
import java.util.Optional;

/**
* Shadow table hint data source mappings retriever.
*/
@HighFrequencyInvocation
@RequiredArgsConstructor
public final class ShadowTableHintDataSourceMappingsRetriever implements ShadowDataSourceMappingsRetriever {

private final ShadowOperationType operationType;

private final boolean isShadow;

private final Map<String, String> tableAliasAndNameMappings;

@Override
public Map<String, String> retrieve(final ShadowRule rule) {
Collection<String> shadowTables = rule.filterShadowTables(tableAliasAndNameMappings.values());
return shadowTables.isEmpty() && isMatchDefaultAlgorithm(rule) ? rule.getAllShadowDataSourceMappings() : findShadowDataSourceMappingsBySQLHints(rule, shadowTables);
}

@SuppressWarnings("unchecked")
private boolean isMatchDefaultAlgorithm(final ShadowRule rule) {
Optional<ShadowAlgorithm> defaultAlgorithm = rule.getDefaultShadowAlgorithm();
if (defaultAlgorithm.isPresent() && defaultAlgorithm.get() instanceof HintShadowAlgorithm<?>) {
return HintShadowAlgorithmDeterminer.isShadow((HintShadowAlgorithm<Comparable<?>>) defaultAlgorithm.get(), new ShadowCondition(), rule, isShadow);
}
return false;
}

private Map<String, String> findShadowDataSourceMappingsBySQLHints(final ShadowRule rule, final Collection<String> shadowTables) {
for (String each : shadowTables) {
if (containsShadowInSQLHints(rule, each, new ShadowCondition(each, operationType))) {
return rule.getShadowDataSourceMappings(each);
}
}
return Collections.emptyMap();
}

private boolean containsShadowInSQLHints(final ShadowRule rule, final String tableName, final ShadowCondition shadowCondition) {
for (HintShadowAlgorithm<Comparable<?>> each : rule.getHintShadowAlgorithms(tableName)) {
if (HintShadowAlgorithmDeterminer.isShadow(each, shadowCondition, rule, isShadow)) {
return true;
}
}
return false;
}
}

0 comments on commit 090fd86

Please sign in to comment.