Update ChangeStreamDao to query different TVF for postgresSQL based on

the change stream partition mode

For MUTABLE_KEY_RANGE change stream, use read_proto_bytes_, else use read_json_

Author: chenxuesdu

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/SpannerIO.java

@@ -46,6 +46,8 @@
 import com.google.cloud.spanner.Options;
 import com.google.cloud.spanner.Options.RpcPriority;
 import com.google.cloud.spanner.PartitionOptions;
+import com.google.cloud.spanner.ReadOnlyTransaction;
+import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.Spanner;
 import com.google.cloud.spanner.SpannerException;
 import com.google.cloud.spanner.SpannerOptions;
@@ -1998,6 +2000,11 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta
       final MapperFactory mapperFactory = new MapperFactory(changeStreamDatabaseDialect);
       final ChangeStreamMetrics metrics = new ChangeStreamMetrics();
       final RpcPriority rpcPriority = MoreObjects.firstNonNull(getRpcPriority(), RpcPriority.HIGH);
+      final SpannerAccessor spannerAccessor =
+          SpannerAccessor.getOrCreate(changeStreamSpannerConfig);
+      final boolean isMutableChangeStream =
+          isMutableChangeStream(
+              spannerAccessor.getDatabaseClient(), changeStreamDatabaseDialect, changeStreamName);
       final DaoFactory daoFactory =
           new DaoFactory(
               changeStreamSpannerConfig,
@@ -2007,7 +2014,8 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta
               rpcPriority,
               input.getPipeline().getOptions().getJobName(),
               changeStreamDatabaseDialect,
-              metadataDatabaseDialect);
+              metadataDatabaseDialect,
+              isMutableChangeStream);
       final ActionFactory actionFactory = new ActionFactory();
 
       final Duration watermarkRefreshRate =
@@ -2688,4 +2696,58 @@ static String resolveSpannerProjectId(SpannerConfig config) {
         ? SpannerOptions.getDefaultProjectId()
         : config.getProjectId().get();
   }
+
+  @VisibleForTesting
+  static boolean isMutableChangeStream(
+      DatabaseClient databaseClient, Dialect dialect, String changeStreamName) {
+    String fetchedPartitionMode = fetchPartitionMode(databaseClient, dialect, changeStreamName);
+    if (fetchedPartitionMode.isEmpty()
+        || fetchedPartitionMode.equalsIgnoreCase("IMMUTABLE_KEY_RANGE")) {
+      return false;
+    }
+    return true;
+  }
+
+  private static String fetchPartitionMode(
+      DatabaseClient databaseClient, Dialect dialect, String changeStreamName) {
+    try (ReadOnlyTransaction tx = databaseClient.readOnlyTransaction()) {
+      Statement statement;
+      if (dialect == Dialect.POSTGRESQL) {
+        statement =
+            Statement.newBuilder(
+                    "select option_value\n"
+                        + "from information_schema.change_stream_options\n"
+                        + "where change_stream_name = $1 and option_name = 'partition_mode'")
+                .bind("p1")
+                .to(changeStreamName)
+                .build();
+      } else {
+        statement =
+            Statement.newBuilder(
+                    "select option_value\n"
+                        + "from information_schema.change_stream_options\n"
+                        + "where change_stream_name = @changeStreamName and  option_name = 'partition_mode'")
+                .bind("changeStreamName")
+                .to(changeStreamName)
+                .build();
+      }
+      ResultSet resultSet = tx.executeQuery(statement);
+      while (resultSet.next()) {
+        String value = resultSet.getString(0);
+        if (value != null) {
+          return value;
+        }
+      }
+      return "";
+    } catch (RuntimeException e) {
+      // Log the failure (with stack trace) but rethrow so the caller still observes
+      // the error.
+      LOG.warn(
+          "Failed to fetch partition_mode for change stream '{}', dialect={} - will propagate exception",
+          changeStreamName,
+          dialect,
+          e);
+      throw e;
+    }
+  }
 }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/ChangeStreamDao.java

@@ -31,12 +31,12 @@
  * as a {@link ResultSet}, which can be consumed until the stream is finished.
  */
 public class ChangeStreamDao {
-
   private final String changeStreamName;
   private final DatabaseClient databaseClient;
   private final RpcPriority rpcPriority;
   private final String jobName;
   private final Dialect dialect;
+  private final boolean isMutableChangeStream;
 
   /**
    * Constructs a change stream dao. All the queries performed by this class will be for the given
@@ -53,12 +53,14 @@ public class ChangeStreamDao {
       DatabaseClient databaseClient,
       RpcPriority rpcPriority,
       String jobName,
-      Dialect dialect) {
+      Dialect dialect,
+      boolean isMutableChangeStream) {
     this.changeStreamName = changeStreamName;
     this.databaseClient = databaseClient;
     this.rpcPriority = rpcPriority;
     this.jobName = jobName;
     this.dialect = dialect;
+    this.isMutableChangeStream = isMutableChangeStream;
   }
 
   /**
@@ -91,8 +93,18 @@ public ChangeStreamResultSet changeStreamQuery(
     String query = "";
     Statement statement;
     if (this.isPostgres()) {
-      query =
-          "SELECT * FROM \"spanner\".\"read_json_" + changeStreamName + "\"($1, $2, $3, $4, null)";
+      // Ensure we have determined whether change stream uses mutable key range
+      if (this.isMutableChangeStream) {
+        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")

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/dao/DaoFactory.java

@@ -49,6 +49,7 @@ public class DaoFactory implements Serializable {
   private final String jobName;
   private final Dialect spannerChangeStreamDatabaseDialect;
   private final Dialect metadataDatabaseDialect;
+  private final boolean isMutableChangeStream;
 
   /**
    * Constructs a {@link DaoFactory} with the configuration to be used for the underlying instances.
@@ -68,7 +69,8 @@ public DaoFactory(
       RpcPriority rpcPriority,
       String jobName,
       Dialect spannerChangeStreamDatabaseDialect,
-      Dialect metadataDatabaseDialect) {
+      Dialect metadataDatabaseDialect,
+      boolean isMutableChangeStream) {
     if (metadataSpannerConfig.getInstanceId() == null) {
       throw new IllegalArgumentException("Metadata instance can not be null");
     }
@@ -83,6 +85,7 @@ public DaoFactory(
     this.jobName = jobName;
     this.spannerChangeStreamDatabaseDialect = spannerChangeStreamDatabaseDialect;
     this.metadataDatabaseDialect = metadataDatabaseDialect;
+    this.isMutableChangeStream = isMutableChangeStream;
   }
 
   /**
@@ -143,7 +146,8 @@ public synchronized ChangeStreamDao getChangeStreamDao() {
               spannerAccessor.getDatabaseClient(),
               rpcPriority,
               jobName,
-              this.spannerChangeStreamDatabaseDialect);
+              this.spannerChangeStreamDatabaseDialect,
+              this.isMutableChangeStream);
     }
     return changeStreamDaoInstance;
   }

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

@@ -19,10 +19,18 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.when;
 
 import com.google.auth.Credentials;
 import com.google.cloud.Timestamp;
+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 com.google.cloud.spanner.Statement;
 import org.apache.beam.sdk.extensions.gcp.auth.TestCredential;
 import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.MetadataSpannerConfigFactory;
@@ -32,6 +40,7 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.Mockito;
 
 @RunWith(JUnit4.class)
 public class SpannerIOReadChangeStreamTest {
@@ -120,7 +129,8 @@ public void testSetSpannerConfigCredential() {
 
   @Test
   public void testWithDefaultCredential() {
-    // Get the default credential, without setting any credentials in the pipeline options or
+    // Get the default credential, without setting any credentials in the pipeline
+    // options or
     // SpannerConfig.
     Credentials defaultCredential =
         testPipeline.getOptions().as(GcpOptions.class).getGcpCredential();
@@ -140,4 +150,53 @@ public void testWithDefaultCredential() {
     assertEquals(defaultCredential, changeStreamSpannerConfigWithCredential.getCredentials().get());
     assertEquals(defaultCredential, metadataSpannerConfigWithCredential.getCredentials().get());
   }
+
+  @Test
+  public void testFetchPartitionModeGoogleSql() {
+    DatabaseClient databaseClient = Mockito.mock(DatabaseClient.class);
+    ReadOnlyTransaction transaction = Mockito.mock(ReadOnlyTransaction.class);
+    ResultSet resultSet = Mockito.mock(ResultSet.class);
+
+    when(databaseClient.readOnlyTransaction()).thenReturn(transaction);
+    when(transaction.executeQuery(any(Statement.class))).thenReturn(resultSet);
+    when(resultSet.next()).thenReturn(true).thenReturn(false);
+    when(resultSet.getString(0)).thenReturn("MUTABLE_KEY_RANGE");
+
+    assertEquals(
+        true,
+        SpannerIO.isMutableChangeStream(
+            databaseClient, Dialect.GOOGLE_STANDARD_SQL, TEST_CHANGE_STREAM));
+  }
+
+  @Test
+  public void testFetchPartitionModePostgres() {
+    DatabaseClient databaseClient = Mockito.mock(DatabaseClient.class);
+    ReadOnlyTransaction transaction = Mockito.mock(ReadOnlyTransaction.class);
+    ResultSet resultSet = Mockito.mock(ResultSet.class);
+
+    when(databaseClient.readOnlyTransaction()).thenReturn(transaction);
+    when(transaction.executeQuery(any(Statement.class))).thenReturn(resultSet);
+    when(resultSet.next()).thenReturn(true).thenReturn(false);
+    when(resultSet.getString(0)).thenReturn("IMMUTABLE_KEY_RANGE");
+
+    assertEquals(
+        false,
+        SpannerIO.isMutableChangeStream(databaseClient, Dialect.POSTGRESQL, TEST_CHANGE_STREAM));
+  }
+
+  @Test
+  public void testFetchPartitionModeNotFound() {
+    DatabaseClient databaseClient = Mockito.mock(DatabaseClient.class);
+    ReadOnlyTransaction transaction = Mockito.mock(ReadOnlyTransaction.class);
+    ResultSet resultSet = Mockito.mock(ResultSet.class);
+
+    when(databaseClient.readOnlyTransaction()).thenReturn(transaction);
+    when(transaction.executeQuery(any(Statement.class))).thenReturn(resultSet);
+    when(resultSet.next()).thenReturn(false);
+
+    assertEquals(
+        false,
+        SpannerIO.isMutableChangeStream(
+            databaseClient, Dialect.GOOGLE_STANDARD_SQL, TEST_CHANGE_STREAM));
+  }
 }

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

@@ -0,0 +1,97 @@
+/*
+ * 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.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+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 com.google.cloud.spanner.Statement;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+
+public class ChangeStreamDaoTest {
+  private DatabaseClient databaseClient;
+  private RpcPriority rpcPriority;
+  private static final String CHANGE_STREAM_NAME = "testCS";
+
+  @Before
+  public void setUp() {
+    databaseClient = mock(DatabaseClient.class);
+    rpcPriority = mock(RpcPriority.class);
+  }
+
+  // New tests for PostgreSQL branch to verify the chosen TVF in the generated
+  // SQL.
+  @Test
+  public void testChangeStreamQueryPostgresMutable() {
+    // Arrange: single-use transaction for the actual change stream query
+    ReadOnlyTransaction singleUseTx = mock(ReadOnlyTransaction.class);
+    when(databaseClient.singleUse()).thenReturn(singleUseTx);
+
+    ChangeStreamDao changeStreamDao =
+        new ChangeStreamDao(
+            CHANGE_STREAM_NAME, databaseClient, rpcPriority, "testjob", Dialect.POSTGRESQL, true);
+
+    // Act: call the method that constructs and executes the statement
+    changeStreamDao.changeStreamQuery(null, null, null, 0L);
+
+    // Assert: capture the Statement passed to singleUse().executeQuery and verify
+    // SQL
+    ArgumentCaptor<Statement> captor = ArgumentCaptor.forClass(Statement.class);
+    verify(singleUseTx).executeQuery(captor.capture(), any(), any());
+    Statement captured = captor.getValue();
+    String sql = captured.getSql(); // adjust if different accessor is used
+    assertTrue(
+        "Expected SQL to contain read_proto_bytes_",
+        sql.contains("read_proto_bytes_" + CHANGE_STREAM_NAME));
+  }
+
+  @Test
+  public void testChangeStreamQueryPostgresImmutable() {
+    // Arrange: single-use transaction for the actual change stream query
+    ReadOnlyTransaction singleUseTx = mock(ReadOnlyTransaction.class);
+    when(databaseClient.singleUse()).thenReturn(singleUseTx);
+
+    ResultSet queryResult = mock(ResultSet.class);
+    when(singleUseTx.executeQuery(any(), any(), any())).thenReturn(queryResult);
+
+    ChangeStreamDao changeStreamDao =
+        new ChangeStreamDao(
+            CHANGE_STREAM_NAME, databaseClient, rpcPriority, "testjob", Dialect.POSTGRESQL, false);
+
+    // Act
+    changeStreamDao.changeStreamQuery(null, null, null, 0L);
+
+    // Assert
+    ArgumentCaptor<Statement> captor = ArgumentCaptor.forClass(Statement.class);
+    verify(singleUseTx).executeQuery(captor.capture(), any(), any());
+    Statement captured = captor.getValue();
+    String sql = captured.getSql();
+    assertTrue(
+        "Expected SQL to contain read_json_", sql.contains("read_json_" + CHANGE_STREAM_NAME));
+  }
+}