Add Cloud SQL read to Spec11Pipeline (#1173)

* Add Cloud SQL read to Spec11Pipeline

* Add database option

* Add database parameter

* Add a test of the full pipeline

* Use DatabaseHelper in tests

* restore the original tm

* More test fixes

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

@@ -23,8 +23,10 @@ import dagger.Component;
 import dagger.Module;
 import dagger.Provides;
 import google.registry.beam.common.RegistryJpaIO;
+import google.registry.beam.common.RegistryJpaIO.Read;
 import google.registry.beam.spec11.SafeBrowsingTransforms.EvaluateSafeBrowsingFn;
 import google.registry.config.RegistryConfig.ConfigModule;
+import google.registry.model.domain.DomainBase;
 import google.registry.model.reporting.Spec11ThreatMatch;
 import google.registry.model.reporting.Spec11ThreatMatch.ThreatType;
 import google.registry.util.Retrier;
@@ -97,20 +99,9 @@ public class Spec11Pipeline implements Serializable {
 
   void setupPipeline(Pipeline pipeline) {
     PCollection<Subdomain> domains =
-        pipeline.apply(
-            "Read active domains from BigQuery",
-            BigQueryIO.read(Subdomain::parseFromRecord)
-                .fromQuery(
-                    SqlTemplate.create(getQueryFromFile(Spec11Pipeline.class, "subdomains.sql"))
-                        .put("PROJECT_ID", options.getProject())
-                        .put("DATASTORE_EXPORT_DATASET", "latest_datastore_export")
-                        .put("REGISTRAR_TABLE", "Registrar")
-                        .put("DOMAIN_BASE_TABLE", "DomainBase")
-                        .build())
-                .withCoder(SerializableCoder.of(Subdomain.class))
-                .usingStandardSql()
-                .withoutValidation()
-                .withTemplateCompatibility());
+        options.getDatabase().equals("DATASTORE")
+            ? readFromBigQuery(options, pipeline)
+            : readFromCloudSql(pipeline);
 
     PCollection<KV<Subdomain, ThreatMatch>> threatMatches =
         domains.apply("Run through SafeBrowsing API", ParDo.of(safeBrowsingFn));
@@ -119,6 +110,47 @@ public class Spec11Pipeline implements Serializable {
     saveToGcs(threatMatches, options);
   }
 
+  static PCollection<Subdomain> readFromCloudSql(Pipeline pipeline) {
+    Read<Object[], Subdomain> read =
+        RegistryJpaIO.read(
+            "select d, r.emailAddress from Domain d join Registrar r on"
+                + " d.currentSponsorClientId = r.clientIdentifier where r.type = 'REAL'"
+                + " and d.deletionTime > now()",
+            Spec11Pipeline::parseRow);
+
+    return pipeline.apply("Read active domains from Cloud SQL", read);
+  }
+
+  static PCollection<Subdomain> readFromBigQuery(Spec11PipelineOptions options, Pipeline pipeline) {
+    return pipeline.apply(
+        "Read active domains from BigQuery",
+        BigQueryIO.read(Subdomain::parseFromRecord)
+            .fromQuery(
+                SqlTemplate.create(getQueryFromFile(Spec11Pipeline.class, "subdomains.sql"))
+                    .put("PROJECT_ID", options.getProject())
+                    .put("DATASTORE_EXPORT_DATASET", "latest_datastore_export")
+                    .put("REGISTRAR_TABLE", "Registrar")
+                    .put("DOMAIN_BASE_TABLE", "DomainBase")
+                    .build())
+            .withCoder(SerializableCoder.of(Subdomain.class))
+            .usingStandardSql()
+            .withoutValidation()
+            .withTemplateCompatibility());
+  }
+
+  private static Subdomain parseRow(Object[] row) {
+    DomainBase domainBase = (DomainBase) row[0];
+    String emailAddress = (String) row[1];
+    if (emailAddress == null) {
+      emailAddress = "";
+    }
+    return Subdomain.create(
+        domainBase.getDomainName(),
+        domainBase.getRepoId(),
+        domainBase.getCurrentSponsorClientId(),
+        emailAddress);
+  }
+
   static void saveToSql(
       PCollection<KV<Subdomain, ThreatMatch>> threatMatches, Spec11PipelineOptions options) {
     String transformId = "Spec11 Threat Matches";

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

@@ -34,4 +34,9 @@ public interface Spec11PipelineOptions extends RegistryPipelineOptions {
   String getReportingBucketUrl();
 
   void setReportingBucketUrl(String value);
+
+  @Description("The database to read data from.")
+  String getDatabase();
+
+  void setDatabase(String value);
 }

core/src/main/java/google/registry/reporting/ReportingModule.java

@@ -14,6 +14,7 @@
 
 package google.registry.reporting;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.RequestParameters.extractOptionalParameter;
 import static google.registry.request.RequestParameters.extractRequiredParameter;
 
@@ -55,6 +56,9 @@ public class ReportingModule {
   /** The request parameter specifying the jobId for a running Dataflow pipeline. */
   public static final String PARAM_JOB_ID = "jobId";
 
+  /** The request parameter for specifying which database reporting actions should read from. */
+  public static final String DATABASE = "database";
+
   /** Provides the Cloud Dataflow jobId for a pipeline. */
   @Provides
   @Parameter(PARAM_JOB_ID)
@@ -62,6 +66,14 @@ public class ReportingModule {
     return extractRequiredParameter(req, PARAM_JOB_ID);
   }
 
+  /** Provides the database for the pipeline to read from. */
+  @Provides
+  @Parameter(DATABASE)
+  static String provideDatabase(HttpServletRequest req) {
+    Optional<String> optionalDatabase = extractOptionalParameter(req, DATABASE);
+    return optionalDatabase.orElse(tm().isOfy() ? "DATASTORE" : "CLOUD_SQL");
+  }
+
   /** Extracts an optional YearMonth in yyyy-MM format from the request. */
   @Provides
   @Parameter(PARAM_YEAR_MONTH)

core/src/main/java/google/registry/reporting/spec11/GenerateSpec11ReportAction.java

@@ -15,6 +15,8 @@
 package google.registry.reporting.spec11;
 
 import static google.registry.beam.BeamUtils.createJobName;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+import static google.registry.reporting.ReportingModule.DATABASE;
 import static google.registry.reporting.ReportingModule.PARAM_DATE;
 import static google.registry.reporting.ReportingUtils.enqueueBeamReportingTask;
 import static google.registry.request.Action.Method.POST;
@@ -69,6 +71,7 @@ public class GenerateSpec11ReportAction implements Runnable {
   private final Clock clock;
   private final Response response;
   private final Dataflow dataflow;
+  private final String database;
 
   @Inject
   GenerateSpec11ReportAction(
@@ -78,15 +81,20 @@ public class GenerateSpec11ReportAction implements Runnable {
       @Config("reportingBucketUrl") String reportingBucketUrl,
       @Key("safeBrowsingAPIKey") String apiKey,
       @Parameter(PARAM_DATE) LocalDate date,
+      @Parameter(DATABASE) String database,
       Clock clock,
       Response response,
       Dataflow dataflow) {
     this.projectId = projectId;
     this.jobRegion = jobRegion;
     this.stagingBucketUrl = stagingBucketUrl;
+    if (tm().isOfy() && database.equals("CLOUD_SQL")) {
+      reportingBucketUrl = reportingBucketUrl.concat("-sql");
+    }
     this.reportingBucketUrl = reportingBucketUrl;
     this.apiKey = apiKey;
     this.date = date;
+    this.database = database;
     this.clock = clock;
     this.response = response;
     this.dataflow = dataflow;
@@ -105,6 +113,8 @@ public class GenerateSpec11ReportAction implements Runnable {
                   ImmutableMap.of(
                       "safeBrowsingApiKey",
                       apiKey,
+                      "database",
+                      database,
                       ReportingModule.PARAM_DATE,
                       date.toString(),
                       "reportingBucketUrl",

core/src/main/resources/google/registry/beam/spec11_pipeline_metadata.json

@@ -61,6 +61,14 @@
       "regexes": [
         "^gs:\\/\\/[^\\n\\r]+$"
       ]
+    },
+    {
+      "name": "database",
+      "label": "Database to read from.",
+      "helpText": "DATASTORE or CLOUD_SQL.",
+      "regexes": [
+        "^DATASTORE|CLOUD_SQL$"
+      ]
     }
   ]
 }