Repository: GoogleCloudPlatform/pgadapter
Pull request #494 (2022-11-18 19:03:19)

fix: pg_catalog tables were not replaced for information_schema queries
Queries in read/write transactions that refer to the information schema were retried using a single use read-only transaction, but not with the replaced pg_catalog tables. This would cause 'Table not found' errors.

codecov[bot]: # [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/494?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#494](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/494?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (b872336) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/e6ffe032d30b762da974cc75fda7e7d3f144b4c6?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (e6ffe03) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@                  Coverage Diff                  @@
##             postgresql-dialect     #494   +/-   ##
=====================================================
  Coverage                 87.51%   87.52%           
  Complexity                 2029     2029           
=====================================================
  Files                       122      122           
  Lines                      6666     6667    +1     
  Branches                    960      960           
=====================================================
+ Hits                       5834     5835    +1     
  Misses                      583      583           
  Partials                    249      249           
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `87.52% <100.00%> (+<0.01%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/494?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [...panner/pgadapter/statements/BackendConnection.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/494/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQmFja2VuZENvbm5lY3Rpb24uamF2YQ==) | `92.09% <100.00%> (+0.01%)` | :arrow_up: |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform)


Patch:
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 7566bae8..13e1cc82 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
@@ -167,6 +167,7 @@ public class BackendConnection {
 
     @Override
     void execute() {
+      Statement updatedStatement = statement;
       try {
         checkConnectionState();
         // TODO(b/235719478): If the statement is a BEGIN statement and there is a COMMIT statement
@@ -212,7 +213,7 @@ public class BackendConnection {
           result.set(ddlExecutor.execute(parsedStatement, statement));
         } else {
           // Potentially replace pg_catalog table references with common table expressions.
-          Statement updatedStatement =
+          updatedStatement =
               sessionState.isReplacePgCatalogTables()
                   ? pgCatalog.replacePgCatalogTables(statement)
                   : statement;
@@ -229,7 +230,7 @@ public class BackendConnection {
                         spannerConnection
                             .getDatabaseClient()
                             .singleUse()
-                            .executeQuery(statement))));
+                            .executeQuery(updatedStatement))));
             return;
           } catch (Exception exception) {
             throw setAndReturn(result, exception);


Test patch:
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 4217e8d2..7707972b 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -2004,6 +2004,82 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
     assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));
   }
 
+  @Test
+  public void testInformationSchemaQueryInTransactionWithReplacedPgCatalogTables()
+      throws SQLException {
+    String sql = "SELECT 1 FROM pg_namespace";
+    String replacedSql =
+        "with pg_namespace as (\n"
+            + "  select case schema_name when 'pg_catalog' then 11 when 'public' then 2200 else 0 end as oid,\n"
+            + "        schema_name as nspname, null as nspowner, null as nspacl\n"
+            + "  from information_schema.schemata\n"
+            + ")\n"
+            + "SELECT 1 FROM pg_namespace";
+    // 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(replacedSql), 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(replacedSql, requests.get(1).getSql());
+    // The INFORMATION_SCHEMA query is then retried using a single-use read-only transaction.
+    assertFalse(requests.get(2).hasTransaction());
+    assertEquals(replacedSql, 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 testShowGuessTypes() throws SQLException {
     try (Connection connection = DriverManager.getConnection(createUrl())) {