[sourcedb-to-spanner] Fix GCS Avro Export flow (#3845)

* initial changes

* null exception

* review changes

Author: aasthabharill

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/AvroDestination.java

@@ -16,14 +16,15 @@
 package com.google.cloud.teleport.v2.templates;
 
 import java.util.Objects;
+import org.apache.avro.reflect.Nullable;
 import org.apache.beam.sdk.coders.DefaultCoder;
 import org.apache.beam.sdk.extensions.avro.coders.AvroCoder;
 
 @DefaultCoder(AvroCoder.class)
 public class AvroDestination {
   public String name;
   public String jsonSchema;
-  public String shardId;
+  @Nullable public String shardId;
 
   // Needed for serialization
   public AvroDestination() {}

v2/sourcedb-to-spanner/src/main/java/com/google/cloud/teleport/v2/templates/MigrateTableTransform.java

@@ -162,9 +162,7 @@ public WriteFilesResult<AvroDestination> writeToGCS(
             .by(
                 (record) ->
                     AvroDestination.of(
-                        record.shardId(),
-                        record.tableName(),
-                        record.getPayload().getSchema().toString()))
+                        record.shardId(), record.tableName(), record.gcsSchema().toString()))
             .via(
                 Contextful.fn(
                     record -> {

v2/sourcedb-to-spanner/src/test/java/com/google/cloud/teleport/v2/templates/MigrateTableTransformTest.java

@@ -22,13 +22,21 @@
 import com.google.cloud.teleport.v2.options.SourceDbToSpannerOptions;
 import com.google.cloud.teleport.v2.source.reader.ReaderImpl;
 import com.google.cloud.teleport.v2.source.reader.io.row.SourceRow;
+import com.google.cloud.teleport.v2.source.reader.io.schema.SchemaTestUtils;
+import com.google.cloud.teleport.v2.source.reader.io.schema.SourceTableSchema;
+import com.google.cloud.teleport.v2.source.reader.io.schema.typemapping.UnifiedTypeMapper;
 import com.google.cloud.teleport.v2.source.reader.io.transform.ReaderTransform;
 import com.google.cloud.teleport.v2.spanner.ddl.Ddl;
 import com.google.cloud.teleport.v2.spanner.migrations.schema.ISchemaMapper;
+import com.google.cloud.teleport.v2.spanner.migrations.schema.SourceColumnType;
+import org.apache.avro.file.DataFileReader;
+import org.apache.avro.generic.GenericDatumReader;
+import org.apache.avro.generic.GenericRecord;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.io.Compression;
 import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.testing.TestPipeline;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
@@ -37,13 +45,16 @@
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionTuple;
 import org.apache.beam.sdk.values.TupleTag;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
 /** Test class for {@link MigrateTableTransform}. */
 @RunWith(JUnit4.class)
 public class MigrateTableTransformTest {
+  @Rule public TemporaryFolder tempFolder = new TemporaryFolder();
 
   /** Tests that metric names are correctly generated, optionally including the shard ID. */
   @Test
@@ -132,4 +143,90 @@ public PCollectionTuple expand(PBegin input) {
     Pipeline p = Pipeline.create();
     migrateTableTransform.expand(PBegin.in(p));
   }
+
+  /**
+   * Tests the writeToGCS method end-to-end to ensure that records are serialized correctly using
+   * the wrapper schema in a non-sharded flow, where shardId is null.
+   */
+  @Test
+  public void testWriteToGCSForNonShardedFlow() throws Exception {
+    runGcsWriteTest(null, "table1");
+  }
+
+  /**
+   * Tests the writeToGCS method end-to-end to ensure that records are serialized correctly using
+   * the wrapper schema in a sharded flow, where shardId is present.
+   */
+  @Test
+  public void testWriteToGCSForShardedFlow() throws Exception {
+    runGcsWriteTest("shard1", "table1/shard1");
+  }
+
+  private void runGcsWriteTest(String shardId, String expectedRelativeSubDir) throws Exception {
+    TestPipeline pipeline = TestPipeline.create().enableAbandonedNodeEnforcement(false);
+
+    final String testTable = "table1";
+    final long testReadTime = 1712751118L;
+    var schemaRef = SchemaTestUtils.generateSchemaReference("public", "mydb");
+    var schema =
+        SourceTableSchema.builder(UnifiedTypeMapper.MapperType.MYSQL)
+            .setTableName(testTable)
+            .addSourceColumnNameToSourceColumnType(
+                "id", new SourceColumnType("INTEGER", new Long[] {}, null))
+            .addSourceColumnNameToSourceColumnType(
+                "firstName", new SourceColumnType("varchar", new Long[] {20L}, null))
+            .build();
+    SourceRow sourceRow =
+        SourceRow.builder(schemaRef, schema, shardId, testReadTime)
+            .setField("id", 123)
+            .setField("firstName", "abc")
+            .build();
+
+    PCollection<SourceRow> sourceRows = pipeline.apply(Create.of(sourceRow));
+    String tempDirPath = tempFolder.getRoot().getAbsolutePath();
+
+    SourceDbToSpannerOptions options = PipelineOptionsFactory.as(SourceDbToSpannerOptions.class);
+    options.setSourceDbDialect("MYSQL");
+    options.setGcsOutputDirectory(tempDirPath);
+
+    SpannerConfig mockSpannerConfig = mock(SpannerConfig.class);
+    Ddl mockDdl = mock(Ddl.class);
+    ISchemaMapper mockSchemaMapper = mock(ISchemaMapper.class);
+    ReaderImpl mockReader = mock(ReaderImpl.class);
+
+    MigrateTableTransform migrateTableTransform =
+        new MigrateTableTransform(
+            options, mockSpannerConfig, mockDdl, mockSchemaMapper, mockReader);
+
+    migrateTableTransform.writeToGCS(sourceRows, tempDirPath);
+
+    // Run the pipeline to execute the Avro write logic on actual data
+    pipeline.run().waitUntilFinish();
+
+    // Verify Avro File was written and matches wrapper schema
+    java.io.File avroSubDir = new java.io.File(tempDirPath, expectedRelativeSubDir);
+    java.io.File[] avroFiles = avroSubDir.listFiles((dir, name) -> name.endsWith(".avro"));
+    assertThat(avroFiles).isNotNull();
+    assertThat(avroFiles).hasLength(1);
+    java.io.File avroFile = avroFiles[0];
+
+    try (DataFileReader<GenericRecord> fileReader =
+        new DataFileReader<>(avroFile, new GenericDatumReader<>())) {
+      assertThat(fileReader.hasNext()).isTrue();
+      GenericRecord record = fileReader.next();
+
+      assertThat(record.get("tableName").toString()).isEqualTo("table1");
+      if (shardId == null) {
+        assertThat(record.get("shardId")).isNull();
+      } else {
+        assertThat(record.get("shardId").toString()).isEqualTo(shardId);
+      }
+
+      GenericRecord payloadRecord = (GenericRecord) record.get("payload");
+      assertThat(payloadRecord.get("id")).isEqualTo(123);
+      assertThat(payloadRecord.get("firstName").toString()).isEqualTo("abc");
+
+      assertThat(fileReader.hasNext()).isFalse();
+    }
+  }
 }