Skip to content

Commit 2ffb290

Browse files
olavloiteolavloite
authored
fix: allow start ddl batch / run batch in one query string (#529)
* fix: allow start ddl batch / run batch in one query string Sending an explicit DDL batch in a single query string caused the error "DDL statements are not allowed in mixed batches or transactions." This change fixes that so it is possible to send such a batch both as an implicit (without `start batch ddl`) as well as an explicit ddl batch. Fixes #528 * test: add more tests
1 parent 3668650 commit 2ffb290

File tree

3 files changed

+135
-3
lines changed

3 files changed

+135
-3
lines changed

‎src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java‎

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@
6464
import com.google.common.collect.ImmutableList;
6565
import com.google.common.collect.ImmutableMap;
6666
import com.google.common.collect.ImmutableMap.Builder;
67+
import com.google.common.collect.ImmutableSet;
6768
import com.google.common.util.concurrent.Futures;
6869
import com.google.common.util.concurrent.ListenableFuture;
6970
import com.google.common.util.concurrent.ListeningExecutorService;
@@ -635,7 +636,10 @@ private void flush(boolean isSync) {
635636
spannerConnection.beginTransaction();
636637
}
637638
boolean canUseBatch = false;
638-
if (bufferedStatement.isBatchingPossible() && index < (getStatementCount() - 1)) {
639+
if (!spannerConnection.isDdlBatchActive()
640+
&& !spannerConnection.isDmlBatchActive()
641+
&& bufferedStatement.isBatchingPossible()
642+
&& index < (getStatementCount() - 1)) {
639643
StatementType statementType = getStatementType(index);
640644
StatementType nextStatementType = getStatementType(index + 1);
641645
canUseBatch = canBeBatchedTogether(statementType, nextStatementType);
@@ -793,7 +797,8 @@ private void prepareExecuteDdl(BufferedStatement<?> bufferedStatement) {
793797
case Batch:
794798
if (spannerConnection.isInTransaction()
795799
|| bufferedStatements.stream()
796-
.anyMatch(statement -> !statement.parsedStatement.isDdl())) {
800+
.anyMatch(
801+
statement -> !isStatementAllowedInDdlBatch(statement.parsedStatement))) {
797802
throw PGExceptionFactory.newPGException(
798803
"DDL statements are not allowed in mixed batches or transactions.",
799804
SQLState.InvalidTransactionState);
@@ -826,6 +831,18 @@ private void prepareExecuteDdl(BufferedStatement<?> bufferedStatement) {
826831
}
827832
}
828833

834+
private static final ImmutableSet<ClientSideStatementType> DDL_BATCH_STATEMENTS =
835+
ImmutableSet.of(
836+
ClientSideStatementType.START_BATCH_DDL,
837+
ClientSideStatementType.RUN_BATCH,
838+
ClientSideStatementType.ABORT_BATCH);
839+
840+
private boolean isStatementAllowedInDdlBatch(ParsedStatement parsedStatement) {
841+
return parsedStatement.isDdl()
842+
|| (parsedStatement.getType() == StatementType.CLIENT_SIDE
843+
&& DDL_BATCH_STATEMENTS.contains(parsedStatement.getClientSideStatementType()));
844+
}
845+
829846
private boolean isBegin(int index) {
830847
return isBegin(bufferedStatements.get(index).parsedStatement);
831848
}

‎src/test/java/com/google/cloud/spanner/pgadapter/JdbcSimpleModeMockServerTest.java‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import com.google.common.collect.ImmutableList;
2727
import com.google.protobuf.ListValue;
2828
import com.google.protobuf.Value;
29+
import com.google.spanner.admin.database.v1.UpdateDatabaseDdlRequest;
2930
import com.google.spanner.v1.CommitRequest;
3031
import com.google.spanner.v1.ExecuteSqlRequest;
3132
import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
@@ -46,6 +47,7 @@
4647
import java.util.Collections;
4748
import java.util.List;
4849
import java.util.TimeZone;
50+
import java.util.stream.Collectors;
4951
import org.junit.BeforeClass;
5052
import org.junit.Test;
5153
import org.junit.runner.RunWith;
@@ -634,4 +636,118 @@ public void testSetInvalidTimezone() throws SQLException {
634636
"ERROR: invalid value for parameter \"TimeZone\": \"foo\"", exception.getMessage());
635637
}
636638
}
639+
640+
@Test
641+
public void testImplicitDdlBatch() throws SQLException {
642+
String sql =
643+
"create table test1 (id bigint primary key); "
644+
+ "create table test2 (id bigint primary key); ";
645+
addDdlResponseToSpannerAdmin();
646+
647+
try (Connection connection = DriverManager.getConnection(createUrl())) {
648+
connection.createStatement().execute(sql);
649+
}
650+
651+
List<UpdateDatabaseDdlRequest> requests =
652+
mockDatabaseAdmin.getRequests().stream()
653+
.filter(message -> message instanceof UpdateDatabaseDdlRequest)
654+
.map(message -> (UpdateDatabaseDdlRequest) message)
655+
.collect(Collectors.toList());
656+
assertEquals(1, requests.size());
657+
assertEquals(2, requests.get(0).getStatementsCount());
658+
assertEquals("create table test1 (id bigint primary key)", requests.get(0).getStatements(0));
659+
assertEquals("create table test2 (id bigint primary key)", requests.get(0).getStatements(1));
660+
}
661+
662+
@Test
663+
public void testDdlBatchWithStartAndRun() throws SQLException {
664+
String sql =
665+
"start batch ddl; "
666+
+ "create table test1 (id bigint primary key); "
667+
+ "create table test2 (id bigint primary key); "
668+
+ "run batch";
669+
addDdlResponseToSpannerAdmin();
670+
671+
try (Connection connection = DriverManager.getConnection(createUrl())) {
672+
connection.createStatement().execute(sql);
673+
}
674+
675+
List<UpdateDatabaseDdlRequest> requests =
676+
mockDatabaseAdmin.getRequests().stream()
677+
.filter(message -> message instanceof UpdateDatabaseDdlRequest)
678+
.map(message -> (UpdateDatabaseDdlRequest) message)
679+
.collect(Collectors.toList());
680+
assertEquals(1, requests.size());
681+
assertEquals(2, requests.get(0).getStatementsCount());
682+
assertEquals("create table test1 (id bigint primary key)", requests.get(0).getStatements(0));
683+
assertEquals("create table test2 (id bigint primary key)", requests.get(0).getStatements(1));
684+
}
685+
686+
@Test
687+
public void testMixedImplicitDdlBatch() throws SQLException {
688+
String sql =
689+
"create table test1 (id bigint primary key); "
690+
+ "show statement_timeout; "
691+
+ "create table test2 (id bigint primary key); ";
692+
addDdlResponseToSpannerAdmin();
693+
694+
try (Connection connection = DriverManager.getConnection(createUrl())) {
695+
SQLException exception =
696+
assertThrows(SQLException.class, () -> connection.createStatement().execute(sql));
697+
assertEquals(
698+
"ERROR: DDL statements are not allowed in mixed batches or transactions.",
699+
exception.getMessage());
700+
}
701+
702+
List<UpdateDatabaseDdlRequest> requests =
703+
mockDatabaseAdmin.getRequests().stream()
704+
.filter(message -> message instanceof UpdateDatabaseDdlRequest)
705+
.map(message -> (UpdateDatabaseDdlRequest) message)
706+
.collect(Collectors.toList());
707+
assertEquals(0, requests.size());
708+
}
709+
710+
@Test
711+
public void testMixedDdlBatchWithStartAndRun() throws SQLException {
712+
String sql =
713+
"start batch ddl; "
714+
+ "create table test1 (id bigint primary key); "
715+
+ "set statement_timeout = '10s'; "
716+
+ "create table test2 (id bigint primary key); "
717+
+ "run batch";
718+
addDdlResponseToSpannerAdmin();
719+
720+
try (Connection connection = DriverManager.getConnection(createUrl())) {
721+
SQLException exception =
722+
assertThrows(SQLException.class, () -> connection.createStatement().execute(sql));
723+
assertEquals(
724+
"ERROR: DDL statements are not allowed in mixed batches or transactions.",
725+
exception.getMessage());
726+
}
727+
728+
List<UpdateDatabaseDdlRequest> requests =
729+
mockDatabaseAdmin.getRequests().stream()
730+
.filter(message -> message instanceof UpdateDatabaseDdlRequest)
731+
.map(message -> (UpdateDatabaseDdlRequest) message)
732+
.collect(Collectors.toList());
733+
assertEquals(0, requests.size());
734+
}
735+
736+
@Test
737+
public void testImplicitBatchOfClientSideStatements() throws SQLException {
738+
String sql = "set statement_timeout = '10s'; " + "show statement_timeout; ";
739+
addDdlResponseToSpannerAdmin();
740+
741+
try (Connection connection = DriverManager.getConnection(createUrl())) {
742+
try (Statement statement = connection.createStatement()) {
743+
assertFalse(statement.execute(sql));
744+
assertTrue(statement.getMoreResults());
745+
try (ResultSet resultSet = statement.getResultSet()) {
746+
assertTrue(resultSet.next());
747+
assertEquals("10s", resultSet.getString(1));
748+
assertFalse(resultSet.next());
749+
}
750+
}
751+
}
752+
}
637753
}

‎src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -444,7 +444,6 @@ public void testDdlExceptionInBatch() {
444444
when(connection.execute(statement1)).thenReturn(NO_RESULT);
445445
RuntimeException error = new RuntimeException("test error");
446446
when(connection.execute(statement2)).thenThrow(error);
447-
when(connection.isDdlBatchActive()).thenReturn(true);
448447

449448
BackendConnection backendConnection =
450449
new BackendConnection(

0 commit comments

Comments
 (0)