feat: Denylist small worker sizes in reverse replication (#3044)

* Add check for small machines

* Remove log

* Apply spotless

* Use Preconditions

* Update IT and LT base classes with workerMachineType

* Add to params instead of env

Author: manitgupta

v2/spanner-common/src/main/java/com/google/cloud/teleport/v2/spanner/migrations/utils/DataflowWorkerMachineTypeValidator.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2025 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.spanner.migrations.utils;
+
+import com.google.common.base.Preconditions;
+import org.apache.commons.lang3.StringUtils;
+
+public class DataflowWorkerMachineTypeValidator {
+
+  public static void validateMachineSpecs(String workerMachineType, Integer minCPUs) {
+    Preconditions.checkArgument(
+        workerMachineType != null && !StringUtils.isBlank(workerMachineType),
+        "Policy Violation: You must specify a workerMachineType with at least %s vCPUs.",
+        minCPUs);
+
+    // Handle custom machine types first, format is custom-{vCPU}-{RAM}
+    if (workerMachineType.startsWith("custom-")) {
+      String[] parts = workerMachineType.split("-");
+      Preconditions.checkArgument(
+          parts.length == 3,
+          "Invalid custom machine type format: '%s'. Expected format: custom-{vCPU}-{RAM}.",
+          workerMachineType);
+      Integer vCpus = null;
+      try {
+        vCpus = Integer.parseInt(parts[1]);
+      } catch (NumberFormatException e) {
+        Preconditions.checkArgument(
+            false, "Invalid vCPU number in custom machine type: '%s'", workerMachineType);
+      }
+      Preconditions.checkArgument(
+          vCpus >= minCPUs,
+          "Policy Violation: Custom machine type '%s' has %s vCPUs. Minimum allowed is %s. Please use a higher machine type.",
+          workerMachineType,
+          vCpus,
+          minCPUs);
+    } else {
+      // Handle standard machine types.
+      java.util.regex.Pattern pattern = java.util.regex.Pattern.compile(".*-(\\d+)$");
+      java.util.regex.Matcher matcher = pattern.matcher(workerMachineType);
+
+      if (matcher.find()) {
+        Integer vCpus = null;
+        try {
+          vCpus = Integer.parseInt(matcher.group(1));
+        } catch (NumberFormatException e) {
+          Preconditions.checkArgument(
+              false, "Invalid vCPU number in machine type: '%s'", workerMachineType);
+        }
+        Preconditions.checkArgument(
+            vCpus >= minCPUs,
+            "Policy Violation: Machine type '%s' has %s vCPUs. Minimum allowed is %s.",
+            workerMachineType,
+            vCpus,
+            minCPUs);
+      } else {
+        Preconditions.checkArgument(
+            false,
+            "Unknown machine type format: '%s'. Please use a standard machine type (e.g., n1-standard-4) or a custom machine type (e.g., custom-4-4096) with at least %s vCPUs.",
+            workerMachineType,
+            minCPUs);
+      }
+    }
+  }
+}

v2/spanner-common/src/test/java/com/google/cloud/teleport/v2/spanner/migrations/utils/DataflowWorkerMachineTypeValidatorTest.java

@@ -0,0 +1,109 @@
+/*
+ * Copyright (C) 2025 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.spanner.migrations.utils;
+
+import static org.junit.Assert.assertThrows;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class DataflowWorkerMachineTypeValidatorTest {
+
+  @Test
+  public void testValidMachineType() {
+    DataflowWorkerMachineTypeValidator.validateMachineSpecs("n1-standard-4", 4);
+  }
+
+  @Test
+  public void testValidMachineTypeHighCpu() {
+    DataflowWorkerMachineTypeValidator.validateMachineSpecs("n1-standard-8", 4);
+  }
+
+  @Test
+  public void testInvalidMachineTypeLowCpu() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          DataflowWorkerMachineTypeValidator.validateMachineSpecs("n1-standard-2", 4);
+        });
+  }
+
+  @Test
+  public void testNullMachineType() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          DataflowWorkerMachineTypeValidator.validateMachineSpecs(null, 4);
+        });
+  }
+
+  @Test
+  public void testEmptyMachineType() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          DataflowWorkerMachineTypeValidator.validateMachineSpecs(" ", 4);
+        });
+  }
+
+  @Test
+  public void testValidCustomMachineType() {
+    DataflowWorkerMachineTypeValidator.validateMachineSpecs("custom-8-12345", 4);
+  }
+
+  @Test
+  public void testValidCustomMachineTypeMinCpu() {
+    DataflowWorkerMachineTypeValidator.validateMachineSpecs("custom-4-12345", 4);
+  }
+
+  @Test
+  public void testInvalidCustomMachineTypeLowCpu() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          DataflowWorkerMachineTypeValidator.validateMachineSpecs("custom-2-12345", 4);
+        });
+  }
+
+  @Test
+  public void testInvalidCustomMachineTypeFormat() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          DataflowWorkerMachineTypeValidator.validateMachineSpecs("custom-2", 4);
+        });
+  }
+
+  @Test
+  public void testInvalidCustomMachineTypeNonNumericCpu() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          DataflowWorkerMachineTypeValidator.validateMachineSpecs("custom-abc-12345", 4);
+        });
+  }
+
+  @Test
+  public void testUnknownMachineType() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          DataflowWorkerMachineTypeValidator.validateMachineSpecs("unknown-machine-type", 4);
+        });
+  }
+}

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

@@ -43,6 +43,7 @@
 import com.google.cloud.teleport.v2.spanner.migrations.transformation.CustomTransformation;
 import com.google.cloud.teleport.v2.spanner.migrations.utils.CassandraConfigFileReader;
 import com.google.cloud.teleport.v2.spanner.migrations.utils.CassandraDriverConfigLoader;
+import com.google.cloud.teleport.v2.spanner.migrations.utils.DataflowWorkerMachineTypeValidator;
 import com.google.cloud.teleport.v2.spanner.migrations.utils.SecretManagerAccessorImpl;
 import com.google.cloud.teleport.v2.spanner.migrations.utils.ShardFileReader;
 import com.google.cloud.teleport.v2.spanner.sourceddl.CassandraInformationSchemaScanner;
@@ -578,6 +579,10 @@ public static PipelineResult run(Options options) {
               + " incease the max shard connections");
     }
 
+    String workerMachineType =
+        pipeline.getOptions().as(DataflowPipelineWorkerPoolOptions.class).getWorkerMachineType();
+    DataflowWorkerMachineTypeValidator.validateMachineSpecs(workerMachineType, 4);
+
     // Prepare Spanner config
     SpannerConfig spannerConfig =
         SpannerConfig.create()

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

@@ -231,6 +231,7 @@ public PipelineLauncher.LaunchInfo launchDataflowJob(
             put("maxNumWorkers", "1");
             put("numWorkers", "1");
             put("sourceType", sourceType);
+            put("workerMachineType", "n2-standard-4");
           }
         };
     if (jobParameters != null) {

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

@@ -216,6 +216,7 @@ public PipelineLauncher.LaunchInfo launchDataflowJob(
             put("deadLetterQueueDirectory", getGcsPath(artifactBucket, "dlq", gcsResourceManager));
             put("maxShardConnections", "100");
             put("sourceType", sourceType);
+            put("workerMachineType", "n2-standard-4");
           }
         };