Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

get splitTable shardKey #92

Merged
merged 1 commit into from
Jul 20, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 62 additions & 40 deletions src/main/java/com/jd/jdbc/tindexes/SplitTableUtil.java
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,9 @@
import com.jd.jdbc.tindexes.config.LogicTableConfig;
import com.jd.jdbc.tindexes.config.SchemaConfig;
import com.jd.jdbc.tindexes.config.SplitTableConfig;
import com.jd.jdbc.vitess.VitessDriver;
import io.vitess.proto.Query;
import java.io.IOException;
import java.io.InputStream;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Iterator;
Expand All @@ -38,64 +39,85 @@
import org.yaml.snakeyaml.constructor.Constructor;
import org.yaml.snakeyaml.error.YAMLException;

public class SplitTableUtil {
public final class SplitTableUtil {

private static final String ACTUAL_TABLE_EXPR_REGEX = "\\$\\{\\d+\\.\\.\\d+}";

private static final Log logger = LogFactory.getLog(SplitTableUtil.class);

private static Map<String, Map<String, LogicTable>> tableIndexesMap;

/**
* 根据分表字段的值,获取分表算法计算后的分表名
*
* @param keyspace 数据库的keyspace
* @param logicTableName 逻辑表名
* @param value 分表字段的值
* @return
*/
private static String preConfigPath;

private SplitTableUtil() {
}

public static String getActualTableName(final String keyspace, final String logicTableName, final Object value) {
return getActualTableName(Constant.DEFAULT_SPLIT_TABLE_CONFIG_PATH, keyspace, logicTableName, value);
}

/**
* 根据分表字段的值,获取分表算法计算后的分表名
*
* @param configPath 分表配置文件路径
* @param keyspace 数据库的keyspace
* @param logicTableName 逻辑表名
* @param value 分表字段的值
* @return
*/
public static String getActualTableName(final String configPath, final String keyspace, final String logicTableName, final Object value) {
ActualTable actualTable = getActualTable(configPath, keyspace, logicTableName, value);
return actualTable.getActualTableName();
}

public static String getShardingColumnName(final String configPath, final String keyspace, final String logicTableName) {
LogicTable logicTable = getLogicTable(configPath, keyspace, logicTableName);
return logicTable.getTindexCol().getColumnName();
}

public static String getShardingColumnName(final String keyspace, final String logicTableName) {
return getShardingColumnName(Constant.DEFAULT_SPLIT_TABLE_CONFIG_PATH, keyspace, logicTableName);
}

private static LogicTable getLogicTable(String configPath, String keyspace, String logicTableName) {
if (StringUtils.isEmpty(keyspace) || StringUtils.isEmpty(logicTableName)) {
return null;
throw new RuntimeException("keyspace or logicTableName should not empty");
}
final Map<String, Map<String, LogicTable>> tableIndexesMap = initTableIndexesMapFromYaml(configPath);
Map<String, Map<String, LogicTable>> tableIndexesMap = getTableIndexesMap(configPath);
if (tableIndexesMap == null || tableIndexesMap.isEmpty()) {
return null;
throw new RuntimeException("cat not find split-table config through configPath=" + configPath);
}
String lowerCaseKeyspace = keyspace.toLowerCase();
String lowerCaseLogicTable = logicTableName.toLowerCase();
if (tableIndexesMap.containsKey(lowerCaseKeyspace) && tableIndexesMap.get(lowerCaseKeyspace).containsKey(lowerCaseLogicTable)) {
final LogicTable logicTable = tableIndexesMap.get(lowerCaseKeyspace).get(lowerCaseLogicTable);
VtValue vtValue;
try {
vtValue = VtValue.toVtValue(value);
} catch (SQLException e) {
throw new RuntimeException(e);
}
final ActualTable actualTable = logicTable.map(vtValue);
return actualTable.getActualTableName();
if (!tableIndexesMap.containsKey(lowerCaseKeyspace)) {
throw new RuntimeException("cat not find keyspace in split-table config, target keyspace=" + keyspace);
}
return null;
if (!tableIndexesMap.get(lowerCaseKeyspace).containsKey(lowerCaseLogicTable)) {
throw new RuntimeException("cat not find logicTable in split-table config, target keyspace=" + keyspace + " ,target logicTable=" + logicTableName);
}
return tableIndexesMap.get(lowerCaseKeyspace).get(lowerCaseLogicTable);
}

private static ActualTable getActualTable(String configPath, String keyspace, String logicTableName, Object value) {
LogicTable logicTable = getLogicTable(configPath, keyspace, logicTableName);
VtValue vtValue;
try {
vtValue = VtValue.toVtValue(value);
} catch (SQLException e) {
throw new RuntimeException(e);
}
final ActualTable actualTable = logicTable.map(vtValue);
if (actualTable == null) {
throw new RuntimeException("cannot calculate split table, logic table: " + logicTable.getLogicTable() + ",current value: " + vtValue);
}
return actualTable;
}

public static Map<String, Map<String, LogicTable>> getTableIndexesMap() {
if (SplitTableUtil.tableIndexesMap == null) {
SplitTableUtil.tableIndexesMap = initTableIndexesMapFromYaml(Constant.DEFAULT_SPLIT_TABLE_CONFIG_PATH);
return getTableIndexesMap(Constant.DEFAULT_SPLIT_TABLE_CONFIG_PATH);
}

public static synchronized Map<String, Map<String, LogicTable>> getTableIndexesMap(final String configPath) {
if (StringUtils.isEmpty(configPath)) {
throw new RuntimeException("configPath should not empty");
}
return SplitTableUtil.tableIndexesMap;
if (Objects.equals(configPath, preConfigPath)) {
return SplitTableUtil.tableIndexesMap;
}
SplitTableUtil.tableIndexesMap = initTableIndexesMapFromYaml(configPath);
preConfigPath = configPath;
return tableIndexesMap;
}

//call by vtdriver-spring-boot-starter
Expand All @@ -105,12 +127,12 @@ public static void setSplitIndexesMapFromSpring(SplitTableConfig splitTableConfi

public static Map<String, Map<String, LogicTable>> initTableIndexesMapFromYaml(final String configPath) {
Yaml yaml = new Yaml(new Constructor(SplitTableConfig.class));
try {
SplitTableConfig splitTableConfig = yaml.load(VitessDriver.class.getClassLoader().getResourceAsStream(configPath));
try (InputStream resourceAsStream = SplitTableUtil.class.getClassLoader().getResourceAsStream(configPath)) {
SplitTableConfig splitTableConfig = yaml.load(resourceAsStream);
return buildTableIndexesMap(splitTableConfig);
} catch (YAMLException e) {
} catch (YAMLException | IOException e) {
if (logger.isDebugEnabled()) {
logger.debug("Init split-table from vtdriver-split-table.yaml fail, caused by:" + e.getMessage());
logger.debug("Init split-table config through configPath=" + configPath + " fail, caused by:" + e.getMessage());
}
}
return null;
Expand Down
21 changes: 17 additions & 4 deletions src/main/java/com/jd/jdbc/vitess/VitessStatement.java
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,8 @@ public class VitessStatement extends AbstractVitessStatement {

private static final String FUNCATION_IDENTITY = "@@IDENTITY";

private static Query.Field[] generatedKeyField;

@Getter
protected final Executor executor;

Expand Down Expand Up @@ -582,18 +584,15 @@ protected VtResultSet getGeneratedKeysInternal(long numKeys) throws SQLException
throw new SQLException("does not support insert multiple rows in one sql statement");
}
long generatedKey = lastInsertId;
VtResultSet vtStaticResultSet = null;
Query.Field field = Query.Field.newBuilder().setName("GENERATED_KEY").setJdbcClassName("java.math.BigInteger").setType(Query.Type.UINT64).setColumnLength(20).setPrecision(20).build();
List<List<VtResultValue>> rows = new ArrayList<>();

VtResultSet vtStaticResultSet = new VtResultSet(getGeneratedKeyField(), rows);
if (!this.resultSets.isEmpty()) {
if (generatedKey < 0) {
throw new SQLException("generatedKey error");
}
if (generatedKey != 0 && (numKeys > 0)) {
List<VtResultValue> row = Collections.singletonList(VtResultValue.newVtResultValue(Query.Type.UINT64, BigInteger.valueOf(generatedKey)));
rows.add(row);
vtStaticResultSet = new VtResultSet(new Query.Field[] {field}, rows);
}
}
return vtStaticResultSet;
Expand Down Expand Up @@ -1145,6 +1144,20 @@ private void errorCount(String sql, Map<String, BindVariable> bindVariableMap, S
StatementCollector.getStatementErrorCounter().labels(connection.getDefaultKeyspace(), VitessJdbcProperyUtil.getRole(connection.getProperties())).inc();
}

private Query.Field[] getGeneratedKeyField() {
if (generatedKeyField == null) {
synchronized (VitessStatement.class) {
if (generatedKeyField == null) {
Query.Field field = Query.Field.newBuilder().setName("GENERATED_KEY").setJdbcClassName("java.math.BigInteger")
.setType(Query.Type.UINT64).setColumnLength(20).setPrecision(20).build();
generatedKeyField = new Query.Field[] {field};
return generatedKeyField;
}
}
}
return generatedKeyField;
}

@Getter
@AllArgsConstructor
public static class ParseResult {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,26 +14,40 @@
limitations under the License.
*/

package com.jd.jdbc.tindexes;
package com.jd.jdbc.table;

import com.google.common.collect.Lists;
import com.jd.jdbc.tindexes.SplitTableUtil;
import java.util.List;
import org.junit.Assert;
import org.junit.Test;

public class SplitTableUtilTest {

@Test
public void getActualTableName() {
public void getActualTableNames() {
List<Integer> ids = Lists.newArrayList(713, 714, 715, 716, 717, 718);
for (Integer id : ids) {
String actualTableName = SplitTableUtil.getActualTableName("vtdriver-split-table.yml", "vtdriver2", "table_engine_test", id);
String actualTableName = SplitTableUtil.getActualTableName("vtdriver-split-table.yml", "commerce", "table_engine_test", id);
System.out.println(String.format("id=%s,actualTableName=%s", id, actualTableName));
}
}

@Test
public void testGetActualTableName() {
String actualTableName = SplitTableUtil.getActualTableName("vtdriver2", "table_engine_test", 111);
String actualTableName = SplitTableUtil.getActualTableName("commerce", "table_engine_test", 111);
System.out.println(actualTableName);
}

@Test
public void getShardingColumnName() {
String shardingColumnName = SplitTableUtil.getShardingColumnName("commerce", "table_engine_test");
Assert.assertTrue("getShardingColumnName error", "f_key".equalsIgnoreCase(shardingColumnName));
}

@Test
public void getShardingColumnName2() {
String shardingColumnName = SplitTableUtil.getShardingColumnName("vtdriver-split-table.yml", "commerce", "table_engine_test");
Assert.assertTrue("getShardingColumnName error", "f_key".equalsIgnoreCase(shardingColumnName));
}
}
16 changes: 7 additions & 9 deletions src/test/java/com/jd/jdbc/table/TableAutoGeneratedKeysTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -25,27 +25,27 @@
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
import testsuite.internal.TestSuiteShardSpec;

public class TableAutoGeneratedKeysTest extends VitessAutoGeneratedKeysTest {

@Rule
public ExpectedException thrown = ExpectedException.none();

@Before
public void init() throws Exception {
@Override
public void testLoadDriver() throws Exception {
getConn();
TableTestUtil.setSplitTableConfig("table/autoGeneratedKeys.yml");
clean();

sql1 = "insert into table_auto (id,ai,email) values(1,1,'x')";
sql100 = "insert into table_auto (id,ai,email) values(1,100,'x')";
sqlx = "insert into table_auto (id,ai,email) values(?,?,'x')";
sqld = "insert into table_auto (id,ai,email) values(%d,%d,'x')";
updateSql = "update table_auto set email = 'zz' where id = %d";
deleteSql = "delete from table_auto where id = %d";
updateSql100 = "update table_auto set email = 'zz' where id = 1";
sql200 = "insert into table_auto (id,ai,email) values(200,200,'x')";
updatesql200 = "update table_auto set email = 'zz' where id = ?";
}

protected void getConn() throws SQLException {
Expand All @@ -55,9 +55,7 @@ protected void getConn() throws SQLException {

@After
public void close() throws Exception {
if (conn != null) {
conn.close();
}
closeConnection(conn);
TableTestUtil.setDefaultTableConfig();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,14 @@ public void test12StatementExecuteBatch() throws Exception {
super.test12StatementExecuteBatch();
}

@Test
@Override
public void test12StatementUpdateExecuteBatch() throws Exception {
thrown.expect(java.sql.SQLFeatureNotSupportedException.class);
thrown.expectMessage("unsupported multiquery");
super.test12StatementUpdateExecuteBatch();
}

@Test
@Override
public void test20PreparedStatementExecuteBatch() throws Exception {
Expand All @@ -70,8 +78,9 @@ public void test22PreparedStatementExecuteBatch() throws Exception {
}

@Test
@Override
public void test23SetNull() throws Exception {
super.test23SetNull();
public void test22PreparedStatementUpdateExecuteBatch() throws Exception {
thrown.expect(java.sql.SQLFeatureNotSupportedException.class);
thrown.expectMessage("unsupported multiquery");
super.test22PreparedStatementUpdateExecuteBatch();
}
}
Loading
Loading