Chain applies together

google · Oct 14, 2022 · 458040e · 458040e
1 parent 29557d0
commit 458040e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 16 deletions.
diff --git a/core/src/main/java/google/registry/beam/spec11/Spec11Pipeline.java b/core/src/main/java/google/registry/beam/spec11/Spec11Pipeline.java
@@ -157,9 +157,10 @@ private static KV<String, String> parseRow(Object[] row) {
   static void saveToSql(
       PCollection<KV<DomainNameInfo, ThreatMatch>> threatMatches, Spec11PipelineOptions options) {
     LocalDate date = LocalDate.parse(options.getDate(), ISODateTimeFormat.date());
+    String transformId = "Spec11 Threat Matches";
 
-    PCollection<Spec11ThreatMatch> spec11ThreatMatches =
-        threatMatches.apply(
+    threatMatches
+        .apply(
             "Construct objects",
             ParDo.of(
                 new DoFn<KV<DomainNameInfo, ThreatMatch>, Spec11ThreatMatch>() {
@@ -179,17 +180,14 @@ public void processElement(
                             .build();
                     output.output(spec11ThreatMatch);
                   }
-                }));
-
-    spec11ThreatMatches.apply("Prevent Fusing", Reshuffle.viaRandomKey());
-    String transformId = "Spec11 Threat Matches";
-
-    spec11ThreatMatches.apply(
-        "Write to Sql: " + transformId,
-        RegistryJpaIO.<Spec11ThreatMatch>write()
-            .withName(transformId)
-            .withBatchSize(options.getSqlWriteBatchSize())
-            .withShards(options.getSqlWriteShards()));
+                }))
+        .apply("Prevent Fusing", Reshuffle.viaRandomKey())
+        .apply(
+            "Write to Sql: " + transformId,
+            RegistryJpaIO.<Spec11ThreatMatch>write()
+                .withName(transformId)
+                .withBatchSize(options.getSqlWriteBatchSize())
+                .withShards(options.getSqlWriteShards()));
   }
 
   static void saveToGcs(

diff --git a/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java b/core/src/test/java/google/registry/beam/spec11/Spec11PipelineTest.java
@@ -288,10 +288,10 @@ private void verifySaveToCloudSql() {
     jpaTm()
         .transact(
             () -> {
-              ImmutableList<Spec11ThreatMatch> sqlThreatMatches =
+              ImmutableList<Spec11ThreatMatch> spec11ThreatMatches =
                   Spec11ThreatMatchDao.loadEntriesByDate(jpaTm(), new LocalDate(2020, 1, 27));
-              assertThat(sqlThreatMatches)
-                  .comparingElementsUsing(immutableObjectCorrespondence())
+              assertThat(spec11ThreatMatches)
+                  .comparingElementsUsing(immutableObjectCorrespondence("id"))
                   .containsExactlyElementsIn(sqlThreatMatches);
             });
   }