Parse proto bytes change stream record in ChangeStreamRecordMapper

Author: chenxuesdu

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

@@ -1977,11 +1977,6 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta
               + changeStreamDatabaseId
               + " has dialect "
               + changeStreamDatabaseDialect);
-      LOG.info(
-          "The Spanner database "
-              + fullPartitionMetadataDatabaseId
-              + " has dialect "
-              + metadataDatabaseDialect);
       PartitionMetadataTableNames partitionMetadataTableNames =
           Optional.ofNullable(getMetadataTable())
               .map(
@@ -2005,6 +2000,7 @@ && getInclusiveStartAt().toSqlTimestamp().after(getInclusiveEndAt().toSqlTimesta
       final boolean isMutableChangeStream =
           isMutableChangeStream(
               spannerAccessor.getDatabaseClient(), changeStreamDatabaseDialect, changeStreamName);
+      LOG.info("The change stream " + changeStreamName + " is mutable: " + isMutableChangeStream);
       final DaoFactory daoFactory =
           new DaoFactory(
               changeStreamSpannerConfig,

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

@@ -20,6 +20,7 @@
 import com.google.cloud.Timestamp;
 import com.google.cloud.spanner.ResultSet;
 import com.google.cloud.spanner.Struct;
+import com.google.protobuf.InvalidProtocolBufferException;
 import org.joda.time.Duration;
 
 /**
@@ -128,6 +129,33 @@ public boolean isProtoChangeRecord() {
         && resultSet.getColumnType(0).getCode() == com.google.cloud.spanner.Type.Code.PROTO;
   }
 
+  /**
+   * Returns the change stream record at the current pointer by parsing the bytes column. It also
+   * updates the timestamp at which the record was read.
+   *
+   * <p>Should only be used for PostgreSQL databases when the change stream record is delivered as
+   * proto bytes.
+   *
+   * @return a change stream record as a proto or null
+   */
+  public com.google.spanner.v1.ChangeStreamRecord getProtoChangeStreamRecordFromBytes() {
+    recordReadAt = Timestamp.now();
+    try {
+      // Use getBytes(0) for the BYTES column returned by read_proto_bytes_ TVF
+      return com.google.spanner.v1.ChangeStreamRecord.parseFrom(
+          resultSet.getBytes(0).toByteArray());
+    } catch (InvalidProtocolBufferException e) {
+      throw new RuntimeException("Failed to parse the proto bytes to ChangeStreamRecord proto", e);
+    }
+  }
+
+  /** Returns true if the result set at the current pointer contain only one bytes change record. */
+  public boolean isProtoBytesChangeRecord() {
+    return resultSet.getColumnCount() == 1
+        && !resultSet.isNull(0)
+        && resultSet.getColumnType(0).getCode() == com.google.cloud.spanner.Type.Code.BYTES;
+  }
+
   /**
    * Returns the record at the current pointer as {@link JsonB}. It also updates the timestamp at
    * which the record was read.

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

@@ -218,18 +218,28 @@ public class ChangeStreamRecordMapper {
    * @param resultSet the change stream result set
    * @param resultSetMetadata the metadata generated when reading the change stream row
    * @return a {@link List} of {@link ChangeStreamRecord} subclasses
+   * @throws InvalidProtocolBufferException
    */
   public List<ChangeStreamRecord> toChangeStreamRecords(
       PartitionMetadata partition,
       ChangeStreamResultSet resultSet,
       ChangeStreamResultSetMetadata resultSetMetadata) {
     if (this.isPostgres()) {
-      // In PostgresQL, change stream records are returned as JsonB.
+      // For `MUTABLE_KEY_RANGE` option, change stream records are returned as protos.
+      if (resultSet.isProtoBytesChangeRecord()) {
+        return Arrays.asList(
+            toChangeStreamRecord(
+                partition, resultSet.getProtoChangeStreamRecordFromBytes(), resultSetMetadata));
+      }
+
+      // For `IMMUTABLE_KEY_RANGE` option, change stream records are returned as
+      // JsonB.
       return Collections.singletonList(
           toChangeStreamRecordJson(partition, resultSet.getPgJsonb(0), resultSetMetadata));
     }
 
-    // In GoogleSQL, for `MUTABLE_KEY_RANGE` option, change stream records are returned as Protos.
+    // In GoogleSQL, for `MUTABLE_KEY_RANGE` option, change stream records are
+    // returned as Protos.
     if (resultSet.isProtoChangeRecord()) {
       return Arrays.asList(
           toChangeStreamRecord(

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

@@ -0,0 +1,61 @@
+/*
+ * 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.apache.beam.sdk.io.gcp.spanner.changestreams.util.TestProtoMapper.recordToProto;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import com.google.cloud.ByteArray;
+import com.google.cloud.Timestamp;
+import com.google.cloud.spanner.ResultSet;
+import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.HeartbeatRecord;
+import org.junit.Test;
+
+public class ChangeStreamResultSetTest {
+
+  @Test
+  public void testGetProtoChangeStreamRecordFromBytes() throws Exception {
+    // 1. Create an expected ChangeStreamRecord proto
+    Timestamp now = Timestamp.now();
+    final HeartbeatRecord heartbeatRecord =
+        new HeartbeatRecord(Timestamp.ofTimeSecondsAndNanos(10L, 20), null);
+    com.google.spanner.v1.ChangeStreamRecord expectedRecord = recordToProto(heartbeatRecord);
+    assertNotNull(expectedRecord);
+
+    // 2. Convert it to bytes (simulating how Spanner PostgreSQL returns it)
+    byte[] protoBytes = expectedRecord.toByteArray();
+
+    // 3. Mock the underlying Spanner ResultSet
+    ResultSet mockResultSet = mock(ResultSet.class);
+    // Simulate column 0 containing the BYTES representation of the proto
+    when(mockResultSet.getBytes(0)).thenReturn(ByteArray.copyFrom(protoBytes));
+
+    // 4. Initialize ChangeStreamResultSet with the mock
+    ChangeStreamResultSet changeStreamResultSet = new ChangeStreamResultSet(mockResultSet);
+
+    // 5. Call the new method and assert it parses correctly
+    // (Note: This assumes you have added getProtoChangeStreamRecordFromBytes to the class)
+    com.google.spanner.v1.ChangeStreamRecord actualRecord =
+        changeStreamResultSet.getProtoChangeStreamRecordFromBytes();
+
+    assertEquals(expectedRecord, actualRecord);
+  }
+}

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

@@ -1039,4 +1039,109 @@ public void testMappingProtoRowToDataChangeRecord() {
         Collections.singletonList(dataChangeRecord),
         mapper.toChangeStreamRecords(partition, resultSet, resultSetMetadata));
   }
+
+  @Test
+  public void testMappingProtoBytesRowToPartitionStartRecord() {
+    final PartitionStartRecord partitionStartRecord =
+        new PartitionStartRecord(
+            Timestamp.MIN_VALUE,
+            "fakeRecordSequence",
+            Arrays.asList("partitionToken1", "partitionToken2"),
+            null);
+    com.google.spanner.v1.ChangeStreamRecord changeStreamRecordProto =
+        recordToProto(partitionStartRecord);
+    assertNotNull(changeStreamRecordProto);
+    ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class);
+
+    when(resultSet.isProtoBytesChangeRecord()).thenReturn(true);
+    when(resultSet.getProtoChangeStreamRecordFromBytes()).thenReturn(changeStreamRecordProto);
+    assertEquals(
+        Collections.singletonList(partitionStartRecord),
+        mapperPostgres.toChangeStreamRecords(partition, resultSet, resultSetMetadata));
+  }
+
+  @Test
+  public void testMappingProtoBytesRowToPartitionEndRecord() {
+    final PartitionEndRecord partitionEndChange =
+        new PartitionEndRecord(Timestamp.MIN_VALUE, "fakeRecordSequence", null);
+    com.google.spanner.v1.ChangeStreamRecord changeStreamRecordProto =
+        recordToProto(partitionEndChange);
+    assertNotNull(changeStreamRecordProto);
+    ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class);
+
+    when(resultSet.isProtoBytesChangeRecord()).thenReturn(true);
+    when(resultSet.getProtoChangeStreamRecordFromBytes()).thenReturn(changeStreamRecordProto);
+    assertEquals(
+        Collections.singletonList(partitionEndChange),
+        mapperPostgres.toChangeStreamRecords(partition, resultSet, resultSetMetadata));
+  }
+
+  @Test
+  public void testMappingProtoBytesRowToPartitionEventRecord() {
+    final PartitionEventRecord partitionEventRecord =
+        new PartitionEventRecord(Timestamp.MIN_VALUE, "fakeRecordSequence", null);
+    com.google.spanner.v1.ChangeStreamRecord changeStreamRecordProto =
+        recordToProto(partitionEventRecord);
+    assertNotNull(changeStreamRecordProto);
+    ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class);
+
+    when(resultSet.isProtoBytesChangeRecord()).thenReturn(true);
+    when(resultSet.getProtoChangeStreamRecordFromBytes()).thenReturn(changeStreamRecordProto);
+    assertEquals(
+        Collections.singletonList(partitionEventRecord),
+        mapperPostgres.toChangeStreamRecords(partition, resultSet, resultSetMetadata));
+  }
+
+  @Test
+  public void testMappingProtoBytesRowToHeartbeatRecord() {
+    final HeartbeatRecord heartbeatRecord =
+        new HeartbeatRecord(Timestamp.ofTimeSecondsAndNanos(10L, 20), null);
+    com.google.spanner.v1.ChangeStreamRecord changeStreamRecordProto =
+        recordToProto(heartbeatRecord);
+    assertNotNull(changeStreamRecordProto);
+    ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class);
+
+    when(resultSet.isProtoBytesChangeRecord()).thenReturn(true);
+    when(resultSet.getProtoChangeStreamRecordFromBytes()).thenReturn(changeStreamRecordProto);
+    assertEquals(
+        Collections.singletonList(heartbeatRecord),
+        mapperPostgres.toChangeStreamRecords(partition, resultSet, resultSetMetadata));
+  }
+
+  @Test
+  public void testMappingProtoBytesRowToDataChangeRecord() {
+    final DataChangeRecord dataChangeRecord =
+        new DataChangeRecord(
+            "partitionToken",
+            Timestamp.ofTimeSecondsAndNanos(10L, 20),
+            "serverTransactionId",
+            true,
+            "1",
+            "tableName",
+            Arrays.asList(
+                new ColumnType("column1", new TypeCode("{\"code\":\"INT64\"}"), true, 1L),
+                new ColumnType("column2", new TypeCode("{\"code\":\"BYTES\"}"), false, 2L)),
+            Collections.singletonList(
+                new Mod(
+                    "{\"column1\":\"value1\"}",
+                    "{\"column2\":\"oldValue2\"}",
+                    "{\"column2\":\"newValue2\"}")),
+            ModType.UPDATE,
+            ValueCaptureType.OLD_AND_NEW_VALUES,
+            10L,
+            2L,
+            "transactionTag",
+            true,
+            null);
+    com.google.spanner.v1.ChangeStreamRecord changeStreamRecordProto =
+        recordToProto(dataChangeRecord);
+    assertNotNull(changeStreamRecordProto);
+    ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class);
+
+    when(resultSet.isProtoBytesChangeRecord()).thenReturn(true);
+    when(resultSet.getProtoChangeStreamRecordFromBytes()).thenReturn(changeStreamRecordProto);
+    assertEquals(
+        Collections.singletonList(dataChangeRecord),
+        mapperPostgres.toChangeStreamRecords(partition, resultSet, resultSetMetadata));
+  }
 }