1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16 package examples.sharding;
17
18 import static examples.sharding.TableCodesDynamicSqlSupport.tableCodes;
19 import static org.assertj.core.api.Assertions.assertThat;
20 import static org.mybatis.dynamic.sql.SqlBuilder.countFrom;
21 import static org.mybatis.dynamic.sql.SqlBuilder.insertInto;
22 import static org.mybatis.dynamic.sql.SqlBuilder.isEqualTo;
23 import static org.mybatis.dynamic.sql.SqlBuilder.select;
24
25 import java.io.InputStream;
26 import java.io.InputStreamReader;
27 import java.sql.Connection;
28 import java.sql.DriverManager;
29 import java.util.HashMap;
30 import java.util.Map;
31
32 import examples.sharding.TableCodesDynamicSqlSupport.TableCodes;
33 import org.apache.ibatis.datasource.unpooled.UnpooledDataSource;
34 import org.apache.ibatis.jdbc.ScriptRunner;
35 import org.apache.ibatis.mapping.Environment;
36 import org.apache.ibatis.session.Configuration;
37 import org.apache.ibatis.session.SqlSession;
38 import org.apache.ibatis.session.SqlSessionFactory;
39 import org.apache.ibatis.session.SqlSessionFactoryBuilder;
40 import org.apache.ibatis.transaction.jdbc.JdbcTransactionFactory;
41 import org.junit.jupiter.api.BeforeEach;
42 import org.junit.jupiter.api.Test;
43 import org.mybatis.dynamic.sql.insert.render.GeneralInsertStatementProvider;
44 import org.mybatis.dynamic.sql.render.RenderingStrategies;
45 import org.mybatis.dynamic.sql.select.render.SelectStatementProvider;
46
47 class ShardingTest {
48 private static final String JDBC_URL = "jdbc:hsqldb:mem:aname";
49 private static final String JDBC_DRIVER = "org.hsqldb.jdbcDriver";
50 private final Map<String, TableCodes> shards = new HashMap<>();
51
52 private SqlSessionFactory sqlSessionFactory;
53
54 @BeforeEach
55 void setup() throws Exception {
56 Class.forName(JDBC_DRIVER);
57 InputStream is = getClass().getResourceAsStream("/examples/sharding/ShardingDB.sql");
58 assert is != null;
59 try (Connection connection = DriverManager.getConnection(JDBC_URL, "sa", "")) {
60 ScriptRunner sr = new ScriptRunner(connection);
61 sr.setLogWriter(null);
62 sr.runScript(new InputStreamReader(is));
63 }
64
65 UnpooledDataSource ds = new UnpooledDataSource(JDBC_DRIVER, JDBC_URL, "sa", "");
66 Environment environment = new Environment("test", new JdbcTransactionFactory(), ds);
67 Configuration config = new Configuration(environment);
68 config.addMapper(ShardedMapper.class);
69 sqlSessionFactory = new SqlSessionFactoryBuilder().build(config);
70 }
71
72 @Test
73 void testShardedSelect() {
74 try (SqlSession sqlSession = sqlSessionFactory.openSession()) {
75 ShardedMapper mapper = sqlSession.getMapper(ShardedMapper.class);
76 TableCodes table = calculateTable(1);
77
78 SelectStatementProvider selectStatement = select(table.description)
79 .from(table)
80 .where(table.id, isEqualTo(1))
81 .build()
82 .render(RenderingStrategies.MYBATIS3);
83
84 assertThat(selectStatement.getSelectStatement())
85 .isEqualTo("select description from tableCodes_odd where id = #{parameters.p1,jdbcType=INTEGER}");
86
87 String description = mapper.selectOneString(selectStatement);
88 assertThat(description).isNull();
89 }
90 }
91
92 @Test
93 void testShardedInserts() {
94 try (SqlSession sqlSession = sqlSessionFactory.openSession()) {
95 ShardedMapper mapper = sqlSession.getMapper(ShardedMapper.class);
96
97 mapper.generalInsert(buildInsert(1, "Description 1"));
98 mapper.generalInsert(buildInsert(2, "Description 2"));
99 mapper.generalInsert(buildInsert(3, "Description 3"));
100 mapper.generalInsert(buildInsert(4, "Description 4"));
101 mapper.generalInsert(buildInsert(5, "Description 5"));
102 mapper.generalInsert(buildInsert(6, "Description 6"));
103 mapper.generalInsert(buildInsert(7, "Description 7"));
104
105 TableCodes oddTable = calculateTable(1);
106 SelectStatementProvider oddCountStatement = countFrom(oddTable)
107 .build()
108 .render(RenderingStrategies.MYBATIS3);
109 assertThat(oddCountStatement.getSelectStatement()).isEqualTo("select count(*) from tableCodes_odd");
110 long oddRows = mapper.count(oddCountStatement);
111 assertThat(oddRows).isEqualTo(4L);
112
113 TableCodes evenTable = calculateTable(2);
114 SelectStatementProvider evenCountStatement = countFrom(evenTable)
115 .build()
116 .render(RenderingStrategies.MYBATIS3);
117 assertThat(evenCountStatement.getSelectStatement()).isEqualTo("select count(*) from tableCodes_even");
118 long evenRows = mapper.count(evenCountStatement);
119 assertThat(evenRows).isEqualTo(3L);
120 }
121 }
122
123 @Test
124 void testShardedSelects() {
125 try (SqlSession sqlSession = sqlSessionFactory.openSession()) {
126 ShardedMapper mapper = sqlSession.getMapper(ShardedMapper.class);
127
128 mapper.generalInsert(buildInsert(1, "Description 1"));
129 mapper.generalInsert(buildInsert(2, "Description 2"));
130
131 assertThat(mapper.selectOneString(buildSelect(1))).isEqualTo("Description 1");
132 assertThat(mapper.selectOneString(buildSelect(2))).isEqualTo("Description 2");
133 assertThat(mapper.selectOneString(buildSelect(3))).isNull();
134 assertThat(mapper.selectOneString(buildSelect(4))).isNull();
135 }
136 }
137
138 private GeneralInsertStatementProvider buildInsert(int id, String description) {
139 TableCodesDynamicSqlSupport.TableCodes table = calculateTable(id);
140 return insertInto(table)
141 .set(table.id).toValue(id)
142 .set(table.description).toValue(description)
143 .build()
144 .render(RenderingStrategies.MYBATIS3);
145 }
146
147 private SelectStatementProvider buildSelect(int id) {
148 TableCodesDynamicSqlSupport.TableCodes table = calculateTable(id);
149 return select(table.description)
150 .from(table)
151 .where(table.id, isEqualTo(id))
152 .build()
153 .render(RenderingStrategies.MYBATIS3);
154 }
155
156 private TableCodes calculateTable(int id) {
157 if (id % 2 == 0) {
158 return shards.computeIfAbsent("even", k -> tableCodes);
159 } else {
160 return shards.computeIfAbsent("odd", k -> tableCodes.withName("tableCodes_odd"));
161 }
162 }
163 }