Migrate invoicing pipeline to read from Cloud SQL (#1220)

* Save entities to Cloud SQL for tests

* Fix merge conflict

* Filter out non-real registrars and non-invoicing TLDs

* Add 1 month filter

* Handle cancellations

* Add to pipeline

* Use database in pipeline

* fix formatting

* Add a full pipeline test

* Fix repo ids in tests

* Move query to separate file

* Remove unused variables

* Remove unnecessary debugging remnant

* Reformat sql file

* Add jpql issue description

* Use DateTimeUtils

* Fix license header year

* Fix SQL formatting

* Use regex pattern

* Fix string building

* Add test for makeCloudSqlQuery

* Add clarifying comment

core/src/main/java/google/registry/beam/invoicing/InvoicingPipeline.java

@@ -14,18 +14,29 @@
 
 package google.registry.beam.invoicing;
 
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.beam.BeamUtils.getQueryFromFile;
 import static org.apache.beam.sdk.values.TypeDescriptors.strings;
 
+import google.registry.beam.common.RegistryJpaIO;
+import google.registry.beam.common.RegistryJpaIO.Read;
 import google.registry.beam.invoicing.BillingEvent.InvoiceGroupingKey;
 import google.registry.beam.invoicing.BillingEvent.InvoiceGroupingKey.InvoiceGroupingKeyCoder;
+import google.registry.model.billing.BillingEvent.Flag;
+import google.registry.model.registrar.Registrar;
+import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
 import google.registry.reporting.billing.BillingModule;
+import google.registry.util.DateTimeUtils;
+import google.registry.util.DomainNameUtils;
 import google.registry.util.SqlTemplate;
 import java.io.Serializable;
 import java.time.LocalDateTime;
 import java.time.LocalTime;
 import java.time.YearMonth;
+import java.time.ZoneId;
 import java.time.format.DateTimeFormatter;
+import java.util.Optional;
+import java.util.regex.Pattern;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.PipelineResult;
 import org.apache.beam.sdk.coders.SerializableCoder;
@@ -58,6 +69,9 @@ public class InvoicingPipeline implements Serializable {
   private static final DateTimeFormatter TIMESTAMP_FORMATTER =
       DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSSSSS");
 
+  private static final Pattern SQL_COMMENT_REGEX =
+      Pattern.compile("^\\s*--.*\\n", Pattern.MULTILINE);
+
   private final InvoicingPipelineOptions options;
 
   InvoicingPipeline(InvoicingPipelineOptions options) {
@@ -71,21 +85,60 @@ public class InvoicingPipeline implements Serializable {
   }
 
   void setupPipeline(Pipeline pipeline) {
+    options.setIsolationOverride(TransactionIsolationLevel.TRANSACTION_READ_COMMITTED);
     PCollection<BillingEvent> billingEvents =
-        pipeline.apply(
-            "Read BillingEvents from Bigquery",
-            BigQueryIO.read(BillingEvent::parseFromRecord)
-                .fromQuery(makeQuery(options.getYearMonth(), options.getProject()))
-                .withCoder(SerializableCoder.of(BillingEvent.class))
-                .usingStandardSql()
-                .withoutValidation()
-                .withTemplateCompatibility());
+        options.getDatabase().equals("DATASTORE")
+            ? readFromBigQuery(options, pipeline)
+            : readFromCloudSql(options, pipeline);
 
     saveInvoiceCsv(billingEvents, options);
 
     saveDetailedCsv(billingEvents, options);
   }
 
+  static PCollection<BillingEvent> readFromBigQuery(
+      InvoicingPipelineOptions options, Pipeline pipeline) {
+    return pipeline.apply(
+        "Read BillingEvents from Bigquery",
+        BigQueryIO.read(BillingEvent::parseFromRecord)
+            .fromQuery(makeQuery(options.getYearMonth(), options.getProject()))
+            .withCoder(SerializableCoder.of(BillingEvent.class))
+            .usingStandardSql()
+            .withoutValidation()
+            .withTemplateCompatibility());
+  }
+
+  static PCollection<BillingEvent> readFromCloudSql(
+      InvoicingPipelineOptions options, Pipeline pipeline) {
+    Read<Object[], BillingEvent> read =
+        RegistryJpaIO.read(
+            makeCloudSqlQuery(options.getYearMonth()), false, InvoicingPipeline::parseRow);
+
+    return pipeline.apply("Read BillingEvents from Cloud SQL", read);
+  }
+
+  private static BillingEvent parseRow(Object[] row) {
+    google.registry.model.billing.BillingEvent.OneTime oneTime =
+        (google.registry.model.billing.BillingEvent.OneTime) row[0];
+    Registrar registrar = (Registrar) row[1];
+    return BillingEvent.create(
+        oneTime.getId(),
+        DateTimeUtils.toZonedDateTime(oneTime.getBillingTime(), ZoneId.of("UTC")),
+        DateTimeUtils.toZonedDateTime(oneTime.getEventTime(), ZoneId.of("UTC")),
+        registrar.getClientId(),
+        registrar.getBillingIdentifier().toString(),
+        registrar.getPoNumber().orElse(""),
+        DomainNameUtils.getTldFromDomainName(oneTime.getTargetId()),
+        oneTime.getReason().toString(),
+        oneTime.getTargetId(),
+        oneTime.getDomainRepoId(),
+        Optional.ofNullable(oneTime.getPeriodYears()).orElse(0),
+        oneTime.getCost().getCurrencyUnit().toString(),
+        oneTime.getCost().getAmount().doubleValue(),
+        String.join(
+            " ", oneTime.getFlags().stream().map(Flag::toString).collect(toImmutableSet())));
+  }
+
   /** Transform that converts a {@code BillingEvent} into an invoice CSV row. */
   private static class GenerateInvoiceRows
       extends PTransform<PCollection<BillingEvent>, PCollection<String>> {
@@ -171,6 +224,21 @@ public class InvoicingPipeline implements Serializable {
         .build();
   }
 
+  /** Create the Cloud SQL query for a given yearMonth at runtime. */
+  static String makeCloudSqlQuery(String yearMonth) {
+    YearMonth endMonth = YearMonth.parse(yearMonth).plusMonths(1);
+    String queryWithComments =
+        SqlTemplate.create(
+                getQueryFromFile(InvoicingPipeline.class, "cloud_sql_billing_events.sql"))
+            .put("FIRST_TIMESTAMP_OF_MONTH", yearMonth.concat("-01"))
+            .put(
+                "LAST_TIMESTAMP_OF_MONTH",
+                String.format("%d-%d-01", endMonth.getYear(), endMonth.getMonthValue()))
+            .build();
+    // Remove the comments from the query string
+    return SQL_COMMENT_REGEX.matcher(queryWithComments).replaceAll("");
+  }
+
   public static void main(String[] args) {
     PipelineOptionsFactory.register(InvoicingPipelineOptions.class);
     InvoicingPipelineOptions options =

core/src/main/java/google/registry/beam/invoicing/InvoicingPipelineOptions.java

@@ -30,6 +30,11 @@ public interface InvoicingPipelineOptions extends RegistryPipelineOptions {
 
   void setInvoiceFilePrefix(String value);
 
+  @Description("The database to read data from.")
+  String getDatabase();
+
+  void setDatabase(String value);
+
   @Description("The GCS bucket URL for invoices and detailed  reports to be uploaded.")
   String getBillingBucketUrl();
 

core/src/main/java/google/registry/reporting/billing/GenerateInvoicesAction.java

@@ -15,6 +15,8 @@
 package google.registry.reporting.billing;
 
 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.ReportingUtils.enqueueBeamReportingTask;
 import static google.registry.reporting.billing.BillingModule.PARAM_SHOULD_PUBLISH;
 import static google.registry.request.Action.Method.POST;
@@ -70,6 +72,7 @@ public class GenerateInvoicesAction implements Runnable {
   private final Clock clock;
   private final Response response;
   private final Dataflow dataflow;
+  private final String database;
 
   @Inject
   GenerateInvoicesAction(
@@ -79,6 +82,7 @@ public class GenerateInvoicesAction implements Runnable {
       @Config("billingBucketUrl") String billingBucketUrl,
       @Config("invoiceFilePrefix") String invoiceFilePrefix,
       @Parameter(PARAM_SHOULD_PUBLISH) boolean shouldPublish,
+      @Parameter(DATABASE) String database,
       YearMonth yearMonth,
       BillingEmailUtils emailUtils,
       Clock clock,
@@ -87,9 +91,15 @@ public class GenerateInvoicesAction implements Runnable {
     this.projectId = projectId;
     this.jobRegion = jobRegion;
     this.stagingBucketUrl = stagingBucketUrl;
+    // When generating the invoices using Cloud SQL before database cutover, save the reports in a
+    // separate bucket so that it does not overwrite the Datastore invoices.
+    if (tm().isOfy() && database.equals("CLOUD_SQL")) {
+      billingBucketUrl = billingBucketUrl.concat("-sql");
+    }
     this.billingBucketUrl = billingBucketUrl;
     this.invoiceFilePrefix = invoiceFilePrefix;
     this.shouldPublish = shouldPublish;
+    this.database = database;
     this.yearMonth = yearMonth;
     this.emailUtils = emailUtils;
     this.clock = clock;
@@ -113,6 +123,8 @@ public class GenerateInvoicesAction implements Runnable {
                       yearMonth.toString("yyyy-MM"),
                       "invoiceFilePrefix",
                       invoiceFilePrefix,
+                      "database",
+                      database,
                       "billingBucketUrl",
                       billingBucketUrl));
       LaunchFlexTemplateResponse launchResponse =

core/src/main/resources/google/registry/beam/invoicing/sql/cloud_sql_billing_events.sql

@@ -0,0 +1,39 @@
+-- Copyright 2021 The Nomulus Authors. All Rights Reserved.
+--
+-- Licensed under the Apache License, Version 2.0 (the "License");
+-- you may not use this file except in compliance with the License.
+-- You may obtain a copy of the License at
+--
+--     http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+
+-- This query gathers all non-canceled billing events for a given
+-- YEAR_MONTH in yyyy-MM format.
+
+-- This query differs from the one for BigQuery because JPQL does not have all
+-- the same functionality as a native query. Since there is no JPQL equivalent
+-- for REGEXP_EXTRACT the tld cannot be extracted from the targetId in the
+-- BillingEvent table. This requires an additional join with the domain table in
+-- order to find the value for the TLD associated with the billing event. Also,
+-- JPQL does not allow subqueries in the SELECT and FROM clauses. This prevents
+-- the query from being able to filter out certain data early in the query to
+-- reduce the size of intermediate results. It may be useful to measure this
+-- query's performance and consider switching to using a native query.
+
+SELECT b, r FROM BillingEvent b
+JOIN Registrar r ON b.clientId = r.clientIdentifier
+JOIN Domain d ON b.domainRepoId = d.repoId
+JOIN Tld t ON t.tldStrId = d.tld
+LEFT JOIN BillingCancellation c ON b.id = c.refOneTime.billingId
+LEFT JOIN BillingCancellation cr ON b.cancellationMatchingBillingEvent = cr.refRecurring.billingId
+WHERE r.billingIdentifier IS NOT NULL
+AND r.type = 'REAL'
+AND t.invoicingEnabled IS TRUE
+AND b.billingTime BETWEEN CAST('%FIRST_TIMESTAMP_OF_MONTH%' AS timestamp) AND CAST('%LAST_TIMESTAMP_OF_MONTH%' AS timestamp)
+AND c.id IS NULL
+AND cr.id IS NULL

core/src/main/resources/google/registry/beam/invoicing_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$"
+      ]
     }
   ]
 }