Skip to content

Commit 8a86467

Browse files
rajatbhattarajatbhatta
authored
feat: add tests for DML with Returning clause (#936)
This PR adds tests for running DML with Returning clause using the JDBC driver, and incorporates the following: - Integration tests for running DML statements with Returning clause using PreparedStatement. - Unit tests for running DML statements with Returning clause using JdbcStatement, for each of the available JDBC APIs `execute`, `executeUpdate`, `executeQuery`, `executeBatchUpdate`. - The JDBC driver does not require any code changes for supporting DML with Returning clause, as all the required changes will be made in the Connection API (Connection API changes are being tracked at https://togithub.com/googleapis/java-spanner/pull/1978).
1 parent 9016c38 commit 8a86467

File tree

3 files changed

+237
-2
lines changed

3 files changed

+237
-2
lines changed

‎src/main/java/com/google/cloud/spanner/jdbc/JdbcStatement.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ public long executeLargeUpdate(String sql) throws SQLException {
8787
switch (result.getResultType()) {
8888
case RESULT_SET:
8989
throw JdbcSqlExceptionFactory.of(
90-
"The statement is not an update or DDL statement", Code.INVALID_ARGUMENT);
90+
"The statement is not a non-returning DML or DDL statement", Code.INVALID_ARGUMENT);
9191
case UPDATE_COUNT:
9292
return result.getUpdateCount();
9393
case NO_RESULT:

‎src/test/java/com/google/cloud/spanner/jdbc/JdbcStatementTest.java

Lines changed: 92 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,12 @@
1717
package com.google.cloud.spanner.jdbc;
1818

1919
import static com.google.common.truth.Truth.assertThat;
20+
import static org.junit.Assert.assertArrayEquals;
21+
import static org.junit.Assert.assertEquals;
22+
import static org.junit.Assert.assertFalse;
23+
import static org.junit.Assert.assertNotNull;
24+
import static org.junit.Assert.assertThrows;
25+
import static org.junit.Assert.assertTrue;
2026
import static org.junit.Assert.fail;
2127
import static org.mockito.Mockito.anyList;
2228
import static org.mockito.Mockito.mock;
@@ -53,6 +59,8 @@ public class JdbcStatementTest {
5359
private static final String SELECT = "SELECT 1";
5460
private static final String UPDATE = "UPDATE FOO SET BAR=1 WHERE BAZ=2";
5561
private static final String LARGE_UPDATE = "UPDATE FOO SET BAR=1 WHERE 1=1";
62+
private static final String DML_RETURNING_GSQL = "UPDATE FOO SET BAR=1 WHERE 1=1 THEN RETURN *";
63+
private static final String DML_RETURNING_PG = "UPDATE FOO SET BAR=1 WHERE 1=1 RETURNING *";
5664
private static final String DDL = "CREATE INDEX FOO ON BAR(ID)";
5765

5866
@Parameter public Dialect dialect;
@@ -62,11 +70,20 @@ public static Object[] data() {
6270
return Dialect.values();
6371
}
6472

73+
private String getDmlReturningSql() {
74+
if (dialect == Dialect.GOOGLE_STANDARD_SQL) {
75+
return DML_RETURNING_GSQL;
76+
}
77+
return DML_RETURNING_PG;
78+
}
79+
6580
@SuppressWarnings("unchecked")
6681
private JdbcStatement createStatement() throws SQLException {
6782
Connection spanner = mock(Connection.class);
6883
when(spanner.getDialect()).thenReturn(dialect);
6984

85+
final String DML_RETURNING_SQL = getDmlReturningSql();
86+
7087
com.google.cloud.spanner.ResultSet resultSet = mock(com.google.cloud.spanner.ResultSet.class);
7188
when(resultSet.next()).thenReturn(true, false);
7289
when(resultSet.getColumnType(0)).thenReturn(Type.int64());
@@ -88,6 +105,19 @@ private JdbcStatement createStatement() throws SQLException {
88105
when(spanner.execute(com.google.cloud.spanner.Statement.of(LARGE_UPDATE)))
89106
.thenReturn(largeUpdateResult);
90107

108+
com.google.cloud.spanner.ResultSet dmlReturningResultSet =
109+
mock(com.google.cloud.spanner.ResultSet.class);
110+
when(dmlReturningResultSet.next()).thenReturn(true, false);
111+
when(dmlReturningResultSet.getColumnCount()).thenReturn(1);
112+
when(dmlReturningResultSet.getColumnType(0)).thenReturn(Type.int64());
113+
when(dmlReturningResultSet.getLong(0)).thenReturn(1L);
114+
115+
StatementResult dmlReturningResult = mock(StatementResult.class);
116+
when(dmlReturningResult.getResultType()).thenReturn(ResultType.RESULT_SET);
117+
when(dmlReturningResult.getResultSet()).thenReturn(dmlReturningResultSet);
118+
when(spanner.execute(com.google.cloud.spanner.Statement.of(DML_RETURNING_SQL)))
119+
.thenReturn(dmlReturningResult);
120+
91121
StatementResult ddlResult = mock(StatementResult.class);
92122
when(ddlResult.getResultType()).thenReturn(ResultType.NO_RESULT);
93123
when(spanner.execute(com.google.cloud.spanner.Statement.of(DDL))).thenReturn(ddlResult);
@@ -96,6 +126,8 @@ private JdbcStatement createStatement() throws SQLException {
96126
when(spanner.executeQuery(com.google.cloud.spanner.Statement.of(UPDATE)))
97127
.thenThrow(
98128
SpannerExceptionFactory.newSpannerException(ErrorCode.INVALID_ARGUMENT, "not a query"));
129+
when(spanner.executeQuery(com.google.cloud.spanner.Statement.of(DML_RETURNING_SQL)))
130+
.thenReturn(dmlReturningResultSet);
99131
when(spanner.executeQuery(com.google.cloud.spanner.Statement.of(DDL)))
100132
.thenThrow(
101133
SpannerExceptionFactory.newSpannerException(ErrorCode.INVALID_ARGUMENT, "not a query"));
@@ -109,6 +141,10 @@ private JdbcStatement createStatement() throws SQLException {
109141
.thenThrow(
110142
SpannerExceptionFactory.newSpannerException(
111143
ErrorCode.INVALID_ARGUMENT, "not an update"));
144+
when(spanner.executeUpdate(com.google.cloud.spanner.Statement.of(DML_RETURNING_SQL)))
145+
.thenThrow(
146+
SpannerExceptionFactory.newSpannerException(
147+
ErrorCode.FAILED_PRECONDITION, "cannot execute dml returning over executeUpdate"));
112148

113149
when(spanner.executeBatchUpdate(anyList()))
114150
.thenAnswer(
@@ -219,6 +255,20 @@ public void testExecuteWithDdlStatement() throws SQLException {
219255
assertThat(statement.getUpdateCount()).isEqualTo(JdbcConstants.STATEMENT_NO_RESULT);
220256
}
221257

258+
@Test
259+
public void testExecuteWithDmlReturningStatement() throws SQLException {
260+
Statement statement = createStatement();
261+
boolean res = statement.execute(getDmlReturningSql());
262+
assertTrue(res);
263+
assertEquals(statement.getUpdateCount(), JdbcConstants.STATEMENT_RESULT_SET);
264+
try (ResultSet rs = statement.getResultSet()) {
265+
assertNotNull(rs);
266+
assertTrue(rs.next());
267+
assertEquals(rs.getLong(1), 1L);
268+
assertFalse(rs.next());
269+
}
270+
}
271+
222272
@Test
223273
public void testExecuteWithGeneratedKeys() throws SQLException {
224274
Statement statement = createStatement();
@@ -257,6 +307,17 @@ public void testExecuteQueryWithUpdateStatement() {
257307
}
258308
}
259309

310+
@Test
311+
public void testExecuteQueryWithDmlReturningStatement() throws SQLException {
312+
Statement statement = createStatement();
313+
try (ResultSet rs = statement.executeQuery(getDmlReturningSql())) {
314+
assertNotNull(rs);
315+
assertTrue(rs.next());
316+
assertEquals(rs.getLong(1), 1L);
317+
assertFalse(rs.next());
318+
}
319+
}
320+
260321
@Test
261322
public void testExecuteQueryWithDdlStatement() {
262323
try {
@@ -353,12 +414,29 @@ public void testExecuteUpdateWithSelectStatement() {
353414
} catch (SQLException e) {
354415
assertThat(
355416
JdbcExceptionMatcher.matchCodeAndMessage(
356-
Code.INVALID_ARGUMENT, "The statement is not an update or DDL statement")
417+
Code.INVALID_ARGUMENT,
418+
"The statement is not a non-returning DML or DDL statement")
357419
.matches(e))
358420
.isTrue();
359421
}
360422
}
361423

424+
@Test
425+
public void testExecuteUpdateWithDmlReturningStatement() {
426+
try {
427+
Statement statement = createStatement();
428+
SQLException e =
429+
assertThrows(SQLException.class, () -> statement.executeUpdate(getDmlReturningSql()));
430+
assertTrue(
431+
JdbcExceptionMatcher.matchCodeAndMessage(
432+
Code.INVALID_ARGUMENT,
433+
"The statement is not a non-returning DML or DDL statement")
434+
.matches(e));
435+
} catch (SQLException e) {
436+
// ignore exception.
437+
}
438+
}
439+
362440
@Test
363441
public void testExecuteUpdateWithDdlStatement() throws SQLException {
364442
Statement statement = createStatement();
@@ -438,6 +516,19 @@ public void testDmlBatch() throws SQLException {
438516
}
439517
}
440518

519+
@Test
520+
public void testDmlBatchWithDmlReturning() throws SQLException {
521+
try (Statement statement = createStatement()) {
522+
// Verify that multiple batches can be executed on the same statement.
523+
for (int i = 0; i < 2; i++) {
524+
statement.addBatch(getDmlReturningSql());
525+
statement.addBatch(getDmlReturningSql());
526+
statement.addBatch(getDmlReturningSql());
527+
assertArrayEquals(statement.executeBatch(), new int[] {1, 1, 1});
528+
}
529+
}
530+
}
531+
441532
@Test
442533
public void testLargeDmlBatch() throws SQLException {
443534
try (Statement statement = createStatement()) {

‎src/test/java/com/google/cloud/spanner/jdbc/it/ITJdbcPreparedStatementTest.java

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import com.google.cloud.spanner.testing.EmulatorSpannerHelper;
3737
import com.google.common.base.Strings;
3838
import com.google.common.io.BaseEncoding;
39+
import com.google.common.io.CharStreams;
3940
import java.io.IOException;
4041
import java.io.InputStream;
4142
import java.io.StringReader;
@@ -263,6 +264,24 @@ private void setPreparedStatement(Connection connection, PreparedStatement ps, D
263264
ps.setArray(6, connection.createArrayOf("INT64", this.ticketPrices));
264265
}
265266
}
267+
268+
private void assertEqualsFields(Connection connection, ResultSet rs, Dialect dialect)
269+
throws SQLException {
270+
assertEquals(rs.getLong(1), this.venueId);
271+
assertEquals(rs.getLong(2), this.singerId);
272+
if (dialect == Dialect.POSTGRESQL) {
273+
assertEquals(rs.getString(3), this.concertDate.toString());
274+
assertEquals(rs.getString(4), this.beginTime.toString());
275+
assertEquals(rs.getString(5), this.endTime.toString());
276+
} else {
277+
assertEquals(rs.getDate(3), this.concertDate);
278+
assertEquals(rs.getTimestamp(4), this.beginTime);
279+
assertEquals(rs.getTimestamp(5), this.endTime);
280+
assertArrayEquals(
281+
(Object[]) rs.getArray(6).getArray(),
282+
(Object[]) connection.createArrayOf("INT64", this.ticketPrices).getArray());
283+
}
284+
}
266285
}
267286

268287
private static Date parseDate(String value) {
@@ -333,6 +352,34 @@ private String getConcertsInsertQuery(Dialect dialect) {
333352
return "INSERT INTO Concerts (VenueId, SingerId, ConcertDate, BeginTime, EndTime, TicketPrices) VALUES (?,?,?,?,?,?);";
334353
}
335354

355+
private String getConcertsInsertReturningQuery(Dialect dialect) {
356+
if (dialect == Dialect.POSTGRESQL) {
357+
return "INSERT INTO Concerts (VenueId, SingerId, ConcertDate, BeginTime, EndTime) VALUES (?,?,?,?,?) RETURNING *;";
358+
}
359+
return "INSERT INTO Concerts (VenueId, SingerId, ConcertDate, BeginTime, EndTime, TicketPrices) VALUES (?,?,?,?,?,?) THEN RETURN *;";
360+
}
361+
362+
private String getSingersInsertReturningQuery(Dialect dialect) {
363+
if (dialect == Dialect.POSTGRESQL) {
364+
return "INSERT INTO Singers (SingerId, FirstName, LastName, SingerInfo, BirthDate) values (?,?,?,?,?) RETURNING *";
365+
}
366+
return "INSERT INTO Singers (SingerId, FirstName, LastName, SingerInfo, BirthDate) values (?,?,?,?,?) THEN RETURN *";
367+
}
368+
369+
private String getAlbumsInsertReturningQuery(Dialect dialect) {
370+
if (dialect == Dialect.POSTGRESQL) {
371+
return "INSERT INTO Albums (SingerId, AlbumId, AlbumTitle, MarketingBudget) VALUES (?,?,?,?) RETURNING *";
372+
}
373+
return "INSERT INTO Albums (SingerId, AlbumId, AlbumTitle, MarketingBudget) VALUES (?,?,?,?) THEN RETURN *";
374+
}
375+
376+
private String getSongsInsertReturningQuery(Dialect dialect) {
377+
if (dialect == Dialect.POSTGRESQL) {
378+
return "INSERT INTO Songs (SingerId, AlbumId, TrackId, SongName, Duration, SongGenre) VALUES (?,?,?,?,?,?) RETURNING *;";
379+
}
380+
return "INSERT INTO Songs (SingerId, AlbumId, TrackId, SongName, Duration, SongGenre) VALUES (?,?,?,?,?,?) THEN RETURN *;";
381+
}
382+
336383
private int getConcertExpectedParamCount(Dialect dialect) {
337384
if (dialect == Dialect.POSTGRESQL) {
338385
return 5;
@@ -1150,6 +1197,103 @@ private void assertDefaultParameterMetaData(ParameterMetaData pmd, int expectedP
11501197
}
11511198
}
11521199

1200+
@Test
1201+
public void test12_InsertReturningTestData() throws SQLException {
1202+
assumeFalse(
1203+
"Emulator does not support DML with returning clause",
1204+
EmulatorSpannerHelper.isUsingEmulator());
1205+
try (Connection connection = createConnection(env, database)) {
1206+
connection.setAutoCommit(false);
1207+
// Delete existing rows from tables populated by other tests,
1208+
// so that this test can populate rows from scratch.
1209+
Statement deleteStatements = connection.createStatement();
1210+
deleteStatements.addBatch("DELETE FROM Concerts WHERE TRUE");
1211+
deleteStatements.addBatch("DELETE FROM Songs WHERE TRUE");
1212+
deleteStatements.addBatch("DELETE FROM Albums WHERE TRUE");
1213+
deleteStatements.addBatch("DELETE FROM Singers WHERE TRUE");
1214+
deleteStatements.executeBatch();
1215+
try (PreparedStatement ps =
1216+
connection.prepareStatement(getSingersInsertReturningQuery(dialect.dialect))) {
1217+
assertDefaultParameterMetaData(ps.getParameterMetaData(), 5);
1218+
for (Singer singer : createSingers()) {
1219+
singer.setPreparedStatement(ps, getDialect());
1220+
assertInsertSingerParameterMetadata(ps.getParameterMetaData());
1221+
ps.addBatch();
1222+
// check that adding the current params to a batch will not reset the metadata
1223+
assertInsertSingerParameterMetadata(ps.getParameterMetaData());
1224+
}
1225+
int[] results = ps.executeBatch();
1226+
for (int res : results) {
1227+
assertEquals(1, res);
1228+
}
1229+
}
1230+
try (PreparedStatement ps =
1231+
connection.prepareStatement(getAlbumsInsertReturningQuery(dialect.dialect))) {
1232+
assertDefaultParameterMetaData(ps.getParameterMetaData(), 4);
1233+
for (Album album : createAlbums()) {
1234+
ps.setLong(1, album.singerId);
1235+
ps.setLong(2, album.albumId);
1236+
ps.setString(3, album.albumTitle);
1237+
ps.setLong(4, album.marketingBudget);
1238+
assertInsertAlbumParameterMetadata(ps.getParameterMetaData());
1239+
try (ResultSet rs = ps.executeQuery()) {
1240+
rs.next();
1241+
assertEquals(rs.getLong(1), album.singerId);
1242+
assertEquals(rs.getLong(2), album.albumId);
1243+
assertEquals(rs.getString(3), album.albumTitle);
1244+
assertEquals(rs.getLong(4), album.marketingBudget);
1245+
}
1246+
// check that calling executeQuery will not reset the metadata
1247+
assertInsertAlbumParameterMetadata(ps.getParameterMetaData());
1248+
}
1249+
}
1250+
try (PreparedStatement ps =
1251+
connection.prepareStatement(getSongsInsertReturningQuery(dialect.dialect))) {
1252+
assertDefaultParameterMetaData(ps.getParameterMetaData(), 6);
1253+
for (Song song : createSongs()) {
1254+
ps.setByte(1, (byte) song.singerId);
1255+
ps.setInt(2, (int) song.albumId);
1256+
ps.setShort(3, (short) song.songId);
1257+
ps.setNString(4, song.songName);
1258+
ps.setLong(5, song.duration);
1259+
ps.setCharacterStream(6, new StringReader(song.songGenre));
1260+
assertInsertSongParameterMetadata(ps.getParameterMetaData());
1261+
try (ResultSet rs = ps.executeQuery()) {
1262+
rs.next();
1263+
assertEquals(rs.getByte(1), (byte) song.singerId);
1264+
assertEquals(rs.getInt(2), (int) song.albumId);
1265+
assertEquals(rs.getShort(3), (short) song.songId);
1266+
assertEquals(rs.getNString(4), song.songName);
1267+
assertEquals(rs.getLong(5), song.duration);
1268+
assertEquals(
1269+
CharStreams.toString(rs.getCharacterStream(6)),
1270+
CharStreams.toString(new StringReader(song.songGenre)));
1271+
}
1272+
// check that calling executeQuery will not reset the metadata
1273+
assertInsertSongParameterMetadata(ps.getParameterMetaData());
1274+
}
1275+
} catch (IOException e) {
1276+
// ignore exception.
1277+
}
1278+
try (PreparedStatement ps =
1279+
connection.prepareStatement(getConcertsInsertReturningQuery(dialect.dialect))) {
1280+
assertDefaultParameterMetaData(
1281+
ps.getParameterMetaData(), getConcertExpectedParamCount(dialect.dialect));
1282+
for (Concert concert : createConcerts()) {
1283+
concert.setPreparedStatement(connection, ps, getDialect());
1284+
assertInsertConcertParameterMetadata(ps.getParameterMetaData());
1285+
try (ResultSet rs = ps.executeQuery()) {
1286+
rs.next();
1287+
concert.assertEqualsFields(connection, rs, dialect.dialect);
1288+
}
1289+
// check that calling executeQuery will not reset the meta data
1290+
assertInsertConcertParameterMetadata(ps.getParameterMetaData());
1291+
}
1292+
}
1293+
connection.commit();
1294+
}
1295+
}
1296+
11531297
private List<String> readValuesFromFile(String filename) {
11541298
StringBuilder builder = new StringBuilder();
11551299
try (InputStream stream = ITJdbcPreparedStatementTest.class.getResourceAsStream(filename)) {

0 commit comments

Comments
 (0)