Flex templates of same dependencies sharing container image

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/TemplateTestBase.java

@@ -26,6 +26,8 @@
 import com.google.api.services.dataflow.model.Job;
 import com.google.auth.Credentials;
 import com.google.cloud.bigquery.TableId;
+import com.google.cloud.storage.Blob;
+import com.google.cloud.storage.Storage;
 import com.google.cloud.teleport.metadata.DirectRunnerTest;
 import com.google.cloud.teleport.metadata.MultiTemplateIntegrationTest;
 import com.google.cloud.teleport.metadata.SkipRunnerV2Test;
@@ -52,6 +54,7 @@
 import java.util.Date;
 import java.util.List;
 import java.util.Random;
+import java.util.UUID;
 import java.util.concurrent.ExecutionException;
 import org.apache.beam.it.common.PipelineLauncher;
 import org.apache.beam.it.common.PipelineLauncher.JobState;
@@ -75,6 +78,7 @@
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification;
 import org.apache.commons.lang3.RandomStringUtils;
+import org.apache.parquet.Strings;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.junit.After;
 import org.junit.Before;
@@ -142,7 +146,14 @@ protected void starting(Description description) {
               })
           .build();
 
+  public static final String STAGING_PREFIX;
+
   static {
+    STAGING_PREFIX =
+        new SimpleDateFormat("yyyy-MM-dd-HH-mm").format(new Date())
+            + "-"
+            + UUID.randomUUID().toString().substring(0, 6)
+            + "_IT";
     Runtime.getRuntime().addShutdownHook(new Thread(stagedTemplates::invalidateAll));
   }
 
@@ -303,36 +314,53 @@ private String getSpecPath(
       LOG.info("A spec path was given, not staging template {}", templateMetadata.name());
       return TestProperties.specPath();
     } else {
-      return stagedTemplates.get(
-          templateMetadata.name(),
+      boolean flex = !Strings.isNullOrEmpty(templateMetadata.flexContainerName());
+
+      // Use bucketName unless only artifactBucket is provided
+      String bucketName;
+      if (TestProperties.hasStageBucket()) {
+        bucketName = TestProperties.stageBucket();
+      } else if (TestProperties.hasArtifactBucket()) {
+        bucketName = TestProperties.artifactBucket();
+        LOG.warn(
+            "-DstageBucket was not specified, using -DartifactBucket ({}) for stage step",
+            bucketName);
+      } else {
+        throw new IllegalArgumentException(
+            "-DstageBucket was not specified, so Template can not be staged. Either give a"
+                + " -DspecPath or provide a proper -DstageBucket for automatic staging.");
+      }
+
+      String blobPath =
+          String.format("%s/%s%s", STAGING_PREFIX, flex ? "flex/" : "", templateMetadata.name());
+      String stagePath = String.format("gs://%s/%s", bucketName, blobPath);
+
+      String identifier = flex ? template.flexContainerName() : templateMetadata.name();
+
+      stagedTemplates.get(
+          identifier,
           () -> {
             LOG.info("Preparing test for {} ({})", templateMetadata.name(), dataflowTemplateClass);
 
-            String prefix = new SimpleDateFormat("yyyy-MM-dd-HH-mm-ss").format(new Date()) + "_IT";
-
             File pom = new File(pomPath).getAbsoluteFile();
             if (!pom.exists()) {
               throw new IllegalArgumentException(
                   "To use tests staging templates, please run in the Maven module directory"
                       + " containing the template.");
             }
 
-            // Use bucketName unless only artifactBucket is provided
-            String bucketName;
-            if (TestProperties.hasStageBucket()) {
-              bucketName = TestProperties.stageBucket();
-            } else if (TestProperties.hasArtifactBucket()) {
-              bucketName = TestProperties.artifactBucket();
-              LOG.warn(
-                  "-DstageBucket was not specified, using -DartifactBucket ({}) for stage step",
-                  bucketName);
-            } else {
-              throw new IllegalArgumentException(
-                  "-DstageBucket was not specified, so Template can not be staged. Either give a"
-                      + " -DspecPath or provide a proper -DstageBucket for automatic staging.");
+            // Check template metadata file existence
+            try (Storage storage = ArtifactUtils.createStorageClient(credentials)) {
+              Blob blob =
+                  storage.get(
+                      bucketName, blobPath, Storage.BlobGetOption.fields(Storage.BlobField.SIZE));
+              if (blob != null && blob.exists() && blob.getSize() > 0) {
+                LOG.info("Find templates at {}", stagePath);
+                return stagePath;
+              }
             }
 
-            String[] mavenCmd = buildMavenStageCommand(prefix, pom, bucketName, template);
+            String[] mavenCmd = buildMavenStageCommand(STAGING_PREFIX, pom, bucketName, template);
             LOG.info("Running command to stage templates: {}", String.join(" ", mavenCmd));
 
             try {
@@ -344,17 +372,12 @@ private String getSpecPath(
                 throw new RuntimeException("Error staging template, check Maven logs.");
               }
 
-              boolean flex =
-                  templateMetadata.flexContainerName() != null
-                      && !templateMetadata.flexContainerName().isEmpty();
-              return String.format(
-                  "gs://%s/%s/%s%s",
-                  bucketName, prefix, flex ? "flex/" : "", templateMetadata.name());
-
+              return stagePath;
             } catch (Exception e) {
               throw new IllegalArgumentException("Error staging template", e);
             }
           });
+      return stagePath;
     }
   }
 
@@ -417,6 +440,14 @@ private String[] buildMavenStageCommand(
     // that will copy only the shaded jar to the docker image.
     boolean skipShade = templateMetadata.type() != TemplateType.XLANG;
 
+    String templateOrContainer;
+    @Nullable String flexContainerName = templateMetadata.flexContainerName();
+    if (Strings.isNullOrEmpty(flexContainerName)) {
+      templateOrContainer = "-DtemplateName=" + templateMetadata.name();
+    } else {
+      templateOrContainer = "-DflexContainerName=" + flexContainerName;
+    }
+
     return new String[] {
       "mvn",
       "compile",
@@ -442,7 +473,7 @@ private String[] buildMavenStageCommand(
       "-DbucketName=" + bucketName,
       "-DgcpTempLocation=" + bucketName,
       "-DstagePrefix=" + prefix,
-      "-DtemplateName=" + templateMetadata.name(),
+      templateOrContainer,
       "-DunifiedWorker=" + System.getProperty("unifiedWorker"),
       // Print stacktrace when command fails
       "-e"
@@ -814,7 +845,7 @@ private static void cleanUpTemplates(String metafileName) {
           if (cmd != null) {
             Process exec = Runtime.getRuntime().exec(cmd);
             if (exec.waitFor() != 0) {
-              LOG.warn("Error deleting staged image {}", imgName);
+              LOG.warn("Error deleting staged image {}. It might already be deleted.", imgName);
             }
           }
         } else {

it/google-cloud-platform/src/main/java/org/apache/beam/it/gcp/dataflow/FlexTemplateDataflowJobResourceManager.java

@@ -36,6 +36,11 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+/**
+ * Used when additional flex template is needed for integration tests (e.g. using another template
+ * to generate data). For generic template integration test, Use TemplateTestBase's subclasses to
+ * manage the templates.
+ */
 public class FlexTemplateDataflowJobResourceManager implements ResourceManager {
 
   private static final Logger LOG =
@@ -49,6 +54,9 @@ public class FlexTemplateDataflowJobResourceManager implements ResourceManager {
   private static final String PROJECT = TestProperties.project();
   private static final String REGION = TestProperties.region();
   private static final Credentials CREDENTIALS = TestProperties.googleCredentials();
+  // TODO(yathu): we should use TemplateTestBase.stagedTemplates to managed all staged templates
+  // during workflow run.
+  //  Currently templates involved here get compiled and staged twice.
   private static Map<String, String> specPaths = new HashMap<>();
 
   private FlexTemplateDataflowJobResourceManager(Builder builder) {
@@ -173,6 +181,11 @@ public FlexTemplateDataflowJobResourceManager build() {
     }
   }
 
+  // TODO(yathu) this method was forked and diverged from TemplateTestBase.buildAndStageTemplate,
+  // causing involved
+  //  templates get compiled and staged twice. We should use TemplateTestBase.stagedTemplates to
+  // managed all staged
+  //  templates during workflow run.
   private void buildAndStageTemplate(
       String templateName, String modulePath, String additionalMavenProfile) {
     LOG.info("Building and Staging {} template", templateName);

plugins/core-plugin/src/main/java/com/google/cloud/teleport/plugin/TemplateSpecsGenerator.java

@@ -26,7 +26,6 @@
 import java.io.File;
 import java.io.FileWriter;
 import java.io.IOException;
-import java.nio.file.Path;
 import java.util.List;
 import java.util.logging.Logger;
 import org.apache.commons.lang3.StringUtils;
@@ -121,14 +120,6 @@ public File saveMetadata(
     }
 
     String imageName = templateDash.toLowerCase();
-    if (StringUtils.isNotEmpty(templateAnnotation.flexContainerName())) {
-      imageName = Path.of(templateAnnotation.flexContainerName()).getFileName().toString();
-    }
-
-    if (!targetDirectory.exists()) {
-      targetDirectory.mkdirs();
-    }
-
     File file = new File(targetDirectory, imageName + "-generated-metadata.json");
     LOG.info("Saving image spec metadata " + file.getAbsolutePath());
 

plugins/core-plugin/src/test/java/com/google/cloud/teleport/plugin/TemplateSpecsGeneratorTest.java

@@ -105,8 +105,8 @@ public void saveMetadataNestedFlex() {
     assertNotNull(saveMetadata);
     assertTrue(saveMetadata.exists());
     assertEquals(
-        saveMetadata.getPath(),
-        outputFolder.toPath().resolve("AtoBNestedFlex-generated-metadata.json").toString());
+        outputFolder.toPath().resolve("atobnestedflex-generated-metadata.json").toString(),
+        saveMetadata.getPath());
   }
 
   @Test

plugins/templates-maven-plugin/src/main/java/com/google/cloud/teleport/plugin/maven/TemplatesReleaseMojo.java

@@ -67,6 +67,9 @@ public class TemplatesReleaseMojo extends TemplatesBaseMojo {
   @Parameter(defaultValue = "${templateName}", readonly = true, required = false)
   protected String templateName;
 
+  @Parameter(defaultValue = "${flexContainerName}", readonly = true, required = false)
+  protected String flexContainerName;
+
   @Parameter(defaultValue = "${bucketName}", readonly = true, required = true)
   protected String bucketName;
 
@@ -202,61 +205,54 @@ public void execute() throws MojoExecutionException {
             "Stage Prefix must be informed for releases, when releasing templates or yaml blueprints.");
       }
 
-      if (!templateDefinitions.isEmpty()) {
-        LOG.info("Found {} templates to release.", templateDefinitions.size());
-        LOG.info("Trying to stage templates...");
-
-        for (TemplateDefinitions definition : templateDefinitions) {
-
-          ImageSpec imageSpec = definition.buildSpecModel(true);
-          String currentTemplateName = imageSpec.getMetadata().getName();
-
-          LOG.info("Staging template {}...", currentTemplateName);
-
-          String useRegion = StringUtils.isNotEmpty(region) ? region : "us-central1";
-
-          // TODO: is there a better way to get the plugin on the _same project_?
-          TemplatesStageMojo configuredMojo =
-              new TemplatesStageMojo(
-                  project,
-                  session,
-                  outputDirectory,
-                  outputClassesDirectory,
-                  resourcesDirectory,
-                  targetDirectory,
-                  projectId,
-                  templateName,
-                  bucketName,
-                  librariesBucketName,
-                  stagePrefix,
-                  useRegion,
-                  artifactRegion,
-                  gcpTempLocation,
-                  baseContainerImage,
-                  basePythonContainerImage,
-                  pythonTemplateLauncherEntryPoint,
-                  javaTemplateLauncherEntryPoint,
-                  pythonVersion,
-                  beamVersion,
-                  artifactRegistry,
-                  stagingArtifactRegistry,
-                  unifiedWorker,
-                  generateSBOM);
-
-          String templatePath = configuredMojo.stageTemplate(definition, imageSpec, pluginManager);
-
-          if (!definition.getTemplateAnnotation().stageImageOnly()) {
-            LOG.info("Template staged: {}", templatePath);
-
-            // Export the specs for collection
-            generator.saveMetadata(definition, imageSpec.getMetadata(), targetDirectory);
-            if (definition.isFlex()) {
-              generator.saveImageSpec(definition, imageSpec, targetDirectory);
-            }
+      String useRegion = StringUtils.isNotEmpty(region) ? region : "us-central1";
+      TemplatesStageMojo configuredMojo =
+          new TemplatesStageMojo(
+              project,
+              session,
+              outputDirectory,
+              outputClassesDirectory,
+              resourcesDirectory,
+              targetDirectory,
+              projectId,
+              templateName,
+              flexContainerName,
+              bucketName,
+              librariesBucketName,
+              stagePrefix,
+              useRegion,
+              artifactRegion,
+              gcpTempLocation,
+              baseContainerImage,
+              basePythonContainerImage,
+              pythonTemplateLauncherEntryPoint,
+              javaTemplateLauncherEntryPoint,
+              pythonVersion,
+              beamVersion,
+              artifactRegistry,
+              stagingArtifactRegistry,
+              unifiedWorker,
+              generateSBOM);
+      configuredMojo.stageCommandSpecs(templateDefinitions);
+
+      for (TemplateDefinitions definition : templateDefinitions) {
+
+        ImageSpec imageSpec = definition.buildSpecModel(true);
+        String currentTemplateName = imageSpec.getMetadata().getName();
+
+        LOG.info("Staging template {}...", currentTemplateName);
+
+        String templatePath = configuredMojo.stageTemplate(definition, imageSpec, pluginManager);
+
+        if (!definition.getTemplateAnnotation().stageImageOnly()) {
+          LOG.info("Template staged: {}", templatePath);
+
+          // Export the specs for collection
+          generator.saveMetadata(definition, imageSpec.getMetadata(), targetDirectory);
+          if (definition.isFlex()) {
+            generator.saveImageSpec(definition, imageSpec, targetDirectory);
           }
         }
-      } else {
-        LOG.warn("Did not find any templates to release in this module.");
       }
 
       if (publishYamlBlueprints) {

plugins/templates-maven-plugin/src/main/java/com/google/cloud/teleport/plugin/maven/TemplatesRunMojo.java

@@ -74,6 +74,9 @@ public class TemplatesRunMojo extends TemplatesBaseMojo {
   @Parameter(defaultValue = "${templateName}", readonly = true, required = false)
   protected String templateName;
 
+  @Parameter(defaultValue = "${flexContainerName}", readonly = true, required = false)
+  protected String flexContainerName;
+
   @Parameter(defaultValue = "${bucketName}", readonly = true, required = true)
   protected String bucketName;
 
@@ -203,6 +206,7 @@ public void execute() throws MojoExecutionException {
               targetDirectory,
               projectId,
               templateName,
+              flexContainerName,
               bucketName,
               bucketName,
               stagePrefix,