Parse v2 record.

pick 4d2a5a7 # Fix API surface test (apache#35028)

Author: changliiu

buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy

@@ -656,6 +656,7 @@ class BeamModulePlugin implements Plugin<Project> {
     def arrow_version = "15.0.2"
     def jmh_version = "1.34"
     def jupiter_version = "5.7.0"
+    def spanner_grpc_proto_version = "6.95.1"
 
     // Export Spark versions, so they are defined in a single place only
     project.ext.spark3_version = spark3_version
@@ -866,7 +867,7 @@ class BeamModulePlugin implements Plugin<Project> {
         proto_google_cloud_pubsub_v1                : "com.google.api.grpc:proto-google-cloud-pubsub-v1", // google_cloud_platform_libraries_bom sets version
         proto_google_cloud_pubsublite_v1            : "com.google.api.grpc:proto-google-cloud-pubsublite-v1", // google_cloud_platform_libraries_bom sets version
         proto_google_cloud_secret_manager_v1        : "com.google.api.grpc:proto-google-cloud-secretmanager-v1", // google_cloud_platform_libraries_bom sets version
-        proto_google_cloud_spanner_v1               : "com.google.api.grpc:proto-google-cloud-spanner-v1", // google_cloud_platform_libraries_bom sets version
+        proto_google_cloud_spanner_v1               : "com.google.api.grpc:proto-google-cloud-spanner-v1:$spanner_grpc_proto_version", // google_cloud_platform_libraries_bom sets version
         proto_google_cloud_spanner_admin_database_v1: "com.google.api.grpc:proto-google-cloud-spanner-admin-database-v1", // google_cloud_platform_libraries_bom sets version
         proto_google_common_protos                  : "com.google.api.grpc:proto-google-common-protos", // google_cloud_platform_libraries_bom sets version
         qpid_jms_client                             : "org.apache.qpid:qpid-jms-client:$qpid_jms_client_version",

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

@@ -40,12 +40,18 @@
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.DataChangeRecord;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.HeartbeatRecord;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.InitialPartition;
+import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionStartRecord;
+import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionEndRecord;
+import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionEventRecord;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.Mod;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.ModType;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.PartitionMetadata;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.TypeCode;
 import org.apache.beam.sdk.io.gcp.spanner.changestreams.model.ValueCaptureType;
+import java.util.Arrays;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * This class is responsible for transforming a {@link Struct} to a {@link List} of {@link
@@ -56,6 +62,7 @@
  */
 public class ChangeStreamRecordMapper {
 
+  private static final Logger LOG = LoggerFactory.getLogger(ChangeStreamRecordMapper.class);
   private static final String DATA_CHANGE_RECORD_COLUMN = "data_change_record";
   private static final String HEARTBEAT_RECORD_COLUMN = "heartbeat_record";
   private static final String CHILD_PARTITIONS_RECORD_COLUMN = "child_partitions_record";
@@ -223,12 +230,45 @@ public List<ChangeStreamRecord> toChangeStreamRecords(
       return Collections.singletonList(
           toChangeStreamRecordJson(partition, resultSet.getPgJsonb(0), resultSetMetadata));
     }
+
+    if (containOnlyProtoType(resultSet.getCurrentRowAsStruct())) {
+      return Arrays.asList(
+          toV2ChangeStreamRecord(partition, resultSetMetadata, resultSet.getCurrentRowAsStruct()));
+    }
+
     // In GoogleSQL, change stream records are returned as an array of structs.
     return resultSet.getCurrentRowAsStruct().getStructList(0).stream()
         .flatMap(struct -> toChangeStreamRecord(partition, struct, resultSetMetadata))
         .collect(Collectors.toList());
   }
 
+  boolean containOnlyProtoType(Struct currentRow){
+    LOG.info("changliiu containOnlyProtoType1");
+    if(currentRow.getColumnCount() !=1 || currentRow.isNull(0)){
+      LOG.info("changliiu containOnlyProtoType2");
+      return false;
+    }
+    com.google.cloud.spanner.Type columnType = currentRow.getColumnType(i);
+    if (columnType.getCode() == com.google.cloud.spanner.Type.Code.PROTO) {
+      LOG.info("changliiu containOnlyProtoType3");
+      return true;
+    }
+    LOG.info("changliiu containOnlyProtoType4");
+    return false;
+  }
+
+  ChangeStreamRecord toV2ChangeStreamRecord(PartitionMetadata partition,
+      ChangeStreamResultSetMetadata resultSetMetadata,
+      Struct currentRow) {
+    com.google.spanner.v1.ChangeStreamRecord changeStreamRecord =
+        currentRow.getProtoMessage(0,
+            com.google.spanner.v1.ChangeStreamRecord.getDefaultInstance());
+    // todo finish this.
+    LOG.info("changliiu toV2ChangeStreamRecord");
+    return new PartitionEndRecord(Timestamp.MIN_VALUE, "xxxx",
+        changeStreamRecordMetadataFrom(partition, Timestamp.MIN_VALUE, resultSetMetadata));
+  }
+
   Stream<ChangeStreamRecord> toChangeStreamRecord(
       PartitionMetadata partition, Struct row, ChangeStreamResultSetMetadata resultSetMetadata) {
 

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

@@ -930,4 +930,24 @@ public void testMappingJsonRowToChildPartitionRecord() {
         Collections.singletonList(childPartitionsRecord),
         mapperPostgres.toChangeStreamRecords(partition, resultSet, resultSetMetadata));
   }
+
+  @Test
+  public void testMappingProtoRowToPartitionStartRecord() {
+    com.google.spanner.v1.ChangeStreamRecord.PartitionStartRecord partitionStartRecord =
+        com.google.spanner.v1.ChangeStreamRecord.PartitionStartRecord.newBuilder()
+            .setStartTimestamp(Timestamp.MIN_VALUE)
+            .setRecordSequence("sequence")
+            .addPartitionTokens("partitionTokens")
+            .build();
+    com.google.spanner.v1.ChangeStreamRecord record = com.google.spanner.v1.ChangeStreamRecord.newBuilder()
+        .setPartitionStartRecord(partitionStartRecord).build();
+
+    assertNotNull(partitionStartRecord);
+    ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class);
+
+    when(resultSet.getPgJsonb(0)).thenReturn(jsonString);
+    assertEquals(
+        Collections.singletonList(partitionStartRecord),
+        mapper.toChangeStreamRecords(partition, struct, resultSetMetadata));
+  }
 }