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 try (InputStream is = getClass().getResourceAsStream("/examples/sharding/ShardingDB.sql")) {
58 assert is != null;
59 try (Connection connection = DriverManager.getConnection(JDBC_URL, "sa", "");
60 InputStreamReader isr = new InputStreamReader(is)) {
61 ScriptRunner sr = new ScriptRunner(connection);
62 sr.setLogWriter(null);
63 sr.runScript(isr);
64 }
65 }
66
67 UnpooledDataSource ds = new UnpooledDataSource(JDBC_DRIVER, JDBC_URL, "sa", "");
68 Environment environment = new Environment("test", new JdbcTransactionFactory(), ds);
69 Configuration config = new Configuration(environment);
70 config.addMapper(ShardedMapper.class);
71 sqlSessionFactory = new SqlSessionFactoryBuilder().build(config);
72 }
73
74 @Test
75 void testShardedSelect() {
76 try (SqlSession sqlSession = sqlSessionFactory.openSession()) {
77 ShardedMapper mapper = sqlSession.getMapper(ShardedMapper.class);
78 TableCodes table = calculateTable(1);
79
80 SelectStatementProvider selectStatement = select(table.description)
81 .from(table)
82 .where(table.id, isEqualTo(1))
83 .build()
84 .render(RenderingStrategies.MYBATIS3);
85
86 assertThat(selectStatement.getSelectStatement())
87 .isEqualTo("select description from tableCodes_odd where id = #{parameters.p1,jdbcType=INTEGER}");
88
89 String description = mapper.selectOneString(selectStatement);
90 assertThat(description).isNull();
91 }
92 }
93
94 @Test
95 void testShardedInserts() {
96 try (SqlSession sqlSession = sqlSessionFactory.openSession()) {
97 ShardedMapper mapper = sqlSession.getMapper(ShardedMapper.class);
98
99 mapper.generalInsert(buildInsert(1, "Description 1"));
100 mapper.generalInsert(buildInsert(2, "Description 2"));
101 mapper.generalInsert(buildInsert(3, "Description 3"));
102 mapper.generalInsert(buildInsert(4, "Description 4"));
103 mapper.generalInsert(buildInsert(5, "Description 5"));
104 mapper.generalInsert(buildInsert(6, "Description 6"));
105 mapper.generalInsert(buildInsert(7, "Description 7"));
106
107 TableCodes oddTable = calculateTable(1);
108 SelectStatementProvider oddCountStatement = countFrom(oddTable)
109 .build()
110 .render(RenderingStrategies.MYBATIS3);
111 assertThat(oddCountStatement.getSelectStatement()).isEqualTo("select count(*) from tableCodes_odd");
112 long oddRows = mapper.count(oddCountStatement);
113 assertThat(oddRows).isEqualTo(4L);
114
115 TableCodes evenTable = calculateTable(2);
116 SelectStatementProvider evenCountStatement = countFrom(evenTable)
117 .build()
118 .render(RenderingStrategies.MYBATIS3);
119 assertThat(evenCountStatement.getSelectStatement()).isEqualTo("select count(*) from tableCodes_even");
120 long evenRows = mapper.count(evenCountStatement);
121 assertThat(evenRows).isEqualTo(3L);
122 }
123 }
124
125 @Test
126 void testShardedSelects() {
127 try (SqlSession sqlSession = sqlSessionFactory.openSession()) {
128 ShardedMapper mapper = sqlSession.getMapper(ShardedMapper.class);
129
130 mapper.generalInsert(buildInsert(1, "Description 1"));
131 mapper.generalInsert(buildInsert(2, "Description 2"));
132
133 assertThat(mapper.selectOneString(buildSelect(1))).isEqualTo("Description 1");
134 assertThat(mapper.selectOneString(buildSelect(2))).isEqualTo("Description 2");
135 assertThat(mapper.selectOneString(buildSelect(3))).isNull();
136 assertThat(mapper.selectOneString(buildSelect(4))).isNull();
137 }
138 }
139
140 private GeneralInsertStatementProvider buildInsert(int id, String description) {
141 TableCodesDynamicSqlSupport.TableCodes table = calculateTable(id);
142 return insertInto(table)
143 .set(table.id).toValue(id)
144 .set(table.description).toValue(description)
145 .build()
146 .render(RenderingStrategies.MYBATIS3);
147 }
148
149 private SelectStatementProvider buildSelect(int id) {
150 TableCodesDynamicSqlSupport.TableCodes table = calculateTable(id);
151 return select(table.description)
152 .from(table)
153 .where(table.id, isEqualTo(id))
154 .build()
155 .render(RenderingStrategies.MYBATIS3);
156 }
157
158 private TableCodes calculateTable(int id) {
159 if (id % 2 == 0) {
160 return shards.computeIfAbsent("even", k -> tableCodes);
161 } else {
162 return shards.computeIfAbsent("odd", k -> tableCodes.withName("tableCodes_odd"));
163 }
164 }
165 }