Repository: GoogleCloudPlatform/pgadapter
Pull request #276 (2022-07-11 14:35:03)

feat: execute INFORMATION_SCHEMA in single-use transaction
Queries referencing the INFORMATION_SCHEMA must use a single-use read-only
transaction. If such a query is executed during a (read/write) transaction, it
will return a specific error. This error is now detected by PGAdapter and the
query is retried in a separate single-use read-only transaction.

This behavior is needed for TypeORM.


Patch:
diff --git a/src/main/java/com/google/cloud/spanner/connection/ResultSetHelper.java b/src/main/java/com/google/cloud/spanner/connection/ResultSetHelper.java
new file mode 100644
index 00000000..6ac71b08
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/connection/ResultSetHelper.java
@@ -0,0 +1,31 @@
+// Copyright 2022 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.cloud.spanner.connection;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.spanner.ResultSet;
+
+@InternalApi
+public class ResultSetHelper {
+
+  /**
+   * Converts the given {@link ResultSet} to a {@link DirectExecuteResultSet}. A {@link
+   * DirectExecuteResultSet} does not defer the execution of the query to the first {@link
+   * ResultSet#next()} call.
+   */
+  public static ResultSet toDirectExecuteResultSet(ResultSet delegate) {
+    return DirectExecuteResultSet.ofResultSet(delegate);
+  }
+}
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
index 4ce6ffe4..1238d56f 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/BackendConnection.java
@@ -40,6 +40,7 @@ import com.google.cloud.spanner.connection.AbstractStatementParser.ParsedStateme
 import com.google.cloud.spanner.connection.AbstractStatementParser.StatementType;
 import com.google.cloud.spanner.connection.Connection;
 import com.google.cloud.spanner.connection.ConnectionOptionsHelper;
+import com.google.cloud.spanner.connection.ResultSetHelper;
 import com.google.cloud.spanner.connection.StatementResult;
 import com.google.cloud.spanner.connection.StatementResult.ClientSideStatementType;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
@@ -184,11 +185,43 @@ public class BackendConnection {
         } else {
           result.set(spannerConnection.execute(statement));
         }
+      } catch (SpannerException spannerException) {
+        // Executing queries against the information schema in a transaction is unsupported.
+        // This ensures that those queries are retried using a separate single-use transaction.
+        if (isUnsupportedConcurrencyModeException(spannerException)) {
+          try {
+            result.set(
+                new QueryResult(
+                    ResultSetHelper.toDirectExecuteResultSet(
+                        spannerConnection
+                            .getDatabaseClient()
+                            .singleUse()
+                            .executeQuery(statement))));
+            return;
+          } catch (Exception exception) {
+            result.setException(exception);
+            throw exception;
+          }
+        }
+        result.setException(spannerException);
+        throw spannerException;
       } catch (Exception exception) {
         result.setException(exception);
         throw exception;
       }
     }
+
+    /**
+     * Returns true if the given exception is the error that is returned by Cloud Spanner when an
+     * INFORMATION_SCHEMA query is not executed in a single-use read-only transaction.
+     */
+    boolean isUnsupportedConcurrencyModeException(SpannerException spannerException) {
+      return spannerException.getErrorCode() == ErrorCode.INVALID_ARGUMENT
+          && spannerException
+              .getMessage()
+              .startsWith(
+                  "INVALID_ARGUMENT: io.grpc.StatusRuntimeException: INVALID_ARGUMENT: Unsupported concurrency mode in query using INFORMATION_SCHEMA.");
+    }
   }
 
   private static final int MAX_PARTITIONS =


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/connection/ResultSetHelperTest.java b/src/test/java/com/google/cloud/spanner/connection/ResultSetHelperTest.java
new file mode 100644
index 00000000..2bb3b85e
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/connection/ResultSetHelperTest.java
@@ -0,0 +1,47 @@
+// Copyright 2022 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.cloud.spanner.connection;
+
+import static com.google.cloud.spanner.connection.ResultSetHelper.toDirectExecuteResultSet;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import com.google.cloud.spanner.ResultSet;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ResultSetHelperTest {
+
+  @Test
+  public void testToDirectExecuteResultSet() {
+    ResultSet resultSet = mock(ResultSet.class);
+    when(resultSet.next()).thenReturn(true, true, false);
+
+    ResultSet directExecuteResultSet = toDirectExecuteResultSet(resultSet);
+
+    // The underlying result set should call next() to execute the query.
+    verify(resultSet).next();
+
+    // The returned result set should have 2 rows.
+    assertTrue(directExecuteResultSet.next());
+    assertTrue(directExecuteResultSet.next());
+    assertFalse(directExecuteResultSet.next());
+  }
+}
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
index 27975c86..ba1e1bca 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -36,6 +36,8 @@ import com.google.cloud.spanner.pgadapter.wireprotocol.ParseMessage;
 import com.google.common.collect.ImmutableList;
 import com.google.protobuf.Value;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlRequest;
+import com.google.spanner.v1.BeginTransactionRequest;
+import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.ExecuteBatchDmlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
@@ -1397,4 +1399,211 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
       assertEquals("COMMIT", parseMessages.get(2).getStatement().getSql());
     }
   }
+
+  @Test
+  public void testInformationSchemaQueryInTransaction() throws SQLException {
+    String sql = "SELECT * FROM INFORMATION_SCHEMA.TABLES";
+    // Register a result for the query. Note that we don't really care what the result is, just that
+    // there is a result.
+    mockSpanner.putStatementResult(StatementResult.query(Statement.of(sql), SELECT1_RESULTSET));
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      // Make sure that we start a transaction.
+      connection.setAutoCommit(false);
+
+      // Execute a query to start the transaction.
+      try (ResultSet resultSet = connection.createStatement().executeQuery(SELECT1.getSql())) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+
+      // This ensures that the following query returns an error the first time it is executed, and
+      // then succeeds the second time. This happens because the exception is 'popped' from the
+      // response queue when it is returned. The next time the query is executed, it will return the
+      // actual result that we set.
+      mockSpanner.setExecuteStreamingSqlExecutionTime(
+          SimulatedExecutionTime.ofException(
+              Status.INVALID_ARGUMENT
+                  .withDescription(
+                      "Unsupported concurrency mode in query using INFORMATION_SCHEMA.")
+                  .asRuntimeException()));
+      try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+
+      // Make sure that the connection is still usable.
+      try (ResultSet resultSet = connection.createStatement().executeQuery(SELECT2.getSql())) {
+        assertTrue(resultSet.next());
+        assertEquals(2L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+      connection.commit();
+    }
+
+    // We should receive the INFORMATION_SCHEMA statement twice on Cloud Spanner:
+    // 1. The first time it returns an error because it is using the wrong concurrency mode.
+    // 2. The specific error will cause the connection to retry the statement using a single-use
+    //    read-only transaction.
+    assertEquals(4, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    // The first statement should start a transaction
+    assertTrue(requests.get(0).getTransaction().hasBegin());
+    // The second statement (the initial attempt of the INFORMATION_SCHEMA query) should try to use
+    // the transaction.
+    assertTrue(requests.get(1).getTransaction().hasId());
+    assertEquals(sql, requests.get(1).getSql());
+    // The INFORMATION_SCHEMA query is then retried using a single-use read-only transaction.
+    assertFalse(requests.get(2).hasTransaction());
+    assertEquals(sql, requests.get(2).getSql());
+    // The last statement should use the transaction.
+    assertTrue(requests.get(3).getTransaction().hasId());
+
+    assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
+    CommitRequest commitRequest = mockSpanner.getRequestsOfType(CommitRequest.class).get(0);
+    assertEquals(commitRequest.getTransactionId(), requests.get(1).getTransaction().getId());
+    assertEquals(commitRequest.getTransactionId(), requests.get(3).getTransaction().getId());
+  }
+
+  @Test
+  public void testInformationSchemaQueryAsFirstQueryInTransaction() throws SQLException {
+    // Running an information_schema query as the first query in a transaction will cause some
+    // additional retrying and transaction magic. This is because:
+    // 1. The first query in a transaction will also try to begin the transaction.
+    // 2. If the first query fails, it will also fail to create a transaction.
+    // 3. If an additional query is executed in the transaction, the entire transaction will be
+    //    retried using an explicit BeginTransaction RPC. This is done so that we can include the
+    //    first query in the transaction, as an error message in itself can give away information
+    //    about the state of the database, and therefore must be included in the transaction to
+    //    guarantee the consistency.
+
+    String sql = "SELECT * FROM INFORMATION_SCHEMA.TABLES";
+    // Register a result for the query. Note that we don't really care what the result is, just that
+    // there is a result.
+    mockSpanner.putStatementResult(StatementResult.query(Statement.of(sql), SELECT1_RESULTSET));
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      // Make sure that we start a transaction.
+      connection.setAutoCommit(false);
+
+      // This makes sure the INFORMATION_SCHEMA query will return an error the first time it is
+      // executed. Then it is retried without a transaction.
+      mockSpanner.setExecuteStreamingSqlExecutionTime(
+          SimulatedExecutionTime.ofException(
+              Status.INVALID_ARGUMENT
+                  .withDescription(
+                      "Unsupported concurrency mode in query using INFORMATION_SCHEMA.")
+                  .asRuntimeException()));
+      try (ResultSet resultSet = connection.createStatement().executeQuery(sql)) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+
+      // This ensures that the next query will once again return the same error. The reason that we
+      // do this is that the following query will cause the entire transaction to be retried, and we
+      // need the first statement (the INFORMATION_SCHEMA query) to return exactly the same result
+      // as the first time in order to make the retry succeed.
+      mockSpanner.setExecuteStreamingSqlExecutionTime(
+          SimulatedExecutionTime.ofException(
+              Status.INVALID_ARGUMENT
+                  .withDescription(
+                      "Unsupported concurrency mode in query using INFORMATION_SCHEMA.")
+                  .asRuntimeException()));
+
+      // Verify that the connection is still usable.
+      try (ResultSet resultSet = connection.createStatement().executeQuery(SELECT2.getSql())) {
+        assertTrue(resultSet.next());
+        assertEquals(2L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+      connection.commit();
+    }
+
+    // We should receive the INFORMATION_SCHEMA statement three times on Cloud Spanner:
+    // 1. The first time it returns an error because it is using the wrong concurrency mode.
+    // 2. The specific error will cause the connection to retry the statement using a single-use
+    //    read-only transaction.
+    // 3. The second query in the transaction will cause the entire transaction to retry, which will
+    //    cause the INFORMATION_SCHEMA query to be executed once more.
+    assertEquals(4, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    List<ExecuteSqlRequest> requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+    // The first statement should try to start a transaction (although it fails).
+    assertTrue(requests.get(0).getTransaction().hasBegin());
+    // The second statement is the INFORMATION_SCHEMA query without a transaction.
+    assertFalse(requests.get(1).hasTransaction());
+    assertEquals(sql, requests.get(1).getSql());
+
+    // The transaction is then retried, which means that we get the INFORMATION_SCHEMA query again.
+    // This time the query tries to use a transaction that has been created using an explicit
+    // BeginTransaction RPC invocation.
+    assertTrue(requests.get(2).getTransaction().hasId());
+    assertEquals(sql, requests.get(2).getSql());
+    // The last query should also use that transaction.
+    assertTrue(requests.get(3).getTransaction().hasId());
+    assertEquals(SELECT2.getSql(), requests.get(3).getSql());
+
+    assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
+    CommitRequest commitRequest = mockSpanner.getRequestsOfType(CommitRequest.class).get(0);
+    assertEquals(commitRequest.getTransactionId(), requests.get(2).getTransaction().getId());
+    assertEquals(commitRequest.getTransactionId(), requests.get(3).getTransaction().getId());
+    // Verify that we also got a BeginTransaction request.
+    assertEquals(1, mockSpanner.countRequestsOfType(BeginTransactionRequest.class));
+  }
+
+  @Test
+  public void testInformationSchemaQueryInTransactionWithErrorDuringRetry() throws SQLException {
+    String sql = "SELECT * FROM INFORMATION_SCHEMA.TABLES";
+    // Register a result for the query. Note that we don't really care what the result is, just that
+    // there is a result.
+    mockSpanner.putStatementResult(StatementResult.query(Statement.of(sql), SELECT1_RESULTSET));
+
+    try (Connection connection = DriverManager.getConnection(createUrl())) {
+      // Make sure that we start a transaction.
+      connection.setAutoCommit(false);
+
+      // Execute a query to start the transaction.
+      try (ResultSet resultSet = connection.createStatement().executeQuery(SELECT1.getSql())) {
+        assertTrue(resultSet.next());
+        assertEquals(1L, resultSet.getLong(1));
+        assertFalse(resultSet.next());
+      }
+
+      // This ensures that the following query returns the specific concurrency error the first time
+      // it is executed, and then a different error.
+      mockSpanner.setExecuteStreamingSqlExecutionTime(
+          SimulatedExecutionTime.ofExceptions(
+              ImmutableList.of(
+                  Status.INVALID_ARGUMENT
+                      .withDescription(
+                          "Unsupported concurrency mode in query using INFORMATION_SCHEMA.")
+                      .asRuntimeException(),
+                  Status.INTERNAL.withDescription("test error").asRuntimeException())));
+      SQLException sqlException =
+          assertThrows(SQLException.class, () -> connection.createStatement().executeQuery(sql));
+      assertEquals(
+          "ERROR: INTERNAL: io.grpc.StatusRuntimeException: INTERNAL: test error - Statement: 'SELECT * FROM INFORMATION_SCHEMA.TABLES'",
+          sqlException.getMessage());
+
+      // Make sure that the connection is now in the aborted state.
+      SQLException abortedException =
+          assertThrows(
+              SQLException.class,
+              () -> connection.createStatement().executeQuery(SELECT2.getSql()));
+      assertEquals(
+          "ERROR: INVALID_ARGUMENT: current transaction is aborted, commands ignored until end of transaction block",
+          abortedException.getMessage());
+    }
+
+    // We should receive the INFORMATION_SCHEMA statement twice on Cloud Spanner:
+    // 1. The first time it returns an error because it is using the wrong concurrency mode.
+    // 2. The specific error will cause the connection to retry the statement using a single-use
+    //    read-only transaction. That will also fail with the second error.
+    // 3. The following SELECT query is never sent to Cloud Spanner, as the transaction is in the
+    //    aborted state.
+    assertEquals(3, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class));
+    assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
+  }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java
index 9f283d16..524ddb69 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/BackendConnectionTest.java
@@ -15,6 +15,7 @@
 package com.google.cloud.spanner.pgadapter.statements;
 
 import static com.google.cloud.spanner.pgadapter.statements.BackendConnection.extractDdlUpdateCounts;
+import static com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.EMPTY_LOCAL_STATEMENTS;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -330,4 +331,26 @@ public class BackendConnectionTest {
     assertThrows(UnsupportedOperationException.class, queryResult::getClientSideStatementType);
     assertThrows(UnsupportedOperationException.class, queryResult::getUpdateCount);
   }
+
+  @Test
+  public void testGeneralException() {
+    Connection connection = mock(Connection.class);
+    ParsedStatement parsedStatement = mock(ParsedStatement.class);
+    Statement statement = Statement.of("select foo from bar");
+    RuntimeException error = new RuntimeException("test error");
+    when(connection.execute(statement)).thenThrow(error);
+
+    BackendConnection backendConnection =
+        new BackendConnection(
+            DatabaseId.of("p", "i", "d"),
+            connection,
+            mock(OptionsMetadata.class),
+            EMPTY_LOCAL_STATEMENTS);
+    Future<StatementResult> resultFuture = backendConnection.execute(parsedStatement, statement);
+    backendConnection.flush();
+
+    ExecutionException executionException =
+        assertThrows(ExecutionException.class, resultFuture::get);
+    assertSame(executionException.getCause(), error);
+  }
 }