Add mutable change stream in IntegrateionTestEnv.java for postgresql (#38057)

* Add mutable change stream in IntegrateionTestEnv.java for postgresql

* Retriggering checks

* Fix tests

Author: chenxuesdu

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/it/IntegrationTestEnv.java

@@ -64,7 +64,8 @@ public class IntegrationTestEnv extends ExternalResource {
   private DatabaseAdminClient databaseAdminClient;
   private DatabaseClient databaseClient;
   private boolean isPostgres;
-  private boolean isPlacementTableBasedChangeStream;
+  private boolean isMutableChangeStream;
+  private boolean isPlacementTable;
   public boolean useSeparateMetadataDb;
 
   @Override
@@ -100,13 +101,18 @@ protected void before() throws Throwable {
 
   IntegrationTestEnv() {
     this.isPostgres = false;
-    this.isPlacementTableBasedChangeStream = false;
+    this.isMutableChangeStream = false;
+    this.isPlacementTable = false;
   }
 
   IntegrationTestEnv(
-      boolean isPostgres, boolean isPlacementTableBasedChangeStream, Optional<String> host) {
+      boolean isPostgres,
+      boolean isMutableChangeStream,
+      boolean isPlacementTable,
+      Optional<String> host) {
     this.isPostgres = isPostgres;
-    this.isPlacementTableBasedChangeStream = isPlacementTableBasedChangeStream;
+    this.isMutableChangeStream = isMutableChangeStream;
+    this.isPlacementTable = isPlacementTable;
     if (host.isPresent()) {
       this.host = host.get();
     }
@@ -155,7 +161,7 @@ protected void after() {
               .get(TIMEOUT_MINUTES, TimeUnit.MINUTES);
         }
       } catch (Exception e) {
-        if (isPlacementTableBasedChangeStream) {
+        if (isPlacementTable) {
           // Drop placement table requires all rows deleted and garbage collected.
           LOG.info("Failed to drop table {}. Skipping...", table, e);
         } else {
@@ -214,7 +220,7 @@ String createSingersTable() throws InterruptedException, ExecutionException, Tim
   }
 
   String createGSQLTableDDL(String tableName) {
-    if (this.isPlacementTableBasedChangeStream) {
+    if (this.isPlacementTable) {
       // create a placement table.
       return "CREATE TABLE "
           + tableName
@@ -245,8 +251,7 @@ String createChangeStreamFor(String tableName)
           .updateDatabaseDdl(
               instanceId,
               databaseId,
-              Collections.singletonList(
-                  "CREATE CHANGE STREAM \"" + changeStreamName + "\" FOR \"" + tableName + "\""),
+              Collections.singletonList(createPostgresChangeStreamDDL(changeStreamName, tableName)),
               null)
           .get(TIMEOUT_MINUTES, TimeUnit.MINUTES);
     } else {
@@ -263,7 +268,7 @@ String createChangeStreamFor(String tableName)
   }
 
   String createGSQLChangeStreamDDL(String changeStreamName, String tableName) {
-    if (this.isPlacementTableBasedChangeStream) {
+    if (this.isMutableChangeStream) {
       // Create a MUTABLE_KEY_RANGE change stream.
       String statement =
           "CREATE CHANGE STREAM "
@@ -276,6 +281,21 @@ String createGSQLChangeStreamDDL(String changeStreamName, String tableName) {
     return "CREATE CHANGE STREAM " + changeStreamName + " FOR " + tableName;
   }
 
+  String createPostgresChangeStreamDDL(String changeStreamName, String tableName) {
+    if (this.isMutableChangeStream) {
+      // Create a MUTABLE_KEY_RANGE change stream.
+      String statement =
+          "CREATE CHANGE STREAM \""
+              + changeStreamName
+              + "\" FOR \""
+              + tableName
+              + "\""
+              + " WITH (partition_mode = 'MUTABLE_KEY_RANGE')";
+      return statement;
+    }
+    return "CREATE CHANGE STREAM \"" + changeStreamName + "\" FOR \"" + tableName + "\"";
+  }
+
   void createRoleAndGrantPrivileges(String table, String changeStream)
       throws InterruptedException, ExecutionException, TimeoutException {
     if (this.isPostgres) {

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/it/SpannerChangeStreamPlacementTableIT.java

@@ -71,7 +71,8 @@ public class SpannerChangeStreamPlacementTableIT {
   public static final IntegrationTestEnv ENV =
       new IntegrationTestEnv(
           /*isPostgres=*/ false,
-          /*isPlacementTableBasedChangeStream=*/ true,
+          /*isMutableChangeStream=*/ true,
+          /*isPlacementTable=*/ true,
           /*host=*/ Optional.empty());
 
   @Rule public final transient TestPipeline pipeline = TestPipeline.create();

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/it/SpannerChangeStreamPlacementTablePostgresIT.java

@@ -65,7 +65,8 @@ public class SpannerChangeStreamPlacementTablePostgresIT {
   public static final IntegrationTestEnv ENV =
       new IntegrationTestEnv(
           /*isPostgres=*/ true,
-          /*isPlacementTableBasedChangeStream=*/ true,
+          /*isMutableChangeStream=*/ true,
+          /*isPlacementTable=*/ true,
           /*host=*/ Optional.empty());
 
   @Rule public final transient TestPipeline pipeline = TestPipeline.create();

sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/it/SpannerChangeStreamPostgresIT.java

@@ -65,7 +65,8 @@ public class SpannerChangeStreamPostgresIT {
   public static final IntegrationTestEnv ENV =
       new IntegrationTestEnv(
           /*isPostgres=*/ true,
-          /*isPlacementTableBasedChangeStream=*/ false,
+          /*isMutableChangeStream=*/ false,
+          /*isPlacementTable=*/ false,
           /*host=*/ Optional.empty());
 
   @Rule public final transient TestPipeline pipeline = TestPipeline.create();