Update ChangeStreamDao to query differnet TVF for postgresSQL based on

chenxuesdu · chenxuesdu · commit dbd78182c0b8 · 2025-10-29T16:42:09.000Z
the change stream partition mode

For MUTABLE_KEY_RANGE change stream, use read_proto_bytes_, else use read_json_
diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamDao.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamDao.java
@@ -22,12 +22,14 @@
 import com.google.cloud.spanner.Dialect;
 import com.google.cloud.spanner.Options;
 import com.google.cloud.spanner.Options.RpcPriority;
+import com.google.cloud.spanner.ReadOnlyTransaction;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.Statement;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.InitialPartition;
 
 /**
- * Responsible for making change stream queries for a given partition. The result will be given back
+ * Responsible for making change stream queries for a given partition. The
+ * result will be given back
  * as a {@link ResultSet}, which can be consumed until the stream is finished.
  */
 public class ChangeStreamDao {
@@ -37,16 +39,19 @@ public class ChangeStreamDao {
   private final RpcPriority rpcPriority;
   private final String jobName;
   private final Dialect dialect;
+  private final boolean isMutableKeyRange;
 
   /**
-   * Constructs a change stream dao. All the queries performed by this class will be for the given
-   * change stream name with the specified rpc priority. The job name will be used to tag all the
+   * Constructs a change stream dao. All the queries performed by this class will
+   * be for the given
+   * change stream name with the specified rpc priority. The job name will be used
+   * to tag all the
    * queries made.
    *
    * @param changeStreamName the name of the change stream to be queried
-   * @param databaseClient a spanner {@link DatabaseClient}
-   * @param rpcPriority the priority to be used for the change stream queries
-   * @param jobName the name of the job performing the query
+   * @param databaseClient   a spanner {@link DatabaseClient}
+   * @param rpcPriority      the priority to be used for the change stream queries
+   * @param jobName          the name of the job performing the query
    */
   ChangeStreamDao(
       String changeStreamName,
@@ -59,83 +64,144 @@ public class ChangeStreamDao {
     this.rpcPriority = rpcPriority;
     this.jobName = jobName;
     this.dialect = dialect;
+    this.isMutableKeyRange = isMutableKeyRangeChangeStream(databaseClient, dialect, changeStreamName);
   }
 
   /**
-   * Performs a change stream query. If the partition token given is the initial partition null will
-   * be used in the query instead. The change stream query will be tagged as following: {@code
-   * "job=<jobName>"}. The result will be given as a {@link ChangeStreamResultSet} which can be
-   * consumed as a stream, yielding records until no more are available for the query made. Note
-   * that one needs to call {@link ChangeStreamResultSet#next()} to initiate the change stream
+   * Performs a change stream query. If the partition token given is the initial
+   * partition null will
+   * be used in the query instead. The change stream query will be tagged as
+   * following: {@code
+   * "job=<jobName>"}. The result will be given as a {@link ChangeStreamResultSet}
+   * which can be
+   * consumed as a stream, yielding records until no more are available for the
+   * query made. Note
+   * that one needs to call {@link ChangeStreamResultSet#next()} to initiate the
+   * change stream
    * query.
    *
-   * @param partitionToken the unique partition token to be queried. If {@link
-   *     InitialPartition#PARTITION_TOKEN} is given, null will be used in the change stream query
-   *     instead.
-   * @param startTimestamp the inclusive start time for the change stream query
-   * @param endTimestamp the inclusive end time for the change stream query
-   * @param heartbeatMillis the number of milliseconds after the stream is idle, which a heartbeat
-   *     record will be emitted in the change stream query
-   * @return a {@link ChangeStreamResultSet} that will produce a stream of records for the change
-   *     stream query
+   * @param partitionToken  the unique partition token to be queried. If {@link
+   *                        InitialPartition#PARTITION_TOKEN} is given, null will
+   *                        be used in the change stream query
+   *                        instead.
+   * @param startTimestamp  the inclusive start time for the change stream query
+   * @param endTimestamp    the inclusive end time for the change stream query
+   * @param heartbeatMillis the number of milliseconds after the stream is idle,
+   *                        which a heartbeat
+   *                        record will be emitted in the change stream query
+   * @return a {@link ChangeStreamResultSet} that will produce a stream of records
+   *         for the change
+   *         stream query
    */
   public ChangeStreamResultSet changeStreamQuery(
       String partitionToken,
       Timestamp startTimestamp,
       Timestamp endTimestamp,
       long heartbeatMillis) {
     // For the initial partition we query with a null partition token
-    final String partitionTokenOrNull =
-        InitialPartition.isInitialPartition(partitionToken) ? null : partitionToken;
+    final String partitionTokenOrNull = InitialPartition.isInitialPartition(partitionToken) ? null : partitionToken;
 
     String query = "";
     Statement statement;
     if (this.isPostgres()) {
-      query =
-          "SELECT * FROM \"spanner\".\"read_json_" + changeStreamName + "\"($1, $2, $3, $4, null)";
-      statement =
-          Statement.newBuilder(query)
-              .bind("p1")
-              .to(startTimestamp)
-              .bind("p2")
-              .to(endTimestamp)
-              .bind("p3")
-              .to(partitionTokenOrNull)
-              .bind("p4")
-              .to(heartbeatMillis)
-              .build();
+      if (this.isMutableKeyRange) {
+        query = "SELECT * FROM \"spanner\".\"read_proto_bytes_"
+            + changeStreamName
+            + "\"($1, $2, $3, $4, null)";
+      } else {
+        query = "SELECT * FROM \"spanner\".\"read_json_"
+            + changeStreamName
+            + "\"($1, $2, $3, $4, null)";
+      }
+      statement = Statement.newBuilder(query)
+          .bind("p1")
+          .to(startTimestamp)
+          .bind("p2")
+          .to(endTimestamp)
+          .bind("p3")
+          .to(partitionTokenOrNull)
+          .bind("p4")
+          .to(heartbeatMillis)
+          .build();
     } else {
-      query =
-          "SELECT * FROM READ_"
-              + changeStreamName
-              + "("
-              + "   start_timestamp => @startTimestamp,"
-              + "   end_timestamp => @endTimestamp,"
-              + "   partition_token => @partitionToken,"
-              + "   read_options => null,"
-              + "   heartbeat_milliseconds => @heartbeatMillis"
-              + ")";
-      statement =
-          Statement.newBuilder(query)
-              .bind("startTimestamp")
-              .to(startTimestamp)
-              .bind("endTimestamp")
-              .to(endTimestamp)
-              .bind("partitionToken")
-              .to(partitionTokenOrNull)
-              .bind("heartbeatMillis")
-              .to(heartbeatMillis)
-              .build();
+      query = "SELECT * FROM READ_"
+          + changeStreamName
+          + "("
+          + "   start_timestamp => @startTimestamp,"
+          + "   end_timestamp => @endTimestamp,"
+          + "   partition_token => @partitionToken,"
+          + "   read_options => null,"
+          + "   heartbeat_milliseconds => @heartbeatMillis"
+          + ")";
+      statement = Statement.newBuilder(query)
+          .bind("startTimestamp")
+          .to(startTimestamp)
+          .bind("endTimestamp")
+          .to(endTimestamp)
+          .bind("partitionToken")
+          .to(partitionTokenOrNull)
+          .bind("heartbeatMillis")
+          .to(heartbeatMillis)
+          .build();
     }
-    final ResultSet resultSet =
-        databaseClient
-            .singleUse()
-            .executeQuery(statement, Options.priority(rpcPriority), Options.tag("job=" + jobName));
+    final ResultSet resultSet = databaseClient
+        .singleUse()
+        .executeQuery(statement, Options.priority(rpcPriority), Options.tag("job=" + jobName));
 
     return new ChangeStreamResultSet(resultSet);
   }
 
   private boolean isPostgres() {
     return this.dialect == Dialect.POSTGRESQL;
   }
+
+  protected boolean isMutableKeyRangeChangeStream() {
+    return this.isMutableKeyRange;
+  }
+
+  // Returns if the change stream with the give name is MUTABLE_KEY_RANGE or not.
+  private static boolean isMutableKeyRangeChangeStream(
+      DatabaseClient databaseClient, Dialect dialect, String changeStreamName) {
+    try (ReadOnlyTransaction tx = databaseClient.readOnlyTransaction()) {
+      ResultSet resultSet = readChangeStreamOptions(tx, dialect, changeStreamName);
+
+      while (resultSet.next()) {
+        String optionName = resultSet.getString(0);
+        if ("partition_mode".equalsIgnoreCase(optionName)) {
+          String optionValue = resultSet.getString(1);
+          return "MUTABLE_KEY_RANGE".equalsIgnoreCase(optionValue);
+        }
+      }
+    }
+    return false;
+  }
+
+  private static ResultSet readChangeStreamOptions(
+      ReadOnlyTransaction tx, Dialect dialect, String changeStreamName) {
+    Statement statement;
+    if (dialect == Dialect.POSTGRESQL) {
+      statement = Statement.newBuilder(
+          "select"
+              + "  option_name,"
+              + "  option_value\n"
+              + "from"
+              + "  information_schema.change_stream_options\n"
+              + "where change_stream_name = $1")
+          .bind("p1")
+          .to(changeStreamName)
+          .build();
+    } else {
+      statement = Statement.newBuilder(
+          "select"
+              + "  option_name,"
+              + "  option_value\n"
+              + "from"
+              + "  information_schema.change_stream_options\n"
+              + "where change_stream_name = @changeStreamName")
+          .bind("changeStreamName")
+          .to(changeStreamName)
+          .build();
+    }
+    return tx.executeQuery(statement);
+  }
 }
diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamDaoTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamDaoTest.java
@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 org.apache.beam.sdk.io.gcp.spanner.changestreams.dao;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.google.cloud.spanner.DatabaseClient;
+import com.google.cloud.spanner.Dialect;
+import com.google.cloud.spanner.Options.RpcPriority;
+import com.google.cloud.spanner.ReadOnlyTransaction;
+import com.google.cloud.spanner.ResultSet;
+import org.junit.Before;
+import org.junit.Test;
+
+public class ChangeStreamDaoTest {
+  private DatabaseClient databaseClient;
+  private RpcPriority rpcPriority;
+  private ChangeStreamDao changeStreamDao;
+  private static final String CHANGE_STREAM_NAME = "testCS";
+
+  @Before
+  public void setUp() {
+    databaseClient = mock(DatabaseClient.class);
+    rpcPriority = mock(RpcPriority.class);
+    changeStreamDao =
+        new ChangeStreamDao(
+            CHANGE_STREAM_NAME,
+            databaseClient,
+            rpcPriority,
+            "testjob",
+            Dialect.GOOGLE_STANDARD_SQL);
+  }
+
+  @Test
+  public void testPartitionOption() {
+    ReadOnlyTransaction readOnlyTransaction = mock(ReadOnlyTransaction.class);
+    when(databaseClient.readOnlyTransaction()).thenReturn(readOnlyTransaction);
+
+    ResultSet resultSet = mock(ResultSet.class);
+    when(readOnlyTransaction.executeQuery(any())).thenReturn(resultSet);
+    when(resultSet.next()).thenReturn(true).thenReturn(false);
+    when(resultSet.getString(0)).thenReturn("partition_mode");
+
+    when(resultSet.getString(1)).thenReturn("MUTABLE_KEY_RANGE");
+    assertEquals(true, changeStreamDao.isMutableKeyRangeChangeStream());
+
+    when(resultSet.getString(1)).thenReturn("IMMUTABLE_KEY_RANGE");
+    assertEquals(false, changeStreamDao.isMutableKeyRangeChangeStream());
+  }
+}
