diff --git a/v2/failure-injection-policies/src/main/java/com/google/cloud/teleport/v2/failureinjection/InitialLimitedDurationErrorInjectionPolicy.java b/v2/failure-injection-policies/src/main/java/com/google/cloud/teleport/v2/failureinjection/InitialLimitedDurationErrorInjectionPolicy.java index fffbe04c40..94feabad77 100644 --- a/v2/failure-injection-policies/src/main/java/com/google/cloud/teleport/v2/failureinjection/InitialLimitedDurationErrorInjectionPolicy.java +++ b/v2/failure-injection-policies/src/main/java/com/google/cloud/teleport/v2/failureinjection/InitialLimitedDurationErrorInjectionPolicy.java @@ -22,6 +22,7 @@ import java.time.Duration; import java.time.Instant; import java.time.format.DateTimeParseException; +import java.util.concurrent.atomic.AtomicLong; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -37,12 +38,12 @@ public class InitialLimitedDurationErrorInjectionPolicy LoggerFactory.getLogger(InitialLimitedDurationErrorInjectionPolicy.class); private static final long serialVersionUID = 1L; - private Instant startTime; + private static volatile Instant startTime = null; + private static final AtomicLong callCount = new AtomicLong(0); private final Duration injectionDuration; private final String effectiveDurationParameter; private String errorCodeToBeInjected; private Clock clock; - private long callCount; private static final String DEFAULT_DURATION = "PT10M"; private static final String DURATION_FIELD_IN_OBJECT = "duration"; @@ -124,22 +125,20 @@ public InitialLimitedDurationErrorInjectionPolicy(JsonNode inputParameter, Clock */ @Override public boolean shouldInjectionError() { - if (this.startTime == null) { - synchronized (this) { - if (this.startTime == null) { - this.startTime = Instant.now(clock); + if (startTime == null) { + synchronized (InitialLimitedDurationErrorInjectionPolicy.class) { + if (startTime == null) { + startTime = Instant.now(clock); LOG.info( "First call detected. Errors will be injected for {} starting from {}.", this.injectionDuration, - this.startTime); + startTime); } } } - synchronized (this) { - ++callCount; - } + long currentCallCount = callCount.incrementAndGet(); - if (callCount < INITIAL_ALLOWED_CALLS_COUNT) { + if (currentCallCount < INITIAL_ALLOWED_CALLS_COUNT) { return false; } @@ -186,6 +185,11 @@ void setClockForTesting(Clock clock) { this.clock = clock; } + public static void resetForTesting() { + startTime = null; + callCount.set(0); + } + @Override public String toString() { return "InitialLimitedDurationErrorInjectionPolicy{" diff --git a/v2/failure-injection-policies/src/test/java/com/google/cloud/teleport/v2/failureinjection/InitialLimitedDurationErrorInjectionPolicyTest.java b/v2/failure-injection-policies/src/test/java/com/google/cloud/teleport/v2/failureinjection/InitialLimitedDurationErrorInjectionPolicyTest.java index fe4a3f8759..407ce12f3d 100644 --- a/v2/failure-injection-policies/src/test/java/com/google/cloud/teleport/v2/failureinjection/InitialLimitedDurationErrorInjectionPolicyTest.java +++ b/v2/failure-injection-policies/src/test/java/com/google/cloud/teleport/v2/failureinjection/InitialLimitedDurationErrorInjectionPolicyTest.java @@ -28,10 +28,16 @@ import java.time.Duration; import java.time.Instant; import java.time.ZoneOffset; +import org.junit.Before; import org.junit.Test; public class InitialLimitedDurationErrorInjectionPolicyTest { + @Before + public void setUp() { + InitialLimitedDurationErrorInjectionPolicy.resetForTesting(); + } + private ObjectNode createInputObject(String duration) { ObjectNode node = JsonNodeFactory.instance.objectNode(); if (duration != null) { @@ -268,4 +274,26 @@ public void shouldInjectError_startTimeDoesNotChangeAfterFirstCall() { Instant thirdStartTime = policy.getStartTime(); assertEquals("Start time should still not change", firstStartTime, thirdStartTime); } + + @Test + public void constructor_shouldParseErrorCode() { + ObjectNode input = JsonNodeFactory.instance.objectNode(); + input.put("duration", "PT5S"); + input.put("errorCode", "UNAVAILABLE"); + Clock clock = Clock.fixed(Instant.EPOCH, ZoneOffset.UTC); + InitialLimitedDurationErrorInjectionPolicy policy = + new InitialLimitedDurationErrorInjectionPolicy(input, clock); + + assertEquals("UNAVAILABLE", policy.getErrorCodeToBeInjected()); + } + + @Test + public void constructor_shouldUseDefaultErrorCodeIfBlank() { + ObjectNode input = JsonNodeFactory.instance.objectNode(); + input.put("duration", "PT5S"); + input.put("errorCode", " "); + Clock clock = Clock.fixed(Instant.EPOCH, ZoneOffset.UTC); + InitialLimitedDurationErrorInjectionPolicy policy = + new InitialLimitedDurationErrorInjectionPolicy(input, clock); + } } diff --git a/v2/gcs-spanner-dv/pom.xml b/v2/gcs-spanner-dv/pom.xml index 01eecdeec1..0fc12753e6 100644 --- a/v2/gcs-spanner-dv/pom.xml +++ b/v2/gcs-spanner-dv/pom.xml @@ -95,9 +95,44 @@ parent POM(s). --> com/google/cloud/teleport/v2/dto/** com/google/cloud/teleport/v2/constants/** + com/google/cloud/teleport/v2/templates/GCSSpannerDV.class + + + useRealSpanner + + true + + !activateFailureInjection + + + + + com.google.cloud.teleport.v2 + real-spanner-service + ${project.version} + + + + + failureInjectionTest + + + activateFailureInjection + true + + + + + com.google.cloud.teleport.v2 + failure-injected-spanner-service + ${project.version} + + + + diff --git a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/templates/GCSSpannerDV.java b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/templates/GCSSpannerDV.java index 35a6011291..6da4c19c84 100644 --- a/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/templates/GCSSpannerDV.java +++ b/v2/gcs-spanner-dv/src/main/java/com/google/cloud/teleport/v2/templates/GCSSpannerDV.java @@ -37,6 +37,7 @@ import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.io.gcp.spanner.SpannerConfig; +import org.apache.beam.sdk.io.gcp.spanner.SpannerServiceFactoryImpl; import org.apache.beam.sdk.options.Default; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; @@ -254,6 +255,16 @@ public interface Options extends PipelineOptions { String getTransformationCustomParameters(); void setTransformationCustomParameters(String value); + + @TemplateParameter.Text( + order = 16, + optional = true, + description = "Failure injection parameter", + helpText = "Failure injection parameter. Only used for testing.") + @Default.String("") + String getFailureInjectionParameter(); + + void setFailureInjectionParameter(String value); } public static void main(String[] args) { @@ -327,11 +338,20 @@ public static PipelineResult run(Options options) { @VisibleForTesting static SpannerConfig createSpannerConfig(Options options) { - return SpannerConfig.create() - .withProjectId(ValueProvider.StaticValueProvider.of(options.getProjectId())) - .withHost(ValueProvider.StaticValueProvider.of(options.getSpannerHost())) - .withInstanceId(ValueProvider.StaticValueProvider.of(options.getInstanceId())) - .withDatabaseId(ValueProvider.StaticValueProvider.of(options.getDatabaseId())) - .withRpcPriority(ValueProvider.StaticValueProvider.of(options.getSpannerPriority())); + SpannerConfig config = + SpannerConfig.create() + .withProjectId(ValueProvider.StaticValueProvider.of(options.getProjectId())) + .withHost(ValueProvider.StaticValueProvider.of(options.getSpannerHost())) + .withInstanceId(ValueProvider.StaticValueProvider.of(options.getInstanceId())) + .withDatabaseId(ValueProvider.StaticValueProvider.of(options.getDatabaseId())) + .withRpcPriority(ValueProvider.StaticValueProvider.of(options.getSpannerPriority())); + + if (options.getFailureInjectionParameter() != null + && !options.getFailureInjectionParameter().isEmpty()) { + config = + SpannerServiceFactoryImpl.createSpannerService( + config, options.getFailureInjectionParameter()); + } + return config; } } diff --git a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVFTBase.java b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVFTBase.java new file mode 100644 index 0000000000..c6c7694a37 --- /dev/null +++ b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVFTBase.java @@ -0,0 +1,121 @@ +/* + * Copyright (C) 2026 Google LLC + * + * Licensed 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 com.google.cloud.teleport.v2.templates; + +import com.google.cloud.teleport.v2.spanner.migrations.transformation.CustomTransformation; +import com.google.common.io.Resources; +import java.io.IOException; +import java.util.Map; +import org.apache.beam.it.common.PipelineLauncher.LaunchInfo; +import org.apache.beam.it.common.utils.PipelineUtils; +import org.apache.beam.it.gcp.dataflow.FlexTemplateDataflowJobResourceManager; +import org.apache.beam.it.gcp.spanner.SpannerResourceManager; + +/** + * Base class for gcs-spanner-dv failure injection integration tests. + * + *

Why is this a separate class from {@link GCSSpannerDVITBase}? + * + *

While {@code GCSSpannerDVITBase} is runner-agnostic and relies on the generic {@code + * PipelineLauncher} (allowing tests to run locally via DirectRunner), failure injection testing + * explicitly requires building a custom Docker image with the {@code failureInjectionTest} Maven + * profile. + * + *

Therefore, tests extending this class are strictly coupled to Dataflow Flex Templates and + * bypass the generic launcher in favor of {@link FlexTemplateDataflowJobResourceManager}. + */ +public abstract class GCSSpannerDVFTBase extends GCSSpannerDVITBase { + + /** + * Launches the Dataflow job with failure injection testing capabilities using + * FlexTemplateDataflowJobResourceManager. + */ + protected LaunchInfo launchFTDataflowJob( + String testId, + String projectId, + SpannerResourceManager spannerResourceManager, + String bigQueryDataset, + String gcsInputDirectory, + String sessionFileResourceName, + String schemaOverridesFileResourceName, + String tableOverrides, + String columnOverrides, + CustomTransformation customTransformation, + String failureInjectionParameter, + Map jobParameters) + throws IOException { + + FlexTemplateDataflowJobResourceManager.Builder flexTemplateBuilder = + FlexTemplateDataflowJobResourceManager.builder(testId) + .withTemplateName("GCS_Spanner_Data_Validator") + .withTemplateModulePath("v2/gcs-spanner-dv") + .withAdditionalMavenProfile("failureInjectionTest") + .addEnvironmentVariable( + "additionalExperiments", java.util.Collections.singletonList("disable_runner_v2")); + + if (failureInjectionParameter != null && !failureInjectionParameter.isEmpty()) { + flexTemplateBuilder.addParameter("failureInjectionParameter", failureInjectionParameter); + } + + flexTemplateBuilder.addParameter("projectId", projectId); + flexTemplateBuilder.addParameter("instanceId", spannerResourceManager.getInstanceId()); + flexTemplateBuilder.addParameter("databaseId", spannerResourceManager.getDatabaseId()); + flexTemplateBuilder.addParameter("bigQueryDataset", bigQueryDataset); + flexTemplateBuilder.addParameter("gcsInputDirectory", gcsInputDirectory); + + if (sessionFileResourceName != null) { + gcsClient.uploadArtifact( + "session.json", Resources.getResource(sessionFileResourceName).getPath()); + flexTemplateBuilder.addParameter("sessionFilePath", getGcsPath("session.json")); + } + + if (schemaOverridesFileResourceName != null) { + gcsClient.uploadArtifact( + "schema_overrides.json", + Resources.getResource(schemaOverridesFileResourceName).getPath()); + flexTemplateBuilder.addParameter( + "schemaOverridesFilePath", getGcsPath("schema_overrides.json")); + } + + if (tableOverrides != null) { + flexTemplateBuilder.addParameter("tableOverrides", tableOverrides); + } + + if (columnOverrides != null) { + flexTemplateBuilder.addParameter("columnOverrides", columnOverrides); + } + + if (customTransformation != null) { + flexTemplateBuilder.addParameter( + "transformationJarPath", getGcsPath(customTransformation.jarPath())); + flexTemplateBuilder.addParameter("transformationClassName", customTransformation.classPath()); + if (customTransformation.customParameters() != null) { + flexTemplateBuilder.addParameter( + "transformationCustomParameters", customTransformation.customParameters()); + } + } + + String runId = PipelineUtils.createJobName(testId); + flexTemplateBuilder.addParameter("runId", runId); + flexTemplateBuilder.addParameter("workerMachineType", "n2-standard-4"); + + if (jobParameters != null) { + jobParameters.forEach(flexTemplateBuilder::addParameter); + } + + return flexTemplateBuilder.build().launchJob(); + } +} diff --git a/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVSpannerReadFT.java b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVSpannerReadFT.java new file mode 100644 index 0000000000..7a1fe3186f --- /dev/null +++ b/v2/gcs-spanner-dv/src/test/java/com/google/cloud/teleport/v2/templates/GCSSpannerDVSpannerReadFT.java @@ -0,0 +1,163 @@ +/* + * Copyright (C) 2026 Google LLC + * + * Licensed 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 com.google.cloud.teleport.v2.templates; + +import com.google.cloud.Timestamp; +import com.google.cloud.spanner.Mutation; +import com.google.cloud.teleport.metadata.SkipDirectRunnerTest; +import com.google.cloud.teleport.metadata.TemplateIntegrationTest; +import com.google.cloud.teleport.v2.templates.GCSSpannerDVAvroSetupHelper.RecordBuilder; +import com.google.cloud.teleport.v2.templates.GCSSpannerDVAvroSetupHelper.TableDef; +import com.google.cloud.teleport.v2.templates.GCSSpannerDVTestAsserts.TableValidationStatsDto; +import com.google.cloud.teleport.v2.templates.GCSSpannerDVTestAsserts.ValidationSummaryDto; +import java.io.IOException; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import org.apache.avro.generic.GenericRecord; +import org.apache.beam.it.common.PipelineLauncher.LaunchInfo; +import org.junit.Before; +import org.junit.Test; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests transient Spanner read failures (via SpannerIO) for GCSSpannerDV pipeline. + * + *

Test cases covered: + * + *

+ */ +@Category({TemplateIntegrationTest.class, SkipDirectRunnerTest.class}) +@TemplateIntegrationTest(GCSSpannerDV.class) +@RunWith(JUnit4.class) +public class GCSSpannerDVSpannerReadFT extends GCSSpannerDVFTBase { + + private static final String SPANNER_DDL_RESOURCE = "GCSSpannerDVSpannerReadFT/spanner-schema.sql"; + private static final int NUM_RECORDS = 500; + + @Before + public void setUp() throws IOException { + spannerResourceManager = setUpSpannerResourceManager(); + createSpannerDDL(spannerResourceManager, SPANNER_DDL_RESOURCE); + bigQueryResourceManager = setUpBigQueryResourceManager(); + bigQueryResourceManager.createDataset(REGION); + } + + // Tests transient UNAVAILABLE errors on SpannerIO.readAll() to ensure native exponential backoff + // succeeds. + @Test + public void testTransientReadFailure() throws IOException, InterruptedException { + List records = new ArrayList<>(); + List mutations = new ArrayList<>(); + Instant now = Instant.now().truncatedTo(java.time.temporal.ChronoUnit.MILLIS); + + // We insert 500 rows to ensure Dataflow has enough data to potentially split bundles + // and exercise the SpannerIO read logic across multiple task execution boundaries, + // rather than trivially passing with a single row. + for (int i = 0; i < NUM_RECORDS; i++) { + long userId = (long) i; + String eventId = "E" + i; + String fullName = "User " + i; + int age = 20 + (i % 30); + + // Avro Record + GenericRecord record = + new RecordBuilder(TableDef.USERS, null) + .set("user_id", userId) + .set("event_id", eventId) + .set("full_name", fullName) + .set("age", age) + .set("created_at", now) + .build(); + records.add(record); + + // Spanner Mutation + mutations.add( + Mutation.newInsertOrUpdateBuilder("Users") + .set("user_id") + .to(userId) + .set("event_id") + .to(eventId) + .set("full_name") + .to(fullName) + .set("age") + .to(age) + .set("created_at") + .to(Timestamp.ofTimeSecondsAndNanos(now.getEpochSecond(), now.getNano())) + .build()); + } + + String gcsInputDirectory = getGcsPath("input"); + uploadAvroFileToGcs("input/users.avro", TableDef.USERS.schema, records); + spannerResourceManager.write(mutations); + + // Injects a 60-second UNAVAILABLE outage specifically on the Spanner workers to trigger + // Dataflow task retries. + String failureInjectionParam = + "{\"policyType\":\"InitialLimitedDurationErrorInjectionPolicy\", \"policyInput\": {\"duration\":\"PT1M\", \"errorCode\":\"UNAVAILABLE\"}}"; + String bqDatasetId = bigQueryResourceManager.getDatasetId(); + + LaunchInfo jobInfo = + launchFTDataflowJob( + testName, + PROJECT, + spannerResourceManager, + bqDatasetId, + gcsInputDirectory, + null, + null, + null, + null, + null, + failureInjectionParam, + new HashMap<>()); + + pipelineOperator().waitUntilDone(createConfig(jobInfo)); + + GCSSpannerDVTestAsserts.assertValidationSummary( + bigQueryResourceManager, + Arrays.asList( + new ValidationSummaryDto( + /* status= */ "MATCH", + /* totalTablesValidated= */ 1L, + /* totalRowsMatched= */ 500L, + /* totalRowsMismatched= */ 0L, + /* tablesWithMismatches= */ ""))); + + GCSSpannerDVTestAsserts.assertTableValidationStats( + bigQueryResourceManager, + Arrays.asList( + new TableValidationStatsDto( + /* schemaName= */ null, + /* tableName= */ "Users", + /* status= */ "MATCH", + /* sourceRowCount= */ 500L, + /* destinationRowCount= */ 500L, + /* matchedRowCount= */ 500L, + /* mismatchRowCount= */ 0L))); + + // No mismatched records should exist + GCSSpannerDVTestAsserts.assertMismatchedRecords(bigQueryResourceManager, Arrays.asList()); + } +} diff --git a/v2/gcs-spanner-dv/src/test/resources/GCSSpannerDVSpannerReadFT/spanner-schema.sql b/v2/gcs-spanner-dv/src/test/resources/GCSSpannerDVSpannerReadFT/spanner-schema.sql new file mode 100644 index 0000000000..815ce495c3 --- /dev/null +++ b/v2/gcs-spanner-dv/src/test/resources/GCSSpannerDVSpannerReadFT/spanner-schema.sql @@ -0,0 +1,12 @@ +CREATE TABLE Users ( + user_id INT64 NOT NULL, + event_id STRING(MAX) NOT NULL, + full_name STRING(MAX), + age INT64, + created_at TIMESTAMP +) PRIMARY KEY (user_id, event_id); + +CREATE TABLE AccountRoles ( + role_id INT64 NOT NULL, + role_name STRING(MAX) +) PRIMARY KEY (role_id);