Repository: GoogleCloudPlatform/pgadapter
Pull request #768 (2023-04-08 09:02:34)

refactor: make function replacement more dynamic
Make function and expression replacement for pg_catalog queries slightly more dynamic:
1. Allow other replacements than only pattern based replacements.
2. Allow replacing to stop after having done a specific replacement, to more easily allow replacement of a complete query.

Required for Ruby-on-Rails support.

codecov[bot]: ## [Codecov](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/768?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) Report
> Merging [#768](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/768?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (7ec619b) into [postgresql-dialect](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/commit/c2aaa363a23c213074c8071dfe39c6bbd4448d96?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) (c2aaa36) will **decrease** coverage by `0.10%`.
> The diff coverage is `87.34%`.

```diff
@@                   Coverage Diff                    @@
##             postgresql-dialect     #768      +/-   ##
========================================================
- Coverage                 89.54%   89.44%   -0.10%     
- Complexity                 2354     2356       +2     
========================================================
  Files                       129      131       +2     
  Lines                      7767     7798      +31     
  Branches                   1112     1115       +3     
========================================================
+ Hits                       6955     6975      +20     
- Misses                      564      573       +9     
- Partials                    248      250       +2     
```

| Flag | Coverage Δ | |
|---|---|---|
| all_tests | `89.44% <87.34%> (-0.10%)` | :arrow_down: |

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/768?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform) | Coverage Δ | |
|---|---|---|
| [...panner/pgadapter/utils/RegexQueryPartReplacer.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/768?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL1JlZ2V4UXVlcnlQYXJ0UmVwbGFjZXIuamF2YQ==) | `66.66% <66.66%> (ø)` | |
| [.../cloud/spanner/pgadapter/statements/PgCatalog.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/768?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3N0YXRlbWVudHMvUGdDYXRhbG9nLmphdmE=) | `93.75% <87.17%> (-3.31%)` | :arrow_down: |
| [...ud/spanner/pgadapter/utils/ClientAutoDetector.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/768?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL0NsaWVudEF1dG9EZXRlY3Rvci5qYXZh) | `93.44% <100.00%> (+0.28%)` | :arrow_up: |
| [...oud/spanner/pgadapter/utils/QueryPartReplacer.java](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/768?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform#diff-c3JjL21haW4vamF2YS9jb20vZ29vZ2xlL2Nsb3VkL3NwYW5uZXIvcGdhZGFwdGVyL3V0aWxzL1F1ZXJ5UGFydFJlcGxhY2VyLmphdmE=) | `100.00% <100.00%> (ø)` | |

... and [1 file with indirect coverage changes](https://codecov.io/gh/GoogleCloudPlatform/pgadapter/pull/768/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=GoogleCloudPlatform)

: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/PgCatalog.java b/src/main/java/com/google/cloud/spanner/pgadapter/statements/PgCatalog.java
index 2ece2e5e..5b2d1ac4 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
@@ -21,6 +21,10 @@ 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.cloud.spanner.pgadapter.utils.QueryPartReplacer;
+import com.google.cloud.spanner.pgadapter.utils.QueryPartReplacer.ReplacementStatus;
+import com.google.cloud.spanner.pgadapter.utils.RegexQueryPartReplacer;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableList;
@@ -31,7 +35,6 @@ import java.util.Locale;
 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;
 
@@ -108,21 +111,23 @@ public class PgCatalog {
           .put(new TableOrIndexName(null, "pg_settings"), new TableOrIndexName(null, "pg_settings"))
           .build();
 
-  private static final ImmutableMap<Pattern, Supplier<String>> DEFAULT_FUNCTION_REPLACEMENTS =
-      ImmutableMap.of(
-          Pattern.compile("pg_catalog.pg_table_is_visible\\s*\\(.+\\)"),
-          Suppliers.ofInstance("true"),
-          Pattern.compile("pg_table_is_visible\\s*\\(.+\\)"),
-          Suppliers.ofInstance("true"),
-          Pattern.compile("=\\s*ANY\\s*\\(current_schemas\\(true\\)\\)"),
-          Suppliers.ofInstance(" IN ('pg_catalog', 'public')"));
+  private static final ImmutableList<QueryPartReplacer> DEFAULT_FUNCTION_REPLACEMENTS =
+      ImmutableList.of(
+          RegexQueryPartReplacer.replace(
+              Pattern.compile("pg_catalog.pg_table_is_visible\\s*\\(.+\\)"),
+              Suppliers.ofInstance("true")),
+          RegexQueryPartReplacer.replace(
+              Pattern.compile("pg_table_is_visible\\s*\\(.+\\)"), Suppliers.ofInstance("true")),
+          RegexQueryPartReplacer.replace(
+              Pattern.compile("=\\s*ANY\\s*\\(current_schemas\\(true\\)\\)"),
+              Suppliers.ofInstance(" IN ('pg_catalog', 'public')")));
 
   private final ImmutableSet<String> checkPrefixes;
 
   private final ImmutableMap<TableOrIndexName, TableOrIndexName> tableReplacements;
   private final ImmutableMap<TableOrIndexName, PgCatalogTable> pgCatalogTables;
 
-  private final ImmutableMap<Pattern, Supplier<String>> functionReplacements;
+  private final ImmutableList<QueryPartReplacer> functionReplacements;
 
   private static final Map<TableOrIndexName, PgCatalogTable> DEFAULT_PG_CATALOG_TABLES =
       ImmutableMap.<TableOrIndexName, PgCatalogTable>builder()
@@ -167,14 +172,21 @@ public class PgCatalog {
     this.pgCatalogTables = pgCatalogTablesBuilder.build();
 
     this.functionReplacements =
-        ImmutableMap.<Pattern, Supplier<String>>builder()
-            .putAll(DEFAULT_FUNCTION_REPLACEMENTS)
-            .put(
-                Pattern.compile("version\\(\\)"), () -> "'" + sessionState.getServerVersion() + "'")
-            .putAll(wellKnownClient.getFunctionReplacements())
+        ImmutableList.<QueryPartReplacer>builder()
+            .addAll(getDefaultFunctionReplacements())
+            .add(
+                RegexQueryPartReplacer.replace(
+                    Pattern.compile("version\\(\\)"),
+                    () -> "'" + sessionState.getServerVersion() + "'"))
+            .addAll(wellKnownClient.getQueryPartReplacements())
             .build();
   }
 
+  @VisibleForTesting
+  ImmutableList<QueryPartReplacer> getDefaultFunctionReplacements() {
+    return DEFAULT_FUNCTION_REPLACEMENTS;
+  }
+
   /** Replace supported pg_catalog tables with Common Table Expressions. */
   public Statement replacePgCatalogTables(Statement statement) {
     // Only replace tables if the statement contains at least one of the known prefixes.
@@ -199,16 +211,16 @@ public class PgCatalog {
     return addCommonTableExpressions(replacedTablesStatement.y(), cteBuilder.build());
   }
 
-  String replaceKnownUnsupportedFunctions(Statement statement) {
+  Tuple<String, ReplacementStatus> replaceKnownUnsupportedFunctions(Statement statement) {
     String sql = statement.getSql();
-    for (Entry<Pattern, Supplier<String>> functionReplacement : functionReplacements.entrySet()) {
-      sql =
-          functionReplacement
-              .getKey()
-              .matcher(sql)
-              .replaceAll(functionReplacement.getValue().get());
+    for (QueryPartReplacer functionReplacement : functionReplacements) {
+      Tuple<String, ReplacementStatus> result = functionReplacement.replace(sql);
+      if (result.y() == ReplacementStatus.STOP) {
+        return result;
+      }
+      sql = result.x();
     }
-    return sql;
+    return Tuple.of(sql, ReplacementStatus.CONTINUE);
   }
 
   Statement addCommonTableExpressions(Statement statement, ImmutableList<String> tableExpressions) {
@@ -216,22 +228,28 @@ public class PgCatalog {
       return statement;
     }
 
-    String sql = replaceKnownUnsupportedFunctions(statement);
-    SimpleParser parser = new SimpleParser(sql);
-    boolean hadCommonTableExpressions = parser.eatKeyword("with");
-    String tableExpressionsSql = String.join(",\n", tableExpressions);
-    Statement.Builder builder =
-        Statement.newBuilder("with ")
-            .append(tableExpressionsSql)
-            .append(hadCommonTableExpressions ? ",\n" : "\n");
-    if (hadCommonTableExpressions) {
-      // Include the entire original statement except the 'with' keyword.
-      builder
-          .append(parser.getSql().substring(0, parser.getPos() - 4))
-          .append(parser.getSql().substring(parser.getPos()));
+    Tuple<String, ReplacementStatus> result = replaceKnownUnsupportedFunctions(statement);
+    String sql = result.x();
+    Statement.Builder builder;
+    if (result.y() == ReplacementStatus.CONTINUE) {
+      SimpleParser parser = new SimpleParser(sql);
+      boolean hadCommonTableExpressions = parser.eatKeyword("with");
+      String tableExpressionsSql = String.join(",\n", tableExpressions);
+      builder =
+          Statement.newBuilder("with ")
+              .append(tableExpressionsSql)
+              .append(hadCommonTableExpressions ? ",\n" : "\n");
+      if (hadCommonTableExpressions) {
+        // Include the entire original statement except the 'with' keyword.
+        builder
+            .append(parser.getSql().substring(0, parser.getPos() - 4))
+            .append(parser.getSql().substring(parser.getPos()));
+      } else {
+        // Include the entire original statement (including any comments at the beginning).
+        builder.append(parser.getSql());
+      }
     } else {
-      // Include the entire original statement (including any comments at the beginning).
-      builder.append(parser.getSql());
+      builder = Statement.newBuilder(sql);
     }
     Map<String, Value> parameters = statement.getParameters();
     for (Entry<String, Value> param : parameters.entrySet()) {
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 4909e310..dc66c759 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
@@ -38,7 +38,6 @@ 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;
@@ -180,18 +179,21 @@ public class ClientAutoDetector {
       }
     },
     NPGSQL {
-      final ImmutableMap<Pattern, Supplier<String>> functionReplacements =
-          ImmutableMap.of(
-              Pattern.compile("elemproc\\.oid = elemtyp\\.typreceive"),
-              Suppliers.ofInstance("false"),
-              Pattern.compile("proc\\.oid = typ\\.typreceive"),
-              Suppliers.ofInstance("false"),
-              Pattern.compile("WHEN proc\\.proname='array_recv' THEN typ\\.typelem"),
-              Suppliers.ofInstance("WHEN substr(typ.typname, 1, 1)='_' THEN typ.typelem"),
-              Pattern.compile(
-                  "WHEN proc\\.proname='array_recv' THEN 'a' ELSE typ\\.typtype END AS typtype"),
-              Suppliers.ofInstance(
-                  "WHEN substr(typ.typname, 1, 1)='_' THEN 'a' ELSE typ.typtype END AS typtype"));
+      final ImmutableList<QueryPartReplacer> functionReplacements =
+          ImmutableList.of(
+              RegexQueryPartReplacer.replace(
+                  Pattern.compile("elemproc\\.oid = elemtyp\\.typreceive"),
+                  Suppliers.ofInstance("false")),
+              RegexQueryPartReplacer.replace(
+                  Pattern.compile("proc\\.oid = typ\\.typreceive"), Suppliers.ofInstance("false")),
+              RegexQueryPartReplacer.replace(
+                  Pattern.compile("WHEN proc\\.proname='array_recv' THEN typ\\.typelem"),
+                  Suppliers.ofInstance("WHEN substr(typ.typname, 1, 1)='_' THEN typ.typelem")),
+              RegexQueryPartReplacer.replace(
+                  Pattern.compile(
+                      "WHEN proc\\.proname='array_recv' THEN 'a' ELSE typ\\.typtype END AS typtype"),
+                  Suppliers.ofInstance(
+                      "WHEN substr(typ.typname, 1, 1)='_' THEN 'a' ELSE typ.typtype END AS typtype")));
 
       @Override
       boolean isClient(List<String> orderedParameterKeys, Map<String, String> parameters) {
@@ -214,17 +216,19 @@ public class ClientAutoDetector {
       }
 
       @Override
-      public ImmutableMap<Pattern, Supplier<String>> getFunctionReplacements() {
+      public ImmutableList<QueryPartReplacer> getQueryPartReplacements() {
         return functionReplacements;
       }
     },
     SQLALCHEMY2 {
-      final ImmutableMap<Pattern, Supplier<String>> functionReplacements =
-          ImmutableMap.of(
-              Pattern.compile("oid::regtype::text AS regtype"),
-              Suppliers.ofInstance("'' as regtype"),
-              Pattern.compile("WHERE t\\.oid = to_regtype\\(\\$1\\)"),
-              Suppliers.ofInstance("WHERE t.typname = \\$1"));
+      final ImmutableList<QueryPartReplacer> functionReplacements =
+          ImmutableList.of(
+              RegexQueryPartReplacer.replace(
+                  Pattern.compile("oid::regtype::text AS regtype"),
+                  Suppliers.ofInstance("'' as regtype")),
+              RegexQueryPartReplacer.replace(
+                  Pattern.compile("WHERE t\\.oid = to_regtype\\(\\$1\\)"),
+                  Suppliers.ofInstance("WHERE t.typname = \\$1")));
 
       @Override
       boolean isClient(List<String> orderedParameterKeys, Map<String, String> parameters) {
@@ -244,7 +248,7 @@ public class ClientAutoDetector {
       }
 
       @Override
-      public ImmutableMap<Pattern, Supplier<String>> getFunctionReplacements() {
+      public ImmutableList<QueryPartReplacer> getQueryPartReplacements() {
         return functionReplacements;
       }
     },
@@ -304,8 +308,8 @@ public class ClientAutoDetector {
       return ImmutableMap.of();
     }
 
-    public ImmutableMap<Pattern, Supplier<String>> getFunctionReplacements() {
-      return ImmutableMap.of();
+    public ImmutableList<QueryPartReplacer> getQueryPartReplacements() {
+      return ImmutableList.of();
     }
 
     /** Creates specific notice messages for a client after startup. */
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/QueryPartReplacer.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/QueryPartReplacer.java
new file mode 100644
index 00000000..a400ce6f
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/QueryPartReplacer.java
@@ -0,0 +1,28 @@
+// Copyright 2023 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.pgadapter.utils;
+
+import com.google.cloud.Tuple;
+
+/** Interface for dynamically replacing a part of a SQL statement. */
+public interface QueryPartReplacer {
+  enum ReplacementStatus {
+    CONTINUE,
+    STOP;
+  }
+
+  /** Replace unsupported SQL parts in the given SQL string and return the new SQL string. */
+  Tuple<String, ReplacementStatus> replace(String sql);
+}
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/utils/RegexQueryPartReplacer.java b/src/main/java/com/google/cloud/spanner/pgadapter/utils/RegexQueryPartReplacer.java
new file mode 100644
index 00000000..614e4704
--- /dev/null
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/utils/RegexQueryPartReplacer.java
@@ -0,0 +1,78 @@
+// Copyright 2023 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.pgadapter.utils;
+
+import com.google.cloud.Tuple;
+import com.google.common.base.Suppliers;
+import java.util.function.Supplier;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+/**
+ * {@link QueryPartReplacer} that replaces a regex {@link java.util.regex.Pattern} with a fixed
+ * string.
+ */
+public class RegexQueryPartReplacer implements QueryPartReplacer {
+  private final Pattern pattern;
+  private final Supplier<String> replacement;
+  private final ReplacementStatus replacementStatus;
+
+  /** Replace all occurrences of the given pattern with the given static replacement. */
+  public static RegexQueryPartReplacer replace(Pattern pattern, String replacement) {
+    return new RegexQueryPartReplacer(
+        pattern, Suppliers.ofInstance(replacement), ReplacementStatus.CONTINUE);
+  }
+
+  /**
+   * Replace all occurrences of the given pattern with the value that is returned by the given
+   * {@link Supplier} function.
+   */
+  public static RegexQueryPartReplacer replace(Pattern pattern, Supplier<String> replacement) {
+    return new RegexQueryPartReplacer(pattern, replacement, ReplacementStatus.CONTINUE);
+  }
+
+  /**
+   * Replace the first occurrence of the given pattern with the given static value and then stop.
+   */
+  public static RegexQueryPartReplacer replaceAndStop(Pattern pattern, String replacement) {
+    return new RegexQueryPartReplacer(
+        pattern, Suppliers.ofInstance(replacement), ReplacementStatus.STOP);
+  }
+
+  /**
+   * Replace the first occurrence of the given pattern with the value that is returned by the given
+   * {@link Supplier} function and then stop.
+   */
+  public static RegexQueryPartReplacer replaceAndStop(
+      Pattern pattern, Supplier<String> replacement) {
+    return new RegexQueryPartReplacer(pattern, replacement, ReplacementStatus.STOP);
+  }
+
+  private RegexQueryPartReplacer(
+      Pattern pattern, Supplier<String> replacement, ReplacementStatus replacementStatus) {
+    this.pattern = pattern;
+    this.replacement = replacement;
+    this.replacementStatus = replacementStatus;
+  }
+
+  @Override
+  public Tuple<String, ReplacementStatus> replace(String sql) {
+    Matcher matcher = pattern.matcher(sql);
+    if (matcher.find()) {
+      return Tuple.of(matcher.replaceAll(replacement.get()), replacementStatus);
+    }
+    return Tuple.of(sql, ReplacementStatus.CONTINUE);
+  }
+}


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 5888e540..4c99de89 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -559,7 +559,7 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
         new String[] {"SELECT version()", "select version()", "select * from version()"}) {
 
       try (Connection connection = DriverManager.getConnection(createUrl())) {
-        String version = null;
+        String version;
         try (ResultSet resultSet =
             connection.createStatement().executeQuery("show server_version")) {
           assertTrue(resultSet.next());
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/statements/PgCatalogTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/statements/PgCatalogTest.java
index 573204ec..d5964589 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/statements/PgCatalogTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/statements/PgCatalogTest.java
@@ -14,12 +14,17 @@
 
 package com.google.cloud.spanner.pgadapter.statements;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertSame;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
+import com.google.cloud.Tuple;
 import com.google.cloud.spanner.Statement;
 import com.google.cloud.spanner.pgadapter.session.SessionState;
 import com.google.cloud.spanner.pgadapter.utils.ClientAutoDetector.WellKnownClient;
+import com.google.cloud.spanner.pgadapter.utils.QueryPartReplacer;
+import com.google.cloud.spanner.pgadapter.utils.QueryPartReplacer.ReplacementStatus;
 import com.google.common.collect.ImmutableList;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -36,4 +41,69 @@ public class PgCatalogTest {
     PgCatalog catalog = new PgCatalog(mock(SessionState.class), WellKnownClient.UNSPECIFIED);
     assertSame(statement, catalog.addCommonTableExpressions(statement, ImmutableList.of()));
   }
+
+  @Test
+  public void testSelectVersion() {
+    Statement statement = Statement.of("select version(), pg_type.* from pg_type");
+    SessionState sessionState = mock(SessionState.class);
+    PgCatalog catalog = new PgCatalog(sessionState, WellKnownClient.UNSPECIFIED);
+
+    when(sessionState.getServerVersion()).thenReturn("14.1");
+    assertEquals(
+        Statement.of("with pg_type\n" + "select '14.1', pg_type.* from pg_type"),
+        catalog.addCommonTableExpressions(statement, ImmutableList.of("pg_type")));
+
+    when(sessionState.getServerVersion()).thenReturn("9.4");
+    assertEquals(
+        Statement.of("with pg_type\n" + "select '9.4', pg_type.* from pg_type"),
+        catalog.addCommonTableExpressions(statement, ImmutableList.of("pg_type")));
+  }
+
+  @Test
+  public void testReplaceKnownUnsupportedFunctions() {
+    SessionState sessionState = mock(SessionState.class);
+    PgCatalog catalog =
+        new PgCatalog(sessionState, WellKnownClient.UNSPECIFIED) {
+          @Override
+          ImmutableList<QueryPartReplacer> getDefaultFunctionReplacements() {
+            ImmutableList<QueryPartReplacer> replacers = super.getDefaultFunctionReplacements();
+
+            return ImmutableList.<QueryPartReplacer>builder()
+                .addAll(replacers)
+                .add(
+                    sql -> {
+                      if ("select * from replace_and_stop".equals(sql)) {
+                        return Tuple.of("select * from replaced", ReplacementStatus.STOP);
+                      }
+                      return Tuple.of(sql, ReplacementStatus.CONTINUE);
+                    })
+                .build();
+          }
+        };
+
+    Tuple<String, ReplacementStatus> result;
+    result = catalog.replaceKnownUnsupportedFunctions(Statement.of("select * from foo"));
+    assertEquals("select * from foo", result.x());
+    assertEquals(ReplacementStatus.CONTINUE, result.y());
+
+    result =
+        catalog.replaceKnownUnsupportedFunctions(Statement.of("select pg_table_is_visible('foo')"));
+    assertEquals("select true", result.x());
+    assertEquals(ReplacementStatus.CONTINUE, result.y());
+
+    result =
+        catalog.replaceKnownUnsupportedFunctions(Statement.of("select pg_table_is_visible('foo')"));
+    assertEquals("select true", result.x());
+    assertEquals(ReplacementStatus.CONTINUE, result.y());
+
+    result =
+        catalog.replaceKnownUnsupportedFunctions(Statement.of("select * from replace_and_stop"));
+    assertEquals("select * from replaced", result.x());
+    assertEquals(ReplacementStatus.STOP, result.y());
+
+    assertEquals(
+        Statement.of("select * from replaced"),
+        catalog.addCommonTableExpressions(
+            Statement.of("select * from replace_and_stop"), ImmutableList.of("replace_and_stop")));
+  }
 }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/utils/RegexQueryPartReplacerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/utils/RegexQueryPartReplacerTest.java
new file mode 100644
index 00000000..b54fc293
--- /dev/null
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/utils/RegexQueryPartReplacerTest.java
@@ -0,0 +1,88 @@
+// Copyright 2023 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.pgadapter.utils;
+
+import static org.junit.Assert.assertEquals;
+
+import com.google.cloud.Tuple;
+import com.google.cloud.spanner.pgadapter.utils.QueryPartReplacer.ReplacementStatus;
+import com.google.common.base.Suppliers;
+import java.util.regex.Pattern;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class RegexQueryPartReplacerTest {
+
+  @Test
+  public void testReplaceString() {
+    QueryPartReplacer replacer =
+        RegexQueryPartReplacer.replace(Pattern.compile("foo\\(\\)"), "bar");
+    assertEquals("select bar", replacer.replace("select foo()").x());
+  }
+
+  @Test
+  public void testReplaceSupplier() {
+    for (String replacement : new String[] {"bar1", "bar2"}) {
+      QueryPartReplacer replacer =
+          RegexQueryPartReplacer.replace(
+              Pattern.compile("foo\\(\\)"), Suppliers.ofInstance(replacement));
+      assertEquals("select " + replacement, replacer.replace("select foo()").x());
+    }
+  }
+
+  @Test
+  public void testReplaceAndStopString() {
+    QueryPartReplacer replacer =
+        RegexQueryPartReplacer.replaceAndStop(Pattern.compile("foo\\(\\)"), "bar");
+    Tuple<String, ReplacementStatus> result = replacer.replace("select foo()");
+    assertEquals("select bar", result.x());
+    assertEquals(ReplacementStatus.STOP, result.y());
+  }
+
+  @Test
+  public void testReplaceAndStopSupplier() {
+    for (String replacement : new String[] {"bar1", "bar2"}) {
+      QueryPartReplacer replacer =
+          RegexQueryPartReplacer.replaceAndStop(
+              Pattern.compile("foo\\(\\)"), Suppliers.ofInstance(replacement));
+      Tuple<String, ReplacementStatus> result = replacer.replace("select foo()");
+      assertEquals("select " + replacement, result.x());
+      assertEquals(ReplacementStatus.STOP, result.y());
+    }
+  }
+
+  @Test
+  public void testReplaceAndStopNotMatchedString() {
+    QueryPartReplacer replacer =
+        RegexQueryPartReplacer.replaceAndStop(Pattern.compile("foo\\(\\)"), "bar");
+    Tuple<String, ReplacementStatus> result = replacer.replace("select baz()");
+    assertEquals("select baz()", result.x());
+    assertEquals(ReplacementStatus.CONTINUE, result.y());
+  }
+
+  @Test
+  public void testReplaceAndStopNotMatchedSupplier() {
+    for (String replacement : new String[] {"bar1", "bar2"}) {
+      QueryPartReplacer replacer =
+          RegexQueryPartReplacer.replaceAndStop(
+              Pattern.compile("foo\\(\\)"), Suppliers.ofInstance(replacement));
+      Tuple<String, ReplacementStatus> result = replacer.replace("select baz()");
+      assertEquals("select baz()", result.x());
+      assertEquals(ReplacementStatus.CONTINUE, result.y());
+    }
+  }
+}