Repository: GoogleCloudPlatform/pgadapter
Pull request #605 (2023-01-20 11:27:29)

feat: make table and function replacements client-aware
The current pg_catalog replacement strategy is fixed for all clients. This means that if a specific client wants to use a pg_catalog table that is currently not supported by Cloud Spanner, we need to add it for all clients. This can be tricky if we cannot really support the table or function. This change makes the pg_catalog implementation in PGAdapter client-aware. This also means that the initialization of pg_catalog must use a lazy-initialized version of the WellKnownClient, as the client may only be known once we get the first query. The detection of the client based on the first query happens before the query is executed. This means that we can safely assume that the WellKnownClient has been initialized when the backend connection needs pg_catalog the first time.

codecov[bot]: # [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/605?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#605](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/605?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (378c6c3) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/8dc7a888614e4cc163735c28ac91f71c2b261601?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (8dc7a88) will **increase** coverage by `0.08%`.
> The diff coverage is `98.00%`.

```diff
@@                   Coverage Diff                    @@
##             postgresql-dialect     #605      +/-   ##
========================================================
+ Coverage                 88.63%   88.71%   +0.08%     
- Complexity                 2135     2137       +2     
========================================================
  Files                       124      125       +1     
  Lines                      7101     7134      +33     
  Branches                   1011     1012       +1     
========================================================
+ Hits                       6294     6329      +35     
+ Misses                      569      568       -1     
+ Partials                    238      237       -1     
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `88.71% <98.00%> (+0.08%)` | :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/605?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [...ud/spanner/pgadapter/metadata/OptionsMetadata.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/605?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL21ldGFkYXRhL09wdGlvbnNNZXRhZGF0YS5qYXZh) | `81.02% <ø> (ø)` | |
| [.../cloud/spanner/pgadapter/session/SessionState.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/605?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3Nlc3Npb24vU2Vzc2lvblN0YXRlLmphdmE=) | `95.18% <0.00%> (-0.52%)` | :arrow_down: |
| [...panner/pgadapter/statements/BackendConnection.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/605?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvQmFja2VuZENvbm5lY3Rpb24uamF2YQ==) | `93.58% <100.00%> (+0.05%)` | :arrow_up: |
| [...apter/statements/ExtendedQueryProtocolHandler.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/605?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvRXh0ZW5kZWRRdWVyeVByb3RvY29sSGFuZGxlci5qYXZh) | `97.72% <100.00%> (+0.05%)` | :arrow_up: |
| [.../cloud/spanner/pgadapter/statements/PgCatalog.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/605?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvUGdDYXRhbG9nLmphdmE=) | `95.83% <100.00%> (+0.83%)` | :arrow_up: |
| [...ud/spanner/pgadapter/utils/ClientAutoDetector.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/605?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL0NsaWVudEF1dG9EZXRlY3Rvci5qYXZh) | `91.80% <100.00%> (+0.73%)` | :arrow_up: |
| [...google/cloud/spanner/pgadapter/utils/LazyInit.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/605?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL0xhenlJbml0LmphdmE=) | `100.00% <100.00%> (ø)` | |
| [...gle/cloud/spanner/pgadapter/ConnectionHandler.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/605?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL0Nvbm5lY3Rpb25IYW5kbGVyLmphdmE=) | `87.22% <0.00%> (+0.83%)` | :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/metadata/OptionsMetadata.java b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
index 6abac46f..f3b595cc 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
@@ -88,7 +88,7 @@ public class OptionsMetadata {
   }
 
   private static final Logger logger = Logger.getLogger(OptionsMetadata.class.getName());
-  private static final String DEFAULT_SERVER_VERSION = "14.1";
+  public static final String DEFAULT_SERVER_VERSION = "14.1";
   private static final String DEFAULT_USER_AGENT = "pg-adapter";
 
   private static final String OPTION_SERVER_PORT = "s";
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java b/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java
index 4692e590..ca37be80 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/session/SessionState.java
@@ -374,6 +374,11 @@ public class SessionState {
     this.transactionSettings = null;
   }
 
+  /** Returns the PostgreSQL version. */
+  public String getServerVersion() {
+    return getStringSetting(null, "server_version", OptionsMetadata.DEFAULT_SERVER_VERSION);
+  }
+
   /**
    * Returns whether transaction statements should be ignored and all statements should be executed
    * in autocommit mode.
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 a5cdf21d..e9171ca5 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
@@ -59,12 +59,14 @@ import com.google.cloud.spanner.pgadapter.statements.DdlExecutor.NotExecuted;
 import com.google.cloud.spanner.pgadapter.statements.SessionStatementParser.SessionStatement;
 import com.google.cloud.spanner.pgadapter.statements.SimpleParser.TableOrIndexName;
 import com.google.cloud.spanner.pgadapter.statements.local.LocalStatement;
+import com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.WellKnownClient;
 import com.google.cloud.spanner.pgadapter.utils.CopyDataReceiver;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
 import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse;
 import com.google.cloud.spanner.pgadapter.wireoutput.ReadyResponse.Status;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
+import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMap.Builder;
@@ -87,6 +89,7 @@ import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 
@@ -208,9 +211,10 @@ public class BackendConnection {
         //  block always ends with a ROLLBACK, PGAdapter should skip the entire execution of that
         //  block.
         SessionStatement sessionStatement = getSessionManagementStatement(parsedStatement);
-        if (!localStatements.isEmpty() && localStatements.containsKey(statement.getSql())) {
+        if (!localStatements.get().isEmpty()
+            && localStatements.get().containsKey(statement.getSql())) {
           result.set(
-              Objects.requireNonNull(localStatements.get(statement.getSql()))
+              Objects.requireNonNull(localStatements.get().get(statement.getSql()))
                   .execute(BackendConnection.this));
         } else if (sessionStatement != null) {
           result.set(sessionStatement.execute(sessionState));
@@ -253,7 +257,7 @@ public class BackendConnection {
           // Potentially replace pg_catalog table references with common table expressions.
           updatedStatement =
               sessionState.isReplacePgCatalogTables()
-                  ? pgCatalog.replacePgCatalogTables(statement)
+                  ? pgCatalog.get().replacePgCatalogTables(statement)
                   : statement;
           updatedStatement = bindStatement(updatedStatement);
           result.set(analyzeOrExecute(updatedStatement));
@@ -567,8 +571,8 @@ public class BackendConnection {
   private static final Statement ROLLBACK = Statement.of("ROLLBACK");
 
   private final SessionState sessionState;
-  private final PgCatalog pgCatalog;
-  private final ImmutableMap<String, LocalStatement> localStatements;
+  private final Supplier<PgCatalog> pgCatalog;
+  private final Supplier<ImmutableMap<String, LocalStatement>> localStatements;
   private ConnectionState connectionState = ConnectionState.IDLE;
   private TransactionMode transactionMode = TransactionMode.IMPLICIT;
   private final String currentSchema = "public";
@@ -581,24 +585,31 @@ public class BackendConnection {
   BackendConnection(
       DatabaseId databaseId,
       Connection spannerConnection,
+      Supplier<WellKnownClient> wellKnownClient,
       OptionsMetadata optionsMetadata,
-      ImmutableList<LocalStatement> localStatements) {
+      Supplier<ImmutableList<LocalStatement>> localStatements) {
     this.sessionState = new SessionState(optionsMetadata);
-    this.pgCatalog = new PgCatalog(this.sessionState);
+    this.pgCatalog =
+        Suppliers.memoize(
+            () -> new PgCatalog(BackendConnection.this.sessionState, wellKnownClient.get()));
     this.spannerConnection = spannerConnection;
     this.databaseId = databaseId;
     this.ddlExecutor = new DdlExecutor(databaseId, this);
-    if (localStatements.isEmpty()) {
-      this.localStatements = EMPTY_LOCAL_STATEMENTS;
-    } else {
-      Builder<String, LocalStatement> builder = ImmutableMap.builder();
-      for (LocalStatement localStatement : localStatements) {
-        for (String sql : localStatement.getSql()) {
-          builder.put(new SimpleImmutableEntry<>(sql, localStatement));
-        }
-      }
-      this.localStatements = builder.build();
-    }
+    this.localStatements =
+        Suppliers.memoize(
+            () -> {
+              if (localStatements.get().isEmpty()) {
+                return EMPTY_LOCAL_STATEMENTS;
+              } else {
+                Builder<String, LocalStatement> builder = ImmutableMap.builder();
+                for (LocalStatement localStatement : localStatements.get()) {
+                  for (String sql : localStatement.getSql()) {
+                    builder.put(new SimpleImmutableEntry<>(sql, localStatement));
+                  }
+                }
+                return builder.build();
+              }
+            });
   }
 
   /** Returns the current connection state. */
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandler.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandler.java
index ae61125d..9d3c5bf7 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandler.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/ExtendedQueryProtocolHandler.java
@@ -40,8 +40,9 @@ public class ExtendedQueryProtocolHandler {
         new BackendConnection(
             connectionHandler.getDatabaseId(),
             connectionHandler.getSpannerConnection(),
+            connectionHandler::getWellKnownClient,
             connectionHandler.getServer().getOptions(),
-            connectionHandler.getWellKnownClient().getLocalStatements(connectionHandler));
+            () -> connectionHandler.getWellKnownClient().getLocalStatements(connectionHandler));
   }
 
   /** Constructor only intended for testing. */
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/statements/PgCatalog.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/PgCatalog.java
index a1d4aac4..0fa0dcbf 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/statements/PgCatalog.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/statements/PgCatalog.java
@@ -20,6 +20,8 @@ import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.Value;
 import com.google.cloud.spanner.pgadapter.session.SessionState;
 import com.google.cloud.spanner.pgadapter.statements.SimpleParser.TableOrIndexName;
+import com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.WellKnownClient;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -27,11 +29,13 @@ import java.util.HashSet;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
+import java.util.function.Supplier;
 import java.util.regex.Pattern;
+import javax.annotation.Nonnull;
 
 @InternalApi
 public class PgCatalog {
-  private static final ImmutableMap<TableOrIndexName, TableOrIndexName> TABLE_REPLACEMENTS =
+  private static final ImmutableMap<TableOrIndexName, TableOrIndexName> DEFAULT_TABLE_REPLACEMENTS =
       ImmutableMap.<TableOrIndexName, TableOrIndexName>builder()
           .put(
               new TableOrIndexName("pg_catalog", "pg_namespace"),
@@ -63,6 +67,12 @@ public class PgCatalog {
           Pattern.compile("pg_table_is_visible\\(.+\\)"), "true",
           Pattern.compile("ANY\\(current_schemas\\(true\\)\\)"), "'public'");
 
+  private final ImmutableSet<String> checkPrefixes;
+
+  private final ImmutableMap<TableOrIndexName, TableOrIndexName> tableReplacements;
+
+  private final ImmutableMap<Pattern, Supplier<String>> functionReplacements;
+
   private final Map<TableOrIndexName, PgCatalogTable> pgCatalogTables =
       ImmutableMap.of(
           new TableOrIndexName(null, "pg_namespace"), new PgNamespace(),
@@ -73,14 +83,31 @@ public class PgCatalog {
           new TableOrIndexName(null, "pg_settings"), new PgSettings());
   private final SessionState sessionState;
 
-  public PgCatalog(SessionState sessionState) {
-    this.sessionState = sessionState;
+  public PgCatalog(@Nonnull SessionState sessionState, @Nonnull WellKnownClient wellKnownClient) {
+    this.sessionState = Preconditions.checkNotNull(sessionState);
+    this.checkPrefixes = wellKnownClient.getPgCatalogCheckPrefixes();
+    ImmutableMap.Builder<TableOrIndexName, TableOrIndexName> builder =
+        ImmutableMap.<TableOrIndexName, TableOrIndexName>builder()
+            .putAll(DEFAULT_TABLE_REPLACEMENTS);
+    wellKnownClient
+        .getTableReplacements()
+        .forEach((k, v) -> builder.put(TableOrIndexName.of(k), TableOrIndexName.of(v)));
+    this.tableReplacements = builder.build();
+    this.functionReplacements =
+        ImmutableMap.<Pattern, Supplier<String>>builder()
+            .put(Pattern.compile("pg_catalog.pg_table_is_visible\\(.+\\)"), () -> "true")
+            .put(Pattern.compile("pg_table_is_visible\\(.+\\)"), () -> "true")
+            .put(Pattern.compile("ANY\\(current_schemas\\(true\\)\\)"), () -> "'public'")
+            .put(
+                Pattern.compile("version\\(\\)"), () -> "'" + sessionState.getServerVersion() + "'")
+            .putAll(wellKnownClient.getFunctionReplacements())
+            .build();
   }
 
   /** Replace supported pg_catalog tables with Common Table Expressions. */
   public Statement replacePgCatalogTables(Statement statement) {
     Tuple<Set<TableOrIndexName>, Statement> replacedTablesStatement =
-        new TableParser(statement).detectAndReplaceTables(TABLE_REPLACEMENTS);
+        new TableParser(statement).detectAndReplaceTables(DEFAULT_TABLE_REPLACEMENTS);
     if (replacedTablesStatement.x().isEmpty()) {
       return replacedTablesStatement.y();
     }
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java
index 78b25ad1..3532d495 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/ClientAutoDetector.java
@@ -31,9 +31,12 @@ import com.google.cloud.spanner.pgadapter.wireoutput.NoticeResponse.NoticeSeveri
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import java.time.Duration;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;
+import java.util.regex.Pattern;
 import javax.annotation.Nonnull;
 import org.postgresql.core.Oid;
 
@@ -51,6 +54,8 @@ public class ClientAutoDetector {
           SelectCurrentCatalogStatement.INSTANCE,
           SelectVersionStatement.INSTANCE,
           DjangoGetTableNamesStatement.INSTANCE);
+  private static final ImmutableSet<String> DEFAULT_CHECK_PG_CATALOG_PREFIXES =
+      ImmutableSet.of("pg_");
   public static final String PGBENCH_USAGE_HINT =
       "See https://github.com/GoogleCloudPlatform/pgadapter/blob/-/docs/pgbench.md for how to use pgbench with PGAdapter";
 
@@ -218,6 +223,18 @@ public class ClientAutoDetector {
       return EMPTY_LOCAL_STATEMENTS;
     }
 
+    public ImmutableSet<String> getPgCatalogCheckPrefixes() {
+      return DEFAULT_CHECK_PG_CATALOG_PREFIXES;
+    }
+
+    public ImmutableMap<String, String> getTableReplacements() {
+      return ImmutableMap.of();
+    }
+
+    public ImmutableMap<Pattern, Supplier<String>> getFunctionReplacements() {
+      return ImmutableMap.of();
+    }
+
     /** Creates specific notice messages for a client after startup. */
     public ImmutableList<NoticeResponse> createStartupNoticeResponses(
         ConnectionHandler connection) {


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/session/SessionStateTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/session/SessionStateTest.java
index dac991ac..d67eb144 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/session/SessionStateTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/session/SessionStateTest.java
@@ -32,6 +32,7 @@ import com.google.cloud.spanner.pgadapter.error.PGException;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata.DdlTransactionMode;
 import com.google.cloud.spanner.pgadapter.statements.PgCatalog;
+import com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.WellKnownClient;
 import com.google.common.collect.ImmutableMap;
 import java.util.List;
 import java.util.Map;
@@ -463,7 +464,7 @@ public class SessionStateTest {
   @Test
   public void testAddSessionState() {
     SessionState state = new SessionState(mock(OptionsMetadata.class));
-    PgCatalog pgCatalog = new PgCatalog(state);
+    PgCatalog pgCatalog = new PgCatalog(state, WellKnownClient.UNSPECIFIED);
     Statement statement = Statement.of("select * from pg_settings");
 
     Statement withSessionState = pgCatalog.replacePgCatalogTables(statement);
@@ -476,7 +477,7 @@ public class SessionStateTest {
   @Test
   public void testAddSessionStateWithParameters() {
     SessionState state = new SessionState(mock(OptionsMetadata.class));
-    PgCatalog pgCatalog = new PgCatalog(state);
+    PgCatalog pgCatalog = new PgCatalog(state, WellKnownClient.UNSPECIFIED);
     Statement statement =
         Statement.newBuilder("select * from pg_settings where name=$1")
             .bind("p1")
@@ -497,7 +498,7 @@ public class SessionStateTest {
   @Test
   public void testAddSessionStateWithoutPgSettings() {
     SessionState state = new SessionState(mock(OptionsMetadata.class));
-    PgCatalog pgCatalog = new PgCatalog(state);
+    PgCatalog pgCatalog = new PgCatalog(state, WellKnownClient.UNSPECIFIED);
     Statement statement = Statement.of("select * from some_table");
 
     Statement withSessionState = pgCatalog.replacePgCatalogTables(statement);
@@ -508,7 +509,7 @@ public class SessionStateTest {
   @Test
   public void testAddSessionStateWithComments() {
     SessionState state = new SessionState(mock(OptionsMetadata.class));
-    PgCatalog pgCatalog = new PgCatalog(state);
+    PgCatalog pgCatalog = new PgCatalog(state, WellKnownClient.UNSPECIFIED);
     Statement statement = Statement.of("/* This comment is preserved */ select * from pg_settings");
 
     Statement withSessionState = pgCatalog.replacePgCatalogTables(statement);
@@ -521,7 +522,7 @@ public class SessionStateTest {
   @Test
   public void testAddSessionStateWithExistingCTE() {
     SessionState state = new SessionState(mock(OptionsMetadata.class));
-    PgCatalog pgCatalog = new PgCatalog(state);
+    PgCatalog pgCatalog = new PgCatalog(state, WellKnownClient.UNSPECIFIED);
     Statement statement =
         Statement.of(
             "with my_cte as (select col1, col2 from foo) select * from pg_settings inner join my_cte on my_cte.col1=pg_settings.name");
@@ -538,7 +539,7 @@ public class SessionStateTest {
   @Test
   public void testAddSessionStateWithCommentsAndExistingCTE() {
     SessionState state = new SessionState(mock(OptionsMetadata.class));
-    PgCatalog pgCatalog = new PgCatalog(state);
+    PgCatalog pgCatalog = new PgCatalog(state, WellKnownClient.UNSPECIFIED);
     Statement statement =
         Statement.of(
             "/* This comment is preserved */ with foo as (select * from bar)\nselect * from pg_settings");
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 f149bcf6..55ca7b76 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
@@ -52,6 +52,7 @@ import com.google.cloud.spanner.pgadapter.statements.BackendConnection.QueryResu
 import com.google.cloud.spanner.pgadapter.statements.DdlExecutor.NotExecuted;
 import com.google.cloud.spanner.pgadapter.statements.local.ListDatabasesStatement;
 import com.google.cloud.spanner.pgadapter.statements.local.LocalStatement;
+import com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.WellKnownClient;
 import com.google.cloud.spanner.pgadapter.utils.CopyDataReceiver;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
 import com.google.common.collect.ImmutableList;
@@ -134,8 +135,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             spannerConnection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            ImmutableList.of());
+            ImmutableList::of);
 
     backendConnection.execute(
         PARSER.parse(Statement.of("CREATE TABLE \"Foo\" (id bigint primary key)")),
@@ -168,8 +170,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             spannerConnection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            ImmutableList.of());
+            ImmutableList::of);
     Future<StatementResult> result =
         backendConnection.executeCopy(parsedStatement, statement, receiver, writer, executor);
     backendConnection.flush();
@@ -205,8 +208,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             spannerConnection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            ImmutableList.of());
+            ImmutableList::of);
     onlyDmlStatements.execute(parsedUpdateStatement, updateStatement, Function.identity());
     onlyDmlStatements.execute(parsedUpdateStatement, updateStatement, Function.identity());
     assertTrue(onlyDmlStatements.hasDmlOrCopyStatementsAfter(0));
@@ -216,8 +220,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             spannerConnection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            ImmutableList.of());
+            ImmutableList::of);
     onlyCopyStatements.executeCopy(parsedCopyStatement, copyStatement, receiver, writer, executor);
     onlyCopyStatements.executeCopy(parsedCopyStatement, copyStatement, receiver, writer, executor);
     assertTrue(onlyCopyStatements.hasDmlOrCopyStatementsAfter(0));
@@ -227,8 +232,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             spannerConnection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            ImmutableList.of());
+            ImmutableList::of);
     dmlAndCopyStatements.execute(parsedUpdateStatement, updateStatement, Function.identity());
     dmlAndCopyStatements.executeCopy(
         parsedCopyStatement, copyStatement, receiver, writer, executor);
@@ -239,8 +245,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             spannerConnection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            ImmutableList.of());
+            ImmutableList::of);
     onlySelectStatements.execute(parsedSelectStatement, selectStatement, Function.identity());
     onlySelectStatements.execute(parsedSelectStatement, selectStatement, Function.identity());
     assertFalse(onlySelectStatements.hasDmlOrCopyStatementsAfter(0));
@@ -250,8 +257,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             spannerConnection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            ImmutableList.of());
+            ImmutableList::of);
     onlyClientSideStatements.execute(
         parsedClientSideStatement, clientSideStatement, Function.identity());
     onlyClientSideStatements.execute(
@@ -263,8 +271,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             spannerConnection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            ImmutableList.of());
+            ImmutableList::of);
     onlyUnknownStatements.execute(parsedUnknownStatement, unknownStatement, Function.identity());
     onlyUnknownStatements.execute(parsedUnknownStatement, unknownStatement, Function.identity());
     assertFalse(onlyUnknownStatements.hasDmlOrCopyStatementsAfter(0));
@@ -274,8 +283,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             spannerConnection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            ImmutableList.of());
+            ImmutableList::of);
     dmlAndSelectStatements.execute(parsedUpdateStatement, updateStatement, Function.identity());
     dmlAndSelectStatements.execute(parsedSelectStatement, selectStatement, Function.identity());
     assertTrue(dmlAndSelectStatements.hasDmlOrCopyStatementsAfter(0));
@@ -285,8 +295,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             spannerConnection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            ImmutableList.of());
+            ImmutableList::of);
     copyAndSelectStatements.executeCopy(
         parsedCopyStatement, copyStatement, receiver, writer, executor);
     copyAndSelectStatements.execute(parsedSelectStatement, selectStatement, Function.identity());
@@ -297,8 +308,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             spannerConnection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            ImmutableList.of());
+            ImmutableList::of);
     copyAndUnknownStatements.executeCopy(
         parsedCopyStatement, copyStatement, receiver, writer, executor);
     copyAndUnknownStatements.execute(parsedUnknownStatement, unknownStatement, Function.identity());
@@ -323,7 +335,11 @@ public class BackendConnectionTest {
 
     BackendConnection backendConnection =
         new BackendConnection(
-            DatabaseId.of("p", "i", "d"), connection, mock(OptionsMetadata.class), localStatements);
+            DatabaseId.of("p", "i", "d"),
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            mock(OptionsMetadata.class),
+            () -> localStatements);
     Future<StatementResult> resultFuture =
         backendConnection.execute(
             parsedListDatabasesStatement,
@@ -355,7 +371,11 @@ public class BackendConnectionTest {
 
     BackendConnection backendConnection =
         new BackendConnection(
-            DatabaseId.of("p", "i", "d"), connection, mock(OptionsMetadata.class), localStatements);
+            DatabaseId.of("p", "i", "d"),
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            mock(OptionsMetadata.class),
+            () -> localStatements);
     Future<StatementResult> resultFuture =
         backendConnection.execute(parsedStatement, statement, Function.identity());
     backendConnection.flush();
@@ -388,8 +408,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             connection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            EMPTY_LOCAL_STATEMENTS);
+            () -> EMPTY_LOCAL_STATEMENTS);
     Future<StatementResult> resultFuture =
         backendConnection.execute(parsedStatement, statement, Function.identity());
     backendConnection.flush();
@@ -412,8 +433,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             connection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            EMPTY_LOCAL_STATEMENTS);
+            () -> EMPTY_LOCAL_STATEMENTS);
     Future<StatementResult> resultFuture =
         backendConnection.execute(parsedStatement, statement, Function.identity());
     backendConnection.flush();
@@ -449,8 +471,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             connection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            EMPTY_LOCAL_STATEMENTS);
+            () -> EMPTY_LOCAL_STATEMENTS);
     Future<StatementResult> resultFuture1 =
         backendConnection.execute(parsedStatement1, statement1, Function.identity());
     backendConnection.execute(parsedStatement2, statement2, Function.identity());
@@ -474,7 +497,11 @@ public class BackendConnectionTest {
 
     BackendConnection backendConnection =
         new BackendConnection(
-            DatabaseId.of("p", "i", "d"), connection, options, EMPTY_LOCAL_STATEMENTS);
+            DatabaseId.of("p", "i", "d"),
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            options,
+            () -> EMPTY_LOCAL_STATEMENTS);
 
     backendConnection.execute(parsedStatement, statement, Function.identity());
     backendConnection.flush();
@@ -517,7 +544,11 @@ public class BackendConnectionTest {
 
     BackendConnection backendConnection =
         new BackendConnection(
-            DatabaseId.of("p", "i", "d"), connection, options, EMPTY_LOCAL_STATEMENTS);
+            DatabaseId.of("p", "i", "d"),
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            options,
+            () -> EMPTY_LOCAL_STATEMENTS);
 
     backendConnection.execute(parsedStatement, statement, Function.identity());
     backendConnection.flush();
@@ -537,8 +568,9 @@ public class BackendConnectionTest {
         new BackendConnection(
             DatabaseId.of("p", "i", "d"),
             connection,
+            () -> WellKnownClient.UNSPECIFIED,
             mock(OptionsMetadata.class),
-            EMPTY_LOCAL_STATEMENTS);
+            () -> EMPTY_LOCAL_STATEMENTS);
 
     backendConnection.execute(parsedStatement, statement, Function.identity());
     backendConnection.flush();
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/StatementTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/StatementTest.java
index 571d33dc..9e0a67fd 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/statements/StatementTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/StatementTest.java
@@ -53,6 +53,7 @@ import com.google.cloud.spanner.pgadapter.error.SQLState;
 import com.google.cloud.spanner.pgadapter.metadata.ConnectionMetadata;
 import com.google.cloud.spanner.pgadapter.metadata.OptionsMetadata;
 import com.google.cloud.spanner.pgadapter.session.SessionState;
+import com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.WellKnownClient;
 import com.google.cloud.spanner.pgadapter.utils.MutationWriter;
 import com.google.cloud.spanner.pgadapter.wireprotocol.ControlMessage;
 import com.google.cloud.spanner.pgadapter.wireprotocol.QueryMessage;
@@ -195,7 +196,11 @@ public class StatementTest {
             ImmutableList.of());
     BackendConnection backendConnection =
         new BackendConnection(
-            connectionHandler.getDatabaseId(), connection, options, EMPTY_LOCAL_STATEMENTS);
+            connectionHandler.getDatabaseId(),
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            options,
+            () -> EMPTY_LOCAL_STATEMENTS);
 
     assertFalse(intermediateStatement.isExecuted());
     assertEquals("UPDATE", intermediateStatement.getCommand());
@@ -281,7 +286,11 @@ public class StatementTest {
             ImmutableList.of());
     BackendConnection backendConnection =
         new BackendConnection(
-            connectionHandler.getDatabaseId(), connection, options, EMPTY_LOCAL_STATEMENTS);
+            connectionHandler.getDatabaseId(),
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            options,
+            () -> EMPTY_LOCAL_STATEMENTS);
 
     intermediateStatement.executeAsync(backendConnection);
     backendConnection.flush();
@@ -316,7 +325,11 @@ public class StatementTest {
             .build();
     BackendConnection backendConnection =
         new BackendConnection(
-            connectionHandler.getDatabaseId(), connection, options, EMPTY_LOCAL_STATEMENTS);
+            connectionHandler.getDatabaseId(),
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            options,
+            () -> EMPTY_LOCAL_STATEMENTS);
 
     IntermediatePreparedStatement intermediateStatement =
         new IntermediatePreparedStatement(
@@ -419,7 +432,11 @@ public class StatementTest {
             ImmutableList.of());
     BackendConnection backendConnection =
         new BackendConnection(
-            connectionHandler.getDatabaseId(), connection, options, EMPTY_LOCAL_STATEMENTS);
+            connectionHandler.getDatabaseId(),
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            options,
+            () -> EMPTY_LOCAL_STATEMENTS);
 
     when(connection.execute(Statement.of(sqlStatement)))
         .thenThrow(
@@ -481,7 +498,11 @@ public class StatementTest {
 
     BackendConnection backendConnection =
         new BackendConnection(
-            DatabaseId.of("p", "i", "d"), connection, options, ImmutableList.of());
+            DatabaseId.of("p", "i", "d"),
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            options,
+            ImmutableList::of);
     statement.executeAsync(backendConnection);
 
     ExecutorService executor = Executors.newSingleThreadExecutor();
@@ -549,7 +570,11 @@ public class StatementTest {
             ImmutableList.of());
     BackendConnection backendConnection =
         new BackendConnection(
-            connectionHandler.getDatabaseId(), connection, options, EMPTY_LOCAL_STATEMENTS);
+            connectionHandler.getDatabaseId(),
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            options,
+            () -> EMPTY_LOCAL_STATEMENTS);
 
     intermediateStatement.executeAsync(backendConnection);
 
@@ -570,7 +595,11 @@ public class StatementTest {
     setupQueryInformationSchemaResults();
     BackendConnection backendConnection =
         new BackendConnection(
-            DatabaseId.of("p", "i", "d"), connection, options, ImmutableList.of());
+            DatabaseId.of("p", "i", "d"),
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            options,
+            ImmutableList::of);
 
     byte[] payload = Files.readAllBytes(Paths.get("./src/test/resources/batch-size-test.txt"));
 
@@ -619,7 +648,11 @@ public class StatementTest {
     setupQueryInformationSchemaResults();
     BackendConnection backendConnection =
         new BackendConnection(
-            DatabaseId.of("p", "i", "d"), connection, options, ImmutableList.of());
+            DatabaseId.of("p", "i", "d"),
+            connection,
+            () -> WellKnownClient.UNSPECIFIED,
+            options,
+            ImmutableList::of);
 
     byte[] payload = "1\t'one'\n2".getBytes();