From 3679cb4a320941ae283aa38c7a1f761f51dcc3db Mon Sep 17 00:00:00 2001 From: aasthabharill Date: Sun, 9 Aug 2026 13:26:45 +0000 Subject: [PATCH] initial --- .../v2/transforms/ReportResultsTransform.java | 40 +++++--- .../ReportResultsTransformTest.java | 99 +++++++++++++++++-- 2 files changed, 121 insertions(+), 18 deletions(-) diff --git a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/transforms/ReportResultsTransform.java b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/transforms/ReportResultsTransform.java index 9c05ab54e4..15261911dd 100644 --- a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/transforms/ReportResultsTransform.java +++ b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/transforms/ReportResultsTransform.java @@ -33,6 +33,7 @@ import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.WriteDisposition; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices; import org.apache.beam.sdk.schemas.NoSuchSchemaException; import org.apache.beam.sdk.schemas.SchemaCoder; import org.apache.beam.sdk.transforms.Combine; @@ -74,6 +75,15 @@ public class ReportResultsTransform extends PTransform private final String runId; private final Instant startTimestamp; + private BigQueryServices testServices; + private String testTempLocation; + + public ReportResultsTransform withTestServices(BigQueryServices testServices, String testTempLocation) { + this.testServices = testServices; + this.testTempLocation = testTempLocation; + return this; + } + public ReportResultsTransform(String bigQueryDataset, String runId, Instant startTimestamp) { this.bigQueryDataset = bigQueryDataset; this.runId = runId; @@ -90,9 +100,7 @@ public ReportResultsTransform(String bigQueryDataset, String runId, Instant star PCollection allMismatches = transformMismatchedRecords(missingInSpanner, missingInSource); - allMismatches.apply( - "WriteMismatchedRecords", - BigQueryIO.write() + BigQueryIO.Write writeMismatches = BigQueryIO.write() .to(String.format("%s.%s", bigQueryDataset, MISMATCHED_RECORDS_TABLE)) .withSchema(BigQuerySchemas.MISMATCHED_RECORDS_SCHEMA) .withFormatFunction( @@ -108,15 +116,17 @@ public ReportResultsTransform(String bigQueryDataset, String runId, Instant star .set(MismatchedRecord.SHARD_ID_COLUMN_NAME, r.getShardId())) .withCreateDisposition(CreateDisposition.CREATE_IF_NEEDED) .withWriteDisposition(WriteDisposition.WRITE_APPEND) - .withMethod(BigQueryIO.Write.Method.FILE_LOADS)); + .withMethod(BigQueryIO.Write.Method.FILE_LOADS); + if (this.testServices != null) { + writeMismatches = writeMismatches.withTestServices(this.testServices).withCustomGcsTempLocation(org.apache.beam.sdk.options.ValueProvider.StaticValueProvider.of(this.testTempLocation)); + } + allMismatches.apply("WriteMismatchedRecords", writeMismatches); // 2. Aggregate Stats PCollection tableStats = calculateTableStats(matched, missingInSpanner, missingInSource); - tableStats.apply( - "WriteTableStats", - BigQueryIO.write() + BigQueryIO.Write writeTableStats = BigQueryIO.write() .to(String.format("%s.%s", bigQueryDataset, TABLE_VALIDATION_STATS_TABLE)) .withSchema(BigQuerySchemas.TABLE_VALIDATION_STATS_SCHEMA) .withFormatFunction( @@ -146,7 +156,11 @@ public ReportResultsTransform(String bigQueryDataset, String runId, Instant star stats.getEndTimestamp().toString())) .withCreateDisposition(CreateDisposition.CREATE_IF_NEEDED) .withWriteDisposition(WriteDisposition.WRITE_APPEND) - .withMethod(BigQueryIO.Write.Method.FILE_LOADS)); + .withMethod(BigQueryIO.Write.Method.FILE_LOADS); + if (this.testServices != null) { + writeTableStats = writeTableStats.withTestServices(this.testServices).withCustomGcsTempLocation(org.apache.beam.sdk.options.ValueProvider.StaticValueProvider.of(this.testTempLocation)); + } + tableStats.apply("WriteTableStats", writeTableStats); // 3. Validation Summary /** @@ -172,9 +186,7 @@ public ReportResultsTransform(String bigQueryDataset, String runId, Instant star PCollection validationSummary = calculateValidationSummary(tableStats); - validationSummary.apply( - "WriteValidationSummary", - BigQueryIO.write() + BigQueryIO.Write writeValidationSummary = BigQueryIO.write() .to(String.format("%s.%s", this.bigQueryDataset, VALIDATION_SUMMARY_TABLE)) .withSchema(BigQuerySchemas.VALIDATION_SUMMARY_SCHEMA) .withFormatFunction( @@ -206,7 +218,11 @@ public ReportResultsTransform(String bigQueryDataset, String runId, Instant star s.getEndTimestamp().toString())) .withCreateDisposition(CreateDisposition.CREATE_IF_NEEDED) .withWriteDisposition(WriteDisposition.WRITE_APPEND) - .withMethod(BigQueryIO.Write.Method.STREAMING_INSERTS)); + .withMethod(BigQueryIO.Write.Method.STREAMING_INSERTS); + if (this.testServices != null) { + writeValidationSummary = writeValidationSummary.withTestServices(this.testServices); + } + validationSummary.apply("WriteValidationSummary", writeValidationSummary); return PDone.in(input.getPipeline()); } diff --git a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/transforms/ReportResultsTransformTest.java b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/transforms/ReportResultsTransformTest.java index 8d54b25239..b61fe13c7e 100644 --- a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/transforms/ReportResultsTransformTest.java +++ b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/transforms/ReportResultsTransformTest.java @@ -15,6 +15,17 @@ */ package com.google.cloud.teleport.v2.transforms; +import static com.google.cloud.teleport.v2.constants.GCSSpannerDVConstants.MATCHED_TAG; +import static com.google.cloud.teleport.v2.constants.GCSSpannerDVConstants.MISSING_IN_SOURCE_TAG; +import static com.google.cloud.teleport.v2.constants.GCSSpannerDVConstants.MISSING_IN_SPANNER_TAG; +import static org.junit.Assert.assertEquals; + +import com.google.api.services.bigquery.model.ErrorProto; +import com.google.api.services.bigquery.model.Table; +import com.google.api.services.bigquery.model.TableDataInsertAllResponse.InsertErrors; +import com.google.api.services.bigquery.model.TableReference; +import com.google.api.services.bigquery.model.TableRow; +import com.google.cloud.teleport.v2.dto.BigQuerySchemas; import com.google.cloud.teleport.v2.dto.Column; import com.google.cloud.teleport.v2.dto.ComparisonRecord; import com.google.cloud.teleport.v2.dto.MismatchedRecord; @@ -23,29 +34,36 @@ import java.io.Serializable; import java.util.Arrays; import java.util.Collections; +import org.apache.beam.sdk.io.gcp.testing.FakeBigQueryServices; +import org.apache.beam.sdk.io.gcp.testing.FakeDatasetService; +import org.apache.beam.sdk.io.gcp.testing.FakeJobService; import org.apache.beam.sdk.schemas.NoSuchSchemaException; import org.apache.beam.sdk.schemas.SchemaCoder; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; import org.apache.beam.sdk.values.TypeDescriptor; +import org.joda.time.DateTimeUtils; import org.joda.time.Instant; import org.junit.Rule; import org.junit.Test; + + + /** - * This unit test does not test the expand() method of the {@link ReportResultsTransform} directly, - * but instead tests all the logical units inside it (transforming records for BQ write, and - * calculation of stats). This is because testing the full expand() method requires instrumenting - * BigQuery writes, which is hard in a unit testing environment. While testing the BQ sink from a - * functionality standpoint is not really a validation concern, we will cover this in the - * integration tests for completeness. + * Tests the logical units of {@link ReportResultsTransform} and its full DAG + * expand() method using {@link FakeBigQueryServices} and a local temporary + * folder to avoid external GCS and BigQuery dependencies. */ public class ReportResultsTransformTest implements Serializable { @Rule public final transient TestPipeline pipeline = TestPipeline.create(); + @Rule public final org.junit.rules.TemporaryFolder tmp = new org.junit.rules.TemporaryFolder(); + @org.junit.Before public void setUp() throws NoSuchSchemaException { pipeline @@ -68,6 +86,75 @@ public void setUp() throws NoSuchSchemaException { pipeline.getSchemaRegistry().getFromRowFunction(ValidationSummary.class))); } + + + + @Test + public void testExpandWithTransientErrors() throws Exception { + DateTimeUtils.setCurrentMillisFixed(1000000L); + try { + FakeDatasetService.setUp(); + FakeDatasetService fakeDatasetService = new FakeDatasetService(); + FakeJobService fakeJobService = new FakeJobService(); + fakeDatasetService.createDataset("project", "dataset", "", "", null); + + TableReference matchedRef = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("ValidationSummary"); + fakeDatasetService.createTable(new Table().setTableReference(matchedRef).setSchema(BigQuerySchemas.VALIDATION_SUMMARY_SCHEMA)); + + TableReference mismatchesRef = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("MismatchedRecords"); + fakeDatasetService.createTable(new Table().setTableReference(mismatchesRef).setSchema(BigQuerySchemas.MISMATCHED_RECORDS_SCHEMA)); + + TableReference statsRef = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("TableValidationStats"); + fakeDatasetService.createTable(new Table().setTableReference(statsRef).setSchema(BigQuerySchemas.TABLE_VALIDATION_STATS_SCHEMA)); + + FakeBigQueryServices fakeBigQueryServices = new FakeBigQueryServices() + .withDatasetService(fakeDatasetService) + .withJobService(fakeJobService); + + Instant now = Instant.now(); + ReportResultsTransform transform = new ReportResultsTransform("project:dataset", "run1", now) + .withTestServices(fakeBigQueryServices, tmp.newFolder("bq-temp").getAbsolutePath()); + + ComparisonRecord matched = ComparisonRecord.builder().setTableName("Table1").setHash("hash1").setPrimaryKeyColumns(Collections.emptyList()).build(); + PCollection pMatched = pipeline.apply("CreateMatched", Create.of(matched)); + PCollection pMissingInSpanner = pipeline.apply("CreateMissingInSpanner", Create.empty(TypeDescriptor.of(ComparisonRecord.class))); + PCollection pMissingInSource = pipeline.apply("CreateMissingInSource", Create.empty(TypeDescriptor.of(ComparisonRecord.class))); + + PCollectionTuple input = PCollectionTuple.of(MATCHED_TAG, pMatched) + .and(MISSING_IN_SPANNER_TAG, pMissingInSpanner) + .and(MISSING_IN_SOURCE_TAG, pMissingInSource); + + input.apply("TestTransform", transform); + + // Inject failure + TableRow expectedSummaryRow = new TableRow() + .set(ValidationSummary.RUN_ID_COLUMN_NAME, "run1") + .set(ValidationSummary.SOURCE_DATABASE_COLUMN_NAME, ReportResultsTransform.GCS_SOURCE) + .set(ValidationSummary.DESTINATION_DATABASE_COLUMN_NAME, ReportResultsTransform.SPANNER_DESTINATION) + .set(ValidationSummary.STATUS_COLUMN_NAME, "MATCH") + .set(ValidationSummary.TOTAL_TABLES_VALIDATED_COLUMN_NAME, 1L) + .set(ValidationSummary.TABLES_WITH_MISMATCHES_COLUMN_NAME, "") + .set(ValidationSummary.TOTAL_ROWS_MATCHED_COLUMN_NAME, 1L) + .set(ValidationSummary.TOTAL_ROWS_MISMATCHED_COLUMN_NAME, 0L) + .set(ValidationSummary.START_TIMESTAMP_COLUMN_NAME, now.toString()) + .set(ValidationSummary.END_TIMESTAMP_COLUMN_NAME, now.toString()); + + InsertErrors error = new InsertErrors().setErrors(Collections.singletonList(new ErrorProto().setReason("timeout"))); + + fakeDatasetService.failOnInsert(Collections.singletonMap(expectedSummaryRow, Collections.singletonList(error))); + + + pipeline.run(); + + assertEquals(1, fakeDatasetService.getAllRows("project", "dataset", "ValidationSummary").size()); + assertEquals(1, fakeDatasetService.getAllRows("project", "dataset", "TableValidationStats").size()); + assertEquals(0, fakeDatasetService.getAllRows("project", "dataset", "MismatchedRecords").size()); + + } finally { + DateTimeUtils.setCurrentMillisSystem(); + } + } + @Test public void testTransformMismatchedRecords() { Instant now = Instant.now();