From 1af8e428c7c53583a6542e61b8d4fa2d51f94805 Mon Sep 17 00:00:00 2001
From: an2x <52892974+an2x@users.noreply.github.com>
Date: Tue, 30 Jan 2024 19:02:23 -0500
Subject: [PATCH] Templates Maven Plugin: Flex template staging improvements.

---
 .../plugin/TemplateSpecsGenerator.java        | 19 ++++++++-
 .../plugin/model/ImageSpecMetadata.java       | 42 +++++++++++--------
 2 files changed, 42 insertions(+), 19 deletions(-)

diff --git a/plugins/core-plugin/src/main/java/com/google/cloud/teleport/plugin/TemplateSpecsGenerator.java b/plugins/core-plugin/src/main/java/com/google/cloud/teleport/plugin/TemplateSpecsGenerator.java
index 4c3a6aeb27..b4ad8e9d83 100644
--- a/plugins/core-plugin/src/main/java/com/google/cloud/teleport/plugin/TemplateSpecsGenerator.java
+++ b/plugins/core-plugin/src/main/java/com/google/cloud/teleport/plugin/TemplateSpecsGenerator.java
@@ -75,8 +75,25 @@ public File saveImageSpec(
         new File(targetDirectory, templateDash.toLowerCase() + "-spec-generated-metadata.json");
     LOG.info("Saving image spec " + file.getAbsolutePath());
 
+    // The serialized image spec should match com.google.api.services.dataflow.model.ContainerSpec
+    // model, ImageSpec contains some extra fields, so we'll pick up only the expected ones.
+
+    ImageSpec is = new ImageSpec();
+    is.setImage(imageSpec.getImage());
+    is.setSdkInfo(imageSpec.getSdkInfo());
+    is.setDefaultEnvironment(imageSpec.getDefaultEnvironment());
+
+    ImageSpecMetadata m = new ImageSpecMetadata();
+    m.setName(imageSpec.getMetadata().getName());
+    m.setDescription(imageSpec.getMetadata().getDescription());
+    m.setParameters(imageSpec.getMetadata().getParameters());
+    m.setStreaming(imageSpec.getMetadata().isStreaming());
+    m.setSupportsAtLeastOnce(imageSpec.getMetadata().isSupportsAtLeastOnce());
+    m.setSupportsExactlyOnce(imageSpec.getMetadata().isSupportsExactlyOnce());
+    is.setMetadata(m);
+
     try (FileWriter writer = new FileWriter(file)) {
-      writer.write(gson.toJson(imageSpec));
+      writer.write(gson.toJson(is));
     } catch (IOException e) {
       throw new RuntimeException("Error writing image spec", e);
     }
diff --git a/plugins/core-plugin/src/main/java/com/google/cloud/teleport/plugin/model/ImageSpecMetadata.java b/plugins/core-plugin/src/main/java/com/google/cloud/teleport/plugin/model/ImageSpecMetadata.java
index b6b49825ab..73bac92e40 100644
--- a/plugins/core-plugin/src/main/java/com/google/cloud/teleport/plugin/model/ImageSpecMetadata.java
+++ b/plugins/core-plugin/src/main/java/com/google/cloud/teleport/plugin/model/ImageSpecMetadata.java
@@ -27,23 +27,23 @@ public class ImageSpecMetadata {
   private String name;
   private String description;
   private String mainClass;
-  private List<ImageSpecParameter> parameters = new ArrayList<>();
-  private Map<String, String> runtimeParameters = new HashMap<>();
+  private List<ImageSpecParameter> parameters;
+  private Map<String, String> runtimeParameters;
   private ImageSpecCategory category;
   private String internalName;
   private String module;
   private String documentationLink;
   private List<String> requirements;
   private List<ImageSpecAdditionalDocumentation> additionalDocumentation;
-  private boolean googleReleased;
-  private boolean preview;
-  private boolean udfSupport;
-  private boolean flexTemplate;
+  private Boolean googleReleased;
+  private Boolean preview;
+  private Boolean udfSupport;
+  private Boolean flexTemplate;
   private String sourceFilePath;
-  private boolean hidden;
-  private boolean streaming;
-  private boolean supportsAtLeastOnce;
-  private boolean supportsExactlyOnce;
+  private Boolean hidden;
+  private Boolean streaming;
+  private Boolean supportsAtLeastOnce;
+  private Boolean supportsExactlyOnce;
 
   public String getDescription() {
     return description;
@@ -70,6 +70,9 @@ public void setMainClass(String mainClass) {
   }
 
   public List<ImageSpecParameter> getParameters() {
+    if (parameters == null) {
+      parameters = new ArrayList<>();
+    }
     return parameters;
   }
 
@@ -78,6 +81,9 @@ public void setParameters(List<ImageSpecParameter> parameters) {
   }
 
   public Map<String, String> getRuntimeParameters() {
+    if (runtimeParameters == null) {
+      runtimeParameters = new HashMap<>();
+    }
     return runtimeParameters;
   }
 
@@ -135,7 +141,7 @@ public void setModule(String module) {
   }
 
   public boolean isGoogleReleased() {
-    return googleReleased;
+    return googleReleased != null && googleReleased;
   }
 
   public void setGoogleReleased(boolean googleReleased) {
@@ -143,7 +149,7 @@ public void setGoogleReleased(boolean googleReleased) {
   }
 
   public boolean isPreview() {
-    return preview;
+    return preview != null && preview;
   }
 
   public void setPreview(boolean preview) {
@@ -151,7 +157,7 @@ public void setPreview(boolean preview) {
   }
 
   public boolean isUdfSupport() {
-    return udfSupport;
+    return udfSupport != null && udfSupport;
   }
 
   public void setUdfSupport(boolean udfSupport) {
@@ -159,7 +165,7 @@ public void setUdfSupport(boolean udfSupport) {
   }
 
   public boolean isFlexTemplate() {
-    return flexTemplate;
+    return flexTemplate != null && flexTemplate;
   }
 
   public void setFlexTemplate(boolean flexTemplate) {
@@ -175,7 +181,7 @@ public void setSourceFilePath(String sourceFilePath) {
   }
 
   public boolean isHidden() {
-    return hidden;
+    return hidden != null && hidden;
   }
 
   public void setHidden(boolean hidden) {
@@ -183,7 +189,7 @@ public void setHidden(boolean hidden) {
   }
 
   public boolean isStreaming() {
-    return streaming;
+    return streaming != null && streaming;
   }
 
   public void setStreaming(boolean streaming) {
@@ -191,7 +197,7 @@ public void setStreaming(boolean streaming) {
   }
 
   public boolean isSupportsAtLeastOnce() {
-    return supportsAtLeastOnce;
+    return supportsAtLeastOnce != null && supportsAtLeastOnce;
   }
 
   public void setSupportsAtLeastOnce(boolean supportsAtLeastOnce) {
@@ -199,7 +205,7 @@ public void setSupportsAtLeastOnce(boolean supportsAtLeastOnce) {
   }
 
   public boolean isSupportsExactlyOnce() {
-    return supportsExactlyOnce;
+    return supportsExactlyOnce != null && supportsExactlyOnce;
   }
 
   public void setSupportsExactlyOnce(boolean supportsExactlyOnce) {