Persist ThreatMatches into Spec11ThreatMatch (#723)

* Replace jpaTm with a JpaSupplierFactory

* Style

* Style

* Pipeline takes in a SerializableSupplier instead

* Change the ordering of imports

* Test a good domain in addition to a bad one

* Rename and check good domain for Transact Answer

* Use standard Mockito verify

* Verify transact call and no more interactions

* Remove Answer comment

* Naming chsnges

* Deploy Spec 11 pipeline correctly

* Fix formatting of deploy file

* Use a file to persist state across Cloud Build steps

Co-authored-by: Gus Brodman <gbrodman@google.com>

core/src/main/java/google/registry/beam/spec11/Spec11Pipeline.java

@@ -19,9 +19,15 @@ import static google.registry.beam.BeamUtils.getQueryFromFile;
 
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableSet;
+import google.registry.backup.AppEngineEnvironment;
+import google.registry.beam.initsql.Transforms.SerializableSupplier;
 import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn;
 import google.registry.config.CredentialModule.LocalCredential;
 import google.registry.config.RegistryConfig.Config;
+import google.registry.model.reporting.Spec11ThreatMatch;
+import google.registry.model.reporting.Spec11ThreatMatch.ThreatType;
+import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.util.GoogleCredentialsBundle;
 import google.registry.util.Retrier;
 import google.registry.util.SqlTemplate;
@@ -37,6 +43,7 @@ import org.apache.beam.sdk.options.Description;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.options.ValueProvider;
 import org.apache.beam.sdk.options.ValueProvider.NestedValueProvider;
+import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.GroupByKey;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.ParDo;
@@ -46,6 +53,7 @@ import org.apache.beam.sdk.values.TypeDescriptor;
 import org.apache.beam.sdk.values.TypeDescriptors;
 import org.joda.time.LocalDate;
 import org.joda.time.YearMonth;
+import org.joda.time.format.ISODateTimeFormat;
 import org.json.JSONArray;
 import org.json.JSONException;
 import org.json.JSONObject;
@@ -86,6 +94,7 @@ public class Spec11Pipeline implements Serializable {
   private final String reportingBucketUrl;
   private final GoogleCredentials googleCredentials;
   private final Retrier retrier;
+  private final SerializableSupplier<JpaTransactionManager> jpaSupplierFactory;
 
   @Inject
   public Spec11Pipeline(
@@ -93,12 +102,14 @@ public class Spec11Pipeline implements Serializable {
       @Config("beamStagingUrl") String beamStagingUrl,
       @Config("spec11TemplateUrl") String spec11TemplateUrl,
       @Config("reportingBucketUrl") String reportingBucketUrl,
+      SerializableSupplier<JpaTransactionManager> jpaSupplierFactory,
       @LocalCredential GoogleCredentialsBundle googleCredentialsBundle,
       Retrier retrier) {
     this.projectId = projectId;
     this.beamStagingUrl = beamStagingUrl;
     this.spec11TemplateUrl = spec11TemplateUrl;
     this.reportingBucketUrl = reportingBucketUrl;
+    this.jpaSupplierFactory = jpaSupplierFactory;
     this.googleCredentials = googleCredentialsBundle.getGoogleCredentials();
     this.retrier = retrier;
   }
@@ -177,12 +188,40 @@ public class Spec11Pipeline implements Serializable {
       EvaluateSafeBrowsingFn evaluateSafeBrowsingFn,
       ValueProvider<String> dateProvider) {
 
+    PCollection<KV<Subdomain, ThreatMatch>> subdomainsSql =
+        domains.apply("Run through SafeBrowsing API", ParDo.of(evaluateSafeBrowsingFn));
+    /* Store ThreatMatch objects in SQL. */
+    subdomainsSql.apply(
+        ParDo.of(
+            new DoFn<KV<Subdomain, ThreatMatch>, Void>() {
+              @ProcessElement
+              public void processElement(ProcessContext context) {
+                // create the Spec11ThreatMatch from Subdomain and ThreatMatch
+                try (AppEngineEnvironment env = new AppEngineEnvironment()) {
+                  Subdomain subdomain = context.element().getKey();
+                  Spec11ThreatMatch threatMatch =
+                      new Spec11ThreatMatch.Builder()
+                          .setThreatTypes(
+                              ImmutableSet.of(
+                                  ThreatType.valueOf(context.element().getValue().threatType())))
+                          .setCheckDate(
+                              LocalDate.parse(dateProvider.get(), ISODateTimeFormat.date()))
+                          .setDomainName(subdomain.domainName())
+                          .setDomainRepoId(subdomain.domainRepoId())
+                          .setRegistrarId(subdomain.registrarId())
+                          .build();
+                  JpaTransactionManager jpaTransactionManager = jpaSupplierFactory.get();
+                  jpaTransactionManager.transact(() -> jpaTransactionManager.saveNew(threatMatch));
+                }
+              }
+            }));
+
     /* Store ThreatMatch objects in JSON. */
     PCollection<KV<Subdomain, ThreatMatch>> subdomainsJson =
         domains.apply("Run through SafeBrowsingAPI", ParDo.of(evaluateSafeBrowsingFn));
     subdomainsJson
         .apply(
-            "Map registrar client ID to email/ThreatMatch pair",
+            "Map registrar ID to email/ThreatMatch pair",
             MapElements.into(
                     TypeDescriptors.kvs(
                         TypeDescriptors.strings(), TypeDescriptor.of(EmailAndThreatMatch.class)))

core/src/main/java/google/registry/model/reporting/Spec11ThreatMatch.java

@@ -140,6 +140,16 @@ public class Spec11ThreatMatch extends ImmutableObject implements Buildable, Sql
       return super.build();
     }
 
+    /**
+     * Manually set the ID for testing or other special circumstances.
+     *
+     * <p>In general the ID is generated by SQL and there should be no need to set it manually.
+     */
+    public Builder setId(Long id) {
+      getInstance().id = id;
+      return this;
+    }
+
     public Builder setDomainName(String domainName) {
       getInstance().domainName = domainName;
       getInstance().tld = DomainNameUtils.getTldFromDomainName(domainName);

core/src/main/java/google/registry/tools/DeploySpec11PipelineCommand.java

@@ -14,7 +14,10 @@
 
 package google.registry.tools;
 
+import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
+import google.registry.beam.initsql.BeamJpaModule.JpaTransactionManagerComponent;
+import google.registry.beam.initsql.JpaSupplierFactory;
 import google.registry.beam.spec11.Spec11Pipeline;
 import google.registry.config.CredentialModule.LocalCredential;
 import google.registry.config.RegistryConfig.Config;
@@ -31,6 +34,12 @@ public class DeploySpec11PipelineCommand implements Command {
   @Config("projectId")
   String projectId;
 
+  @Parameter(
+      names = {"-p", "--project"},
+      description = "Cloud KMS project ID",
+      required = true)
+  String cloudKmsProjectId;
+
   @Inject
   @Config("beamStagingUrl")
   String beamStagingUrl;
@@ -53,12 +62,19 @@ public class DeploySpec11PipelineCommand implements Command {
 
   @Override
   public void run() {
+    JpaSupplierFactory jpaSupplierFactory =
+        new JpaSupplierFactory(
+            sqlAccessInfoFile,
+            cloudKmsProjectId,
+            JpaTransactionManagerComponent::cloudSqlJpaTransactionManager);
+
     Spec11Pipeline pipeline =
         new Spec11Pipeline(
             projectId,
             beamStagingUrl,
             spec11TemplateUrl,
             reportingBucketUrl,
+            jpaSupplierFactory,
             googleCredentialsBundle,
             retrier);
     pipeline.deploy();