Skip to content

Commit 12ac28b

Browse files
committed
Updated tests for generated keys
1 parent f82b438 commit 12ac28b

File tree

4 files changed

+265
-92
lines changed

4 files changed

+265
-92
lines changed

jdbc/src/main/java/tech/ydb/jdbc/impl/BaseYdbStatement.java

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -153,8 +153,7 @@ protected boolean updateState(YdbQueryResult result) throws SQLException {
153153
return state.hasResultSets();
154154
}
155155

156-
protected YdbQueryResult executeBulkUpsert(YdbQuery query, String tablePath, ListValue rows)
157-
throws SQLException {
156+
protected YdbQueryResult executeBulkUpsert(YdbQuery query, String tablePath, ListValue rows) throws SQLException {
158157
ensureOpened();
159158

160159
if (connection.getExecutor().isInsideTransaction()) {

jdbc/src/main/java/tech/ydb/jdbc/impl/YdbPreparedStatementImpl.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -93,10 +93,12 @@ public int[] executeBatch() throws SQLException {
9393
try {
9494
if (query.getType() == QueryType.BULK_QUERY && (prepared instanceof BulkUpsertQuery)) {
9595
BulkUpsertQuery bulk = (BulkUpsertQuery) prepared;
96-
executeBulkUpsert(query, bulk.getTablePath(), bulk.getBatchedBulk());
96+
YdbQueryResult newState = executeBulkUpsert(query, bulk.getTablePath(), bulk.getBatchedBulk());
97+
updateState(newState);
9798
} else {
9899
for (Params prm: prepared.getBatchParams()) {
99-
executeDataQuery(query, prepared.getBatchText(prm), prm);
100+
YdbQueryResult newState = executeDataQuery(query, prepared.getBatchText(prm), prm);
101+
updateState(newState);
100102
}
101103
}
102104
} finally {

jdbc/src/test/java/tech/ydb/jdbc/YdbSerialTypesTest.java

Lines changed: 0 additions & 88 deletions
This file was deleted.
Lines changed: 260 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,260 @@
1+
package tech.ydb.jdbc.impl;
2+
3+
4+
import java.sql.Connection;
5+
import java.sql.DriverManager;
6+
import java.sql.PreparedStatement;
7+
import java.sql.ResultSet;
8+
import java.sql.SQLException;
9+
import java.sql.Statement;
10+
11+
import org.junit.jupiter.api.Assertions;
12+
import org.junit.jupiter.api.BeforeAll;
13+
import org.junit.jupiter.api.BeforeEach;
14+
import org.junit.jupiter.api.Disabled;
15+
import org.junit.jupiter.api.Test;
16+
import org.junit.jupiter.api.extension.RegisterExtension;
17+
18+
import tech.ydb.jdbc.impl.helper.JdbcUrlHelper;
19+
import tech.ydb.jdbc.impl.helper.TestTxTracer;
20+
import tech.ydb.test.junit5.YdbHelperExtension;
21+
22+
/**
23+
*
24+
* @author Aleksandr Gorshenin
25+
*/
26+
public class GeneratedKeysTest {
27+
@RegisterExtension
28+
private static final YdbHelperExtension ydb = new YdbHelperExtension();
29+
30+
private static final JdbcUrlHelper jdbcURL = new JdbcUrlHelper(ydb)
31+
.withArg("usePrefixPath", "serial_types")
32+
.withArg("enableTxTracer", "true");
33+
34+
@BeforeAll
35+
public static void prepareTable() throws SQLException {
36+
try (Connection connection = DriverManager.getConnection(jdbcURL.build())) {
37+
try (Statement st = connection.createStatement()) {
38+
try {
39+
st.execute("DROP TABLE serial_test");
40+
} catch (SQLException e) {
41+
// ignore
42+
}
43+
st.execute("CREATE TABLE serial_test(id Serial NOT NULL, value Text, PRIMARY KEY(id))");
44+
}
45+
}
46+
}
47+
48+
@BeforeEach
49+
public void cleanTable() throws SQLException {
50+
try (Connection connection = DriverManager.getConnection(jdbcURL.build())) {
51+
try (Statement st = connection.createStatement()) {
52+
st.execute("DELETE FROM serial_test");
53+
}
54+
}
55+
}
56+
57+
@Test
58+
public void inMemoryQueriesTest() throws SQLException {
59+
try (Connection connection = DriverManager.getConnection(jdbcURL.build())) {
60+
TestTxTracer tracer = YdbTracerImpl.use(new TestTxTracer());
61+
62+
try (PreparedStatement ps = connection.prepareStatement(
63+
"$i=select 1; INSERT INTO serial_test(value) VALUES (?)", Statement.RETURN_GENERATED_KEYS)) {
64+
ps.setString(1, "first value");
65+
66+
Assertions.assertEquals(1, ps.executeUpdate());
67+
tracer.assertQueriesCount(1);
68+
tracer.assertLastQueryContains("RETURNING *");
69+
70+
try (ResultSet rs = ps.getGeneratedKeys()) {
71+
Assertions.assertEquals(2, rs.getMetaData().getColumnCount());
72+
Assertions.assertTrue(rs.next());
73+
Assertions.assertTrue(rs.getInt("id") != 0);
74+
Assertions.assertEquals("first value", rs.getString("value"));
75+
76+
Assertions.assertFalse(rs.next());
77+
}
78+
}
79+
80+
try (PreparedStatement ps = connection.prepareStatement(
81+
"$i=select 1; INSERT INTO serial_test(value) VALUES (?)", new String[] { "id" })) {
82+
ps.setString(1, "first value");
83+
84+
Assertions.assertEquals(1, ps.executeUpdate());
85+
tracer.assertQueriesCount(1);
86+
tracer.assertLastQueryContains("RETURNING `id`");
87+
88+
try (ResultSet rs = ps.getGeneratedKeys()) {
89+
Assertions.assertEquals(1, rs.getMetaData().getColumnCount());
90+
91+
Assertions.assertTrue(rs.next());
92+
Assertions.assertTrue(rs.getInt("id") != 0);
93+
94+
Assertions.assertFalse(rs.next());
95+
}
96+
}
97+
98+
try (PreparedStatement ps = connection.prepareStatement(
99+
"$i=select 1; UPDATE serial_test SET value = ? WHERE value = ?", new String[] { "id", "value" })) {
100+
ps.setString(1, "second value");
101+
ps.setString(2, "first value");
102+
103+
Assertions.assertEquals(1, ps.executeUpdate());
104+
tracer.assertQueriesCount(1);
105+
tracer.assertLastQueryContains("RETURNING `id`, `value`");
106+
107+
try (ResultSet rs = ps.getGeneratedKeys()) {
108+
Assertions.assertEquals(2, rs.getMetaData().getColumnCount());
109+
110+
Assertions.assertTrue(rs.next());
111+
Assertions.assertTrue(rs.getInt("id") != 0);
112+
Assertions.assertEquals("second value", rs.getString("value"));
113+
114+
Assertions.assertTrue(rs.next());
115+
Assertions.assertTrue(rs.getInt("id") != 0);
116+
Assertions.assertEquals("second value", rs.getString("value"));
117+
118+
Assertions.assertFalse(rs.next());
119+
}
120+
}
121+
122+
try (PreparedStatement ps = connection.prepareStatement(
123+
"$i=select 1; DELETE FROM serial_test", Statement.RETURN_GENERATED_KEYS)) {
124+
125+
Assertions.assertEquals(1, ps.executeUpdate());
126+
tracer.assertQueriesCount(1);
127+
tracer.assertLastQueryContains("RETURNING *");
128+
129+
try (ResultSet rs = ps.getGeneratedKeys()) {
130+
Assertions.assertEquals(2, rs.getMetaData().getColumnCount());
131+
132+
Assertions.assertTrue(rs.next());
133+
Assertions.assertTrue(rs.getInt("id") != 0);
134+
Assertions.assertEquals("second value", rs.getString("value"));
135+
136+
Assertions.assertTrue(rs.next());
137+
Assertions.assertTrue(rs.getInt("id") != 0);
138+
Assertions.assertEquals("second value", rs.getString("value"));
139+
140+
Assertions.assertFalse(rs.next());
141+
}
142+
}
143+
}
144+
}
145+
146+
@Test
147+
@Disabled
148+
public void inMemoryBatchQueriesTest() throws SQLException {
149+
try (Connection connection = DriverManager.getConnection(jdbcURL.build())) {
150+
TestTxTracer tracer = YdbTracerImpl.use(new TestTxTracer());
151+
152+
try (PreparedStatement ps = connection.prepareStatement(
153+
"$i=select 1; INSERT INTO serial_test(value) VALUES (?)", Statement.RETURN_GENERATED_KEYS)) {
154+
ps.setString(1, "first value");
155+
ps.addBatch();
156+
157+
ps.setString(1, "second value");
158+
ps.addBatch();
159+
160+
Assertions.assertEquals(2, ps.executeBatch().length);
161+
162+
tracer.assertQueriesCount(2);
163+
tracer.assertLastQueryContains("RETURNING *");
164+
165+
try (ResultSet rs = ps.getGeneratedKeys()) {
166+
Assertions.assertEquals(2, rs.getMetaData().getColumnCount());
167+
Assertions.assertTrue(rs.next());
168+
Assertions.assertTrue(rs.getInt("id") != 0);
169+
Assertions.assertEquals("first value", rs.getString("value"));
170+
171+
Assertions.assertTrue(rs.next());
172+
Assertions.assertTrue(rs.getInt("id") != 0);
173+
Assertions.assertEquals("second value", rs.getString("value"));
174+
175+
Assertions.assertFalse(rs.next());
176+
}
177+
}
178+
}
179+
}
180+
181+
@Test
182+
public void standardQueriesTest() throws SQLException {
183+
try (Connection connection = DriverManager.getConnection(jdbcURL.build())) {
184+
TestTxTracer tracer = YdbTracerImpl.use(new TestTxTracer());
185+
186+
try (PreparedStatement ps = connection.prepareStatement(
187+
"INSERT INTO serial_test(value) VALUES (?)", Statement.RETURN_GENERATED_KEYS)) {
188+
ps.setString(1, "first value");
189+
ps.addBatch();
190+
ps.setString(1, "second value");
191+
ps.addBatch();
192+
ps.setString(1, "third value");
193+
ps.addBatch();
194+
195+
Assertions.assertEquals(3, ps.executeBatch().length);
196+
tracer.assertQueriesCount(1);
197+
tracer.assertLastQueryContains("FROM AS_TABLE($batch)");
198+
tracer.assertLastQueryContains("RETURNING *");
199+
200+
try (ResultSet rs = ps.getGeneratedKeys()) {
201+
Assertions.assertEquals(2, rs.getMetaData().getColumnCount());
202+
Assertions.assertTrue(rs.next());
203+
Assertions.assertTrue(rs.getInt("id") != 0);
204+
Assertions.assertEquals("first value", rs.getString("value"));
205+
Assertions.assertTrue(rs.next());
206+
Assertions.assertTrue(rs.getInt("id") != 0);
207+
Assertions.assertEquals("second value", rs.getString("value"));
208+
Assertions.assertTrue(rs.next());
209+
Assertions.assertTrue(rs.getInt("id") != 0);
210+
Assertions.assertEquals("third value", rs.getString("value"));
211+
212+
Assertions.assertFalse(rs.next());
213+
}
214+
}
215+
216+
try (PreparedStatement ps = connection.prepareStatement(
217+
"UPDATE serial_test SET value = ? WHERE value = ?", new String[] { "id", "value" })) {
218+
ps.setString(1, "second value");
219+
ps.setString(2, "first value");
220+
221+
Assertions.assertEquals(1, ps.executeUpdate());
222+
tracer.assertQueriesCount(1);
223+
tracer.assertLastQueryContains("RETURNING `id`, `value`");
224+
225+
try (ResultSet rs = ps.getGeneratedKeys()) {
226+
Assertions.assertEquals(2, rs.getMetaData().getColumnCount());
227+
228+
Assertions.assertTrue(rs.next());
229+
Assertions.assertTrue(rs.getInt("id") != 0);
230+
Assertions.assertEquals("second value", rs.getString("value"));
231+
232+
Assertions.assertFalse(rs.next());
233+
}
234+
}
235+
236+
try (PreparedStatement ps = connection.prepareStatement(
237+
"$i=select 1; DELETE FROM serial_test WHERE value=?", Statement.RETURN_GENERATED_KEYS)) {
238+
ps.setString(1, "second value");
239+
240+
Assertions.assertEquals(1, ps.executeUpdate());
241+
tracer.assertQueriesCount(1);
242+
tracer.assertLastQueryContains("RETURNING *");
243+
244+
try (ResultSet rs = ps.getGeneratedKeys()) {
245+
Assertions.assertEquals(2, rs.getMetaData().getColumnCount());
246+
247+
Assertions.assertTrue(rs.next());
248+
Assertions.assertTrue(rs.getInt("id") != 0);
249+
Assertions.assertEquals("second value", rs.getString("value"));
250+
251+
Assertions.assertTrue(rs.next());
252+
Assertions.assertTrue(rs.getInt("id") != 0);
253+
Assertions.assertEquals("second value", rs.getString("value"));
254+
255+
Assertions.assertFalse(rs.next());
256+
}
257+
}
258+
}
259+
}
260+
}

0 commit comments

Comments
 (0)