Skip to content

Commit 1953ea2

Browse files
olavloiteolavloite
authored
chore: method for adding returning clause to statements (#1311)
* chore: method for adding returning clause to statements Adds a method to JdbcStatement for appending a THEN RETURN/RETURNING clause to the statement. This will be used to modify statements that request generated keys to be returned. * feat: support return all columns * fix: only add THEN RETURN * from DML
1 parent aaf89de commit 1953ea2

File tree

2 files changed

+237
-0
lines changed

2 files changed

+237
-0
lines changed

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

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,19 +23,25 @@
2323
import com.google.cloud.spanner.Statement;
2424
import com.google.cloud.spanner.Type;
2525
import com.google.cloud.spanner.Type.StructField;
26+
import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStatement;
2627
import com.google.cloud.spanner.connection.StatementResult;
2728
import com.google.common.annotations.VisibleForTesting;
2829
import com.google.common.base.Preconditions;
30+
import com.google.common.collect.ImmutableList;
2931
import com.google.rpc.Code;
3032
import java.sql.ResultSet;
3133
import java.sql.SQLException;
3234
import java.util.ArrayList;
3335
import java.util.Arrays;
3436
import java.util.Collections;
3537
import java.util.List;
38+
import java.util.stream.Collectors;
39+
import javax.annotation.Nullable;
3640

3741
/** Implementation of {@link java.sql.Statement} for Google Cloud Spanner. */
3842
class JdbcStatement extends AbstractJdbcStatement {
43+
static final ImmutableList<String> ALL_COLUMNS = ImmutableList.of("*");
44+
3945
enum BatchType {
4046
NONE,
4147
DML,
@@ -98,6 +104,81 @@ public long executeLargeUpdate(String sql) throws SQLException {
98104
}
99105
}
100106

107+
/**
108+
* Adds a THEN RETURN/RETURNING clause to the given statement if the following conditions are all
109+
* met:
110+
*
111+
* <ol>
112+
* <li>The generatedKeysColumns is not null or empty
113+
* <li>The statement is a DML statement
114+
* <li>The DML statement does not already contain a THEN RETURN/RETURNING clause
115+
* </ol>
116+
*/
117+
Statement addReturningToStatement(
118+
Statement statement, @Nullable ImmutableList<String> generatedKeysColumns)
119+
throws SQLException {
120+
if (generatedKeysColumns == null || generatedKeysColumns.isEmpty()) {
121+
return statement;
122+
}
123+
// Check if the statement is a DML statement or not.
124+
ParsedStatement parsedStatement = getConnection().getParser().parse(statement);
125+
if (parsedStatement.isUpdate() && !parsedStatement.hasReturningClause()) {
126+
if (generatedKeysColumns.size() == 1
127+
&& ALL_COLUMNS.get(0).equals(generatedKeysColumns.get(0))) {
128+
// Add a 'THEN RETURN/RETURNING *' clause to the statement.
129+
return statement
130+
.toBuilder()
131+
.replace(statement.getSql() + getReturningAllColumnsClause())
132+
.build();
133+
}
134+
// Add a 'THEN RETURN/RETURNING col1, col2, ...' to the statement.
135+
// The column names will be quoted using the dialect-specific identifier quoting character.
136+
return statement
137+
.toBuilder()
138+
.replace(
139+
generatedKeysColumns.stream()
140+
.map(this::quoteColumn)
141+
.collect(
142+
Collectors.joining(
143+
", ", statement.getSql() + getReturningClause() + " ", "")))
144+
.build();
145+
}
146+
return statement;
147+
}
148+
149+
/** Returns the dialect-specific clause for returning values from a DML statement. */
150+
String getReturningAllColumnsClause() {
151+
switch (getConnection().getDialect()) {
152+
case POSTGRESQL:
153+
return "\nRETURNING *";
154+
case GOOGLE_STANDARD_SQL:
155+
default:
156+
return "\nTHEN RETURN *";
157+
}
158+
}
159+
160+
/** Returns the dialect-specific clause for returning values from a DML statement. */
161+
String getReturningClause() {
162+
switch (getConnection().getDialect()) {
163+
case POSTGRESQL:
164+
return "\nRETURNING";
165+
case GOOGLE_STANDARD_SQL:
166+
default:
167+
return "\nTHEN RETURN";
168+
}
169+
}
170+
171+
/** Adds dialect-specific quotes to the given column name. */
172+
String quoteColumn(String column) {
173+
switch (getConnection().getDialect()) {
174+
case POSTGRESQL:
175+
return "\"" + column + "\"";
176+
case GOOGLE_STANDARD_SQL:
177+
default:
178+
return "`" + column + "`";
179+
}
180+
}
181+
101182
@Override
102183
public boolean execute(String sql) throws SQLException {
103184
checkClosed();

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

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import static org.junit.Assert.assertEquals;
2222
import static org.junit.Assert.assertFalse;
2323
import static org.junit.Assert.assertNotNull;
24+
import static org.junit.Assert.assertSame;
2425
import static org.junit.Assert.assertThrows;
2526
import static org.junit.Assert.assertTrue;
2627
import static org.junit.Assert.fail;
@@ -39,6 +40,7 @@
3940
import com.google.cloud.spanner.connection.StatementResult;
4041
import com.google.cloud.spanner.connection.StatementResult.ResultType;
4142
import com.google.cloud.spanner.jdbc.JdbcSqlExceptionFactory.JdbcSqlExceptionImpl;
43+
import com.google.common.collect.ImmutableList;
4244
import com.google.rpc.Code;
4345
import java.sql.ResultSet;
4446
import java.sql.SQLException;
@@ -47,6 +49,7 @@
4749
import java.util.Arrays;
4850
import java.util.List;
4951
import java.util.concurrent.TimeUnit;
52+
import javax.annotation.Nullable;
5053
import org.junit.Test;
5154
import org.junit.runner.RunWith;
5255
import org.junit.runners.Parameterized;
@@ -599,4 +602,157 @@ public void testConvertUpdateCountsToSuccessNoInfo() throws SQLException {
599602
(long) Statement.SUCCESS_NO_INFO);
600603
}
601604
}
605+
606+
@Test
607+
public void testAddReturningToStatement() throws SQLException {
608+
JdbcConnection connection = mock(JdbcConnection.class);
609+
when(connection.getDialect()).thenReturn(dialect);
610+
when(connection.getParser()).thenReturn(AbstractStatementParser.getInstance(dialect));
611+
try (JdbcStatement statement = new JdbcStatement(connection)) {
612+
assertAddReturningSame(statement, "insert into test (id, value) values (1, 'One')", null);
613+
assertAddReturningSame(
614+
statement, "insert into test (id, value) values (1, 'One')", ImmutableList.of());
615+
assertAddReturningEquals(
616+
statement,
617+
dialect == Dialect.POSTGRESQL
618+
? "insert into test (id, value) values (1, 'One')\nRETURNING \"id\""
619+
: "insert into test (id, value) values (1, 'One')\nTHEN RETURN `id`",
620+
"insert into test (id, value) values (1, 'One')",
621+
ImmutableList.of("id"));
622+
assertAddReturningEquals(
623+
statement,
624+
dialect == Dialect.POSTGRESQL
625+
? "insert into test (id, value) values (1, 'One')\nRETURNING \"id\", \"value\""
626+
: "insert into test (id, value) values (1, 'One')\nTHEN RETURN `id`, `value`",
627+
"insert into test (id, value) values (1, 'One')",
628+
ImmutableList.of("id", "value"));
629+
assertAddReturningEquals(
630+
statement,
631+
dialect == Dialect.POSTGRESQL
632+
? "insert into test (id, value) values (1, 'One')\nRETURNING *"
633+
: "insert into test (id, value) values (1, 'One')\nTHEN RETURN *",
634+
"insert into test (id, value) values (1, 'One')",
635+
ImmutableList.of("*"));
636+
// Requesting generated keys for a DML statement that already contains a returning clause is a
637+
// no-op.
638+
assertAddReturningSame(
639+
statement,
640+
"insert into test (id, value) values (1, 'One') "
641+
+ statement.getReturningClause()
642+
+ " value",
643+
ImmutableList.of("id"));
644+
// Requesting generated keys for a query is a no-op.
645+
for (ImmutableList<String> keys :
646+
ImmutableList.of(
647+
ImmutableList.of("id"), ImmutableList.of("id", "value"), ImmutableList.of("*"))) {
648+
assertAddReturningSame(statement, "select id, value from test", keys);
649+
}
650+
651+
// Update statements may also request generated keys.
652+
assertAddReturningSame(statement, "update test set value='Two' where id=1", null);
653+
assertAddReturningSame(
654+
statement, "update test set value='Two' where id=1", ImmutableList.of());
655+
assertAddReturningEquals(
656+
statement,
657+
dialect == Dialect.POSTGRESQL
658+
? "update test set value='Two' where id=1\nRETURNING \"value\""
659+
: "update test set value='Two' where id=1\nTHEN RETURN `value`",
660+
"update test set value='Two' where id=1",
661+
ImmutableList.of("value"));
662+
assertAddReturningEquals(
663+
statement,
664+
dialect == Dialect.POSTGRESQL
665+
? "update test set value='Two' where id=1\nRETURNING \"value\", \"id\""
666+
: "update test set value='Two' where id=1\nTHEN RETURN `value`, `id`",
667+
"update test set value='Two' where id=1",
668+
ImmutableList.of("value", "id"));
669+
assertAddReturningEquals(
670+
statement,
671+
dialect == Dialect.POSTGRESQL
672+
? "update test set value='Two' where id=1\nRETURNING *"
673+
: "update test set value='Two' where id=1\nTHEN RETURN *",
674+
"update test set value='Two' where id=1",
675+
ImmutableList.of("*"));
676+
// Requesting generated keys for a DML statement that already contains a returning clause is a
677+
// no-op.
678+
assertAddReturningSame(
679+
statement,
680+
"update test set value='Two' where id=1 " + statement.getReturningClause() + " value",
681+
ImmutableList.of("value"));
682+
683+
// Delete statements may also request generated keys.
684+
assertAddReturningSame(statement, "delete test where id=1", null);
685+
assertAddReturningSame(statement, "delete test where id=1", ImmutableList.of());
686+
assertAddReturningEquals(
687+
statement,
688+
dialect == Dialect.POSTGRESQL
689+
? "delete test where id=1\nRETURNING \"value\""
690+
: "delete test where id=1\nTHEN RETURN `value`",
691+
"delete test where id=1",
692+
ImmutableList.of("value"));
693+
assertAddReturningEquals(
694+
statement,
695+
dialect == Dialect.POSTGRESQL
696+
? "delete test where id=1\nRETURNING \"id\", \"value\""
697+
: "delete test where id=1\nTHEN RETURN `id`, `value`",
698+
"delete test where id=1",
699+
ImmutableList.of("id", "value"));
700+
assertAddReturningEquals(
701+
statement,
702+
dialect == Dialect.POSTGRESQL
703+
? "delete test where id=1\nRETURNING *"
704+
: "delete test where id=1\nTHEN RETURN *",
705+
"delete test where id=1",
706+
ImmutableList.of("*"));
707+
// Requesting generated keys for a DML statement that already contains a returning clause is a
708+
// no-op.
709+
for (ImmutableList<String> keys :
710+
ImmutableList.of(
711+
ImmutableList.of("id"), ImmutableList.of("id", "value"), ImmutableList.of("*"))) {
712+
assertAddReturningSame(
713+
statement,
714+
"delete test where id=1 "
715+
+ (dialect == Dialect.POSTGRESQL
716+
? "delete test where id=1\nRETURNING"
717+
: "delete test where id=1\nTHEN RETURN")
718+
+ " value",
719+
keys);
720+
}
721+
722+
// Requesting generated keys for DDL is a no-op.
723+
for (ImmutableList<String> keys :
724+
ImmutableList.of(
725+
ImmutableList.of("id"), ImmutableList.of("id", "value"), ImmutableList.of("*"))) {
726+
assertAddReturningSame(
727+
statement,
728+
dialect == Dialect.POSTGRESQL
729+
? "create table test (id bigint primary key, value varchar)"
730+
: "create table test (id int64, value string(max)) primary key (id)",
731+
keys);
732+
}
733+
}
734+
}
735+
736+
private void assertAddReturningSame(
737+
JdbcStatement statement, String sql, @Nullable ImmutableList<String> generatedKeysColumns)
738+
throws SQLException {
739+
com.google.cloud.spanner.Statement spannerStatement =
740+
com.google.cloud.spanner.Statement.of(sql);
741+
assertSame(
742+
spannerStatement,
743+
statement.addReturningToStatement(spannerStatement, generatedKeysColumns));
744+
}
745+
746+
private void assertAddReturningEquals(
747+
JdbcStatement statement,
748+
String expectedSql,
749+
String sql,
750+
@Nullable ImmutableList<String> generatedKeysColumns)
751+
throws SQLException {
752+
com.google.cloud.spanner.Statement spannerStatement =
753+
com.google.cloud.spanner.Statement.of(sql);
754+
assertEquals(
755+
com.google.cloud.spanner.Statement.of(expectedSql),
756+
statement.addReturningToStatement(spannerStatement, generatedKeysColumns));
757+
}
602758
}

0 commit comments

Comments
 (0)