Skip to content

[Data Validaion] BigQuery Failure UT #4122

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
aasthabharill wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[Data Validaion] BigQuery Failure UT #4122

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 28 additions & 12 deletions ...nner-dv/src/main/java/com/google/cloud/teleport/v2/transforms/ReportResultsTransform.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -74,6 +75,15 @@ public class ReportResultsTransform extends PTransform<PCollectionTuple, PDone>
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;
Expand All @@ -90,9 +100,7 @@ public ReportResultsTransform(String bigQueryDataset, String runId, Instant star
PCollection<MismatchedRecord> allMismatches =
transformMismatchedRecords(missingInSpanner, missingInSource);

allMismatches.apply(
"WriteMismatchedRecords",
BigQueryIO.<MismatchedRecord>write()
BigQueryIO.Write<MismatchedRecord> writeMismatches = BigQueryIO.<MismatchedRecord>write()
.to(String.format("%s.%s", bigQueryDataset, MISMATCHED_RECORDS_TABLE))
.withSchema(BigQuerySchemas.MISMATCHED_RECORDS_SCHEMA)
.withFormatFunction(
Expand All @@ -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<TableValidationStats> tableStats =
calculateTableStats(matched, missingInSpanner, missingInSource);

tableStats.apply(
"WriteTableStats",
BigQueryIO.<TableValidationStats>write()
BigQueryIO.Write<TableValidationStats> writeTableStats = BigQueryIO.<TableValidationStats>write()
.to(String.format("%s.%s", bigQueryDataset, TABLE_VALIDATION_STATS_TABLE))
.withSchema(BigQuerySchemas.TABLE_VALIDATION_STATS_SCHEMA)
.withFormatFunction(
Expand Down Expand Up @@ -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
/**
Expand All @@ -172,9 +186,7 @@ public ReportResultsTransform(String bigQueryDataset, String runId, Instant star

PCollection<ValidationSummary> validationSummary = calculateValidationSummary(tableStats);

validationSummary.apply(
"WriteValidationSummary",
BigQueryIO.<ValidationSummary>write()
BigQueryIO.Write<ValidationSummary> writeValidationSummary = BigQueryIO.<ValidationSummary>write()
.to(String.format("%s.%s", this.bigQueryDataset, VALIDATION_SUMMARY_TABLE))
.withSchema(BigQuerySchemas.VALIDATION_SUMMARY_SCHEMA)
.withFormatFunction(
Expand Down Expand Up @@ -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());
}
Expand Down
99 changes: 93 additions & 6 deletions ...-dv/src/test/java/com/google/cloud/teleport/v2/transforms/ReportResultsTransformTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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
Expand All @@ -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<ComparisonRecord> pMatched = pipeline.apply("CreateMatched", Create.of(matched));
PCollection<ComparisonRecord> pMissingInSpanner = pipeline.apply("CreateMissingInSpanner", Create.empty(TypeDescriptor.of(ComparisonRecord.class)));
PCollection<ComparisonRecord> 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();
Expand Down
Loading