Simplify dependency manager: remove unused code

GoogleCloudPlatform · Dec 16, 2024 · 4b0d74c · 4b0d74c
1 parent cd28d01
commit 4b0d74c
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 66 deletions.
diff --git a/...-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/templates/GoogleCloudToNeo4j.java b/...-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/templates/GoogleCloudToNeo4j.java
@@ -292,8 +292,7 @@ public void run() {
               String.format("Metadata for source %s", source.getName()), provider.queryMetadata());
       sourceRows.add(sourceMetadata);
       Schema sourceBeamSchema = sourceMetadata.getSchema();
-      processingQueue.addToQueue(
-          ArtifactType.source, false, source.getName(), defaultActionContext, sourceMetadata);
+      processingQueue.addToQueue(ArtifactType.source, source.getName(), defaultActionContext);
       PCollection<Row> nullableSourceBeamRows = null;
 
       ////////////////////////////
@@ -338,7 +337,7 @@ public void run() {
         List<PCollection<?>> dependencies =
             new ArrayList<>(preActionRows.getOrDefault(ActionStage.PRE_NODES, List.of()));
         dependencies.add(
-            processingQueue.waitOnCollections(target.getDependencies(), nodeStepDescription));
+            processingQueue.resolveOutputs(target.getDependencies(), nodeStepDescription));
 
         PCollection<Row> blockingReturn =
             preInsertBeamRows
@@ -364,8 +363,7 @@ public void run() {
             .computeIfAbsent(TargetType.NODE, (type) -> new ArrayList<>(nodeTargets.size()))
             .add(blockingReturn);
 
-        processingQueue.addToQueue(
-            ArtifactType.node, false, target.getName(), blockingReturn, preInsertBeamRows);
+        processingQueue.addToQueue(ArtifactType.node, target.getName(), blockingReturn);
       }
 
       ////////////////////////////
@@ -406,7 +404,7 @@ public void run() {
         dependencyNames.add(target.getStartNodeReference());
         dependencyNames.add(target.getEndNodeReference());
         dependencies.add(
-            processingQueue.waitOnCollections(dependencyNames, relationshipStepDescription));
+            processingQueue.resolveOutputs(dependencyNames, relationshipStepDescription));
 
         PCollection<Row> blockingReturn =
             preInsertBeamRows
@@ -433,8 +431,7 @@ public void run() {
                 TargetType.RELATIONSHIP, (type) -> new ArrayList<>(relationshipTargets.size()))
             .add(blockingReturn);
         // serialize relationships
-        processingQueue.addToQueue(
-            ArtifactType.edge, false, target.getName(), blockingReturn, preInsertBeamRows);
+        processingQueue.addToQueue(ArtifactType.edge, target.getName(), blockingReturn);
       }
       ////////////////////////////
       // Custom query targets
@@ -452,8 +449,7 @@ public void run() {
         List<PCollection<?>> dependencies =
             new ArrayList<>(preActionRows.getOrDefault(ActionStage.PRE_QUERIES, List.of()));
         dependencies.add(
-            processingQueue.waitOnCollections(
-                target.getDependencies(), customQueryStepDescription));
+            processingQueue.resolveOutputs(target.getDependencies(), customQueryStepDescription));
 
         PCollection<Row> blockingReturn =
             nullableSourceBeamRows
@@ -478,12 +474,7 @@ public void run() {
         targetRows
             .computeIfAbsent(TargetType.QUERY, (type) -> new ArrayList<>(customQueryTargets.size()))
             .add(blockingReturn);
-        processingQueue.addToQueue(
-            ArtifactType.custom_query,
-            false,
-            target.getName(),
-            blockingReturn,
-            nullableSourceBeamRows);
+        processingQueue.addToQueue(ArtifactType.custom_query, target.getName(), blockingReturn);
       }
     }
 

diff --git a/...ooglecloud-to-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/utils/BeamBlock.java b/...ooglecloud-to-neo4j/src/main/java/com/google/cloud/teleport/v2/neo4j/utils/BeamBlock.java
@@ -32,59 +32,20 @@
 public class BeamBlock {
 
   private static final Logger LOG = LoggerFactory.getLogger(BeamBlock.class);
-  private final List<PCollection<Row>> sourceQueue = new ArrayList<>();
-  private final List<PCollection<Row>> preloadActionQueue = new ArrayList<>();
-  private final List<PCollection<Row>> processActionQueue = new ArrayList<>();
-  private final List<PCollection<Row>> nodeQueue = new ArrayList<>();
-  private final List<PCollection<Row>> edgeQueue = new ArrayList<>();
-  private final List<PCollection<Row>> customQueue = new ArrayList<>();
-  private final Map<String, PCollection<Row>> executeAfterNamedQueue = new HashMap<>();
-  private final Map<String, PCollection<Row>> executionContexts = new HashMap<>();
-  private PCollection<Row> defaultCollection;
+  private final Map<String, PCollection<Row>> outputs = new HashMap<>();
+  private final PCollection<Row> defaultCollection;
 
   public BeamBlock(PCollection<Row> defaultCollection) {
     this.defaultCollection = defaultCollection;
   }
 
-  public void addToQueue(
-      ArtifactType artifactType, boolean preload, String name, PCollection<Row> blockingReturn) {
-    addToQueue(artifactType, preload, name, blockingReturn, defaultCollection);
+  public void addToQueue(ArtifactType artifactType, String name, PCollection<Row> output) {
+    outputs.put(artifactType.name() + ":" + name, output);
   }
 
-  public void addToQueue(
-      ArtifactType artifactType,
-      boolean preload,
-      String name,
-      PCollection<Row> blockingReturn,
-      PCollection<Row> executionContext) {
-    switch (artifactType) {
-      case action:
-        if (preload) {
-          preloadActionQueue.add(blockingReturn);
-        } else {
-          processActionQueue.add(blockingReturn);
-        }
-        break;
-      case source:
-        sourceQueue.add(blockingReturn);
-        break;
-      case node:
-        nodeQueue.add(blockingReturn);
-        break;
-      case edge:
-        edgeQueue.add(blockingReturn);
-        break;
-      case custom_query:
-        customQueue.add(blockingReturn);
-        break;
-    }
-    executeAfterNamedQueue.put(artifactType.name() + ":" + name, blockingReturn);
-    executionContexts.put(artifactType.name() + ":" + name, executionContext);
-  }
-
-  public PCollection<Row> waitOnCollections(
+  public PCollection<Row> resolveOutputs(
       Collection<String> dependencies, String queuingDescription) {
-    List<PCollection<Row>> waitOnQueues = populateQueueForTargets(dependencies);
+    List<PCollection<Row>> waitOnQueues = resolveOutputs(dependencies);
     if (waitOnQueues.isEmpty()) {
       waitOnQueues.add(defaultCollection);
     }
@@ -103,15 +64,16 @@ public PCollection<Row> waitOnCollections(
             Flatten.pCollections());
   }
 
-  private List<PCollection<Row>> populateQueueForTargets(Collection<String> dependencies) {
-    List<PCollection<Row>> waitOnQueues = new ArrayList<>();
+  private List<PCollection<Row>> resolveOutputs(Collection<String> dependencies) {
+    List<PCollection<Row>> outputs = new ArrayList<>();
     for (String dependency : dependencies) {
       for (ArtifactType type : ArtifactType.values()) {
-        if (executeAfterNamedQueue.containsKey(type + ":" + dependency)) {
-          waitOnQueues.add(executeAfterNamedQueue.get(type + ":" + dependency));
+        if (this.outputs.containsKey(type + ":" + dependency)) {
+          outputs.add(this.outputs.get(type + ":" + dependency));
+          break;
         }
       }
     }
-    return waitOnQueues;
+    return outputs;
   }
 }