Migrate DNS query table (#2543)

Co-authored-by: Lai Jiang <jianglai@google.com>

core/src/main/java/google/registry/config/files/default-config.yaml

@@ -84,7 +84,7 @@ registryPolicy:
 
   # Custom logic class for handling DNS query count reporting for ICANN.
   # See reporting/icann/DnsCountQueryCoordinator.java
-  dnsCountQueryCoordinatorClass: google.registry.reporting.icann.BasicDnsCountQueryCoordinator
+  dnsCountQueryCoordinatorClass: google.registry.reporting.icann.DummyDnsCountQueryCoordinator
 
   # Length of time after which contact transfers automatically conclude.
   contactAutomaticTransferDays: 5

core/src/main/java/google/registry/module/RequestComponent.java

@@ -81,7 +81,7 @@ import google.registry.reporting.billing.BillingModule;
 import google.registry.reporting.billing.CopyDetailReportsAction;
 import google.registry.reporting.billing.GenerateInvoicesAction;
 import google.registry.reporting.billing.PublishInvoicesAction;
-import google.registry.reporting.icann.DnsCountQueryCoordinatorModule;
+import google.registry.reporting.icann.DnsCountQueryCoordinator.DnsCountQueryCoordinatorModule;
 import google.registry.reporting.icann.IcannReportingModule;
 import google.registry.reporting.icann.IcannReportingStagingAction;
 import google.registry.reporting.icann.IcannReportingUploadAction;

core/src/main/java/google/registry/module/backend/BackendRequestComponent.java

@@ -59,7 +59,7 @@ import google.registry.reporting.billing.BillingModule;
 import google.registry.reporting.billing.CopyDetailReportsAction;
 import google.registry.reporting.billing.GenerateInvoicesAction;
 import google.registry.reporting.billing.PublishInvoicesAction;
-import google.registry.reporting.icann.DnsCountQueryCoordinatorModule;
+import google.registry.reporting.icann.DnsCountQueryCoordinator.DnsCountQueryCoordinatorModule;
 import google.registry.reporting.icann.IcannReportingModule;
 import google.registry.reporting.icann.IcannReportingStagingAction;
 import google.registry.reporting.icann.IcannReportingUploadAction;

core/src/main/java/google/registry/reporting/icann/ActivityReportingQueryBuilder.java

@@ -77,7 +77,7 @@ public final class ActivityReportingQueryBuilder implements QueryBuilder {
     queriesBuilder.put(
         getTableName(REGISTRAR_OPERATING_STATUS, yearMonth), operationalRegistrarsQuery);
 
-    String dnsCountsQuery = dnsCountQueryCoordinator.createQuery(yearMonth);
+    String dnsCountsQuery = dnsCountQueryCoordinator.createQuery();
     queriesBuilder.put(getTableName(DNS_COUNTS, yearMonth), dnsCountsQuery);
 
     // Convert reportingMonth into YYYYMMDD format for Bigquery table partition pattern-matching.

core/src/main/java/google/registry/reporting/icann/CloudDnsCountQueryCoordinator.java

@@ -0,0 +1,82 @@
+// Copyright 2024 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.
+package google.registry.reporting.icann;
+
+import com.google.common.flogger.FluentLogger;
+import com.google.common.io.Resources;
+import google.registry.bigquery.BigqueryUtils.TableType;
+import google.registry.util.ResourceUtils;
+import google.registry.util.SqlTemplate;
+import java.util.concurrent.ExecutionException;
+import org.joda.time.YearMonth;
+import org.joda.time.format.DateTimeFormat;
+
+/**
+ * DNS Count query that relies on a table Cloud DNS publishes internally to Google.
+ *
+ * <p>The internal Plx table is exposed as a BigQuery table via BQ-TS federation. This is not
+ * applicable to external users who also happen to use Cloud DNS as the plx table is specific to
+ * Google Registry's zones. External (non-Google) users must re-implement the abstract class and
+ * configure the usage of the new class using the `registryPolicy.dnsCountQueryCoordinatorClass`
+ * field in the config file.
+ */
+public class CloudDnsCountQueryCoordinator extends DnsCountQueryCoordinator {
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+  private static final String PLX_DNS_TABLE_NAME = "dns_counts_from_plx";
+  private static final String TABLE_ID = "zoneman_daily_query_counts";
+
+  @Override
+  public String createQuery() {
+    String template =
+        ResourceUtils.readResourceUtf8(
+            Resources.getResource(this.getClass(), "sql/dns_counts_cloud.sql"));
+    return SqlTemplate.create(template)
+        .put("PROJECT_ID", projectId)
+        .put("ICANN_REPORTING_DATA_SET", icannReportingDataSet)
+        .put("DNS_TABLE_NAME", PLX_DNS_TABLE_NAME)
+        .build();
+  }
+
+  @Override
+  public void prepareForQuery(YearMonth yearMonth) throws InterruptedException {
+    logger.atInfo().log("Generating intermediary table dns_counts");
+    String query = getPlxDnsTableQuery(yearMonth);
+    try {
+      bigquery
+          .startQuery(
+              query,
+              bigquery
+                  .buildDestinationTable(PLX_DNS_TABLE_NAME)
+                  .description("A table holding DNS query counts to generate ACTIVITY reports.")
+                  .type(TableType.TABLE)
+                  .build())
+          .get();
+    } catch (ExecutionException e) {
+      throw new RuntimeException("Error while running BigQuery query", e.getCause());
+    }
+  }
+
+  String getPlxDnsTableQuery(YearMonth yearMonth) {
+    String template =
+        ResourceUtils.readResourceUtf8(
+            Resources.getResource(this.getClass(), "sql/prepare_dns_counts_internal.sql"));
+    SqlTemplate queryTemplate =
+        SqlTemplate.create(template)
+            .put("PROJECT_ID", projectId)
+            .put("DATASET_ID", icannReportingDataSet)
+            .put("TABLE_ID", TABLE_ID)
+            .put("YEAR_MONTH", DateTimeFormat.forPattern("yyyyMM").print(yearMonth));
+    return queryTemplate.build();
+  }
+}

core/src/main/java/google/registry/reporting/icann/DnsCountQueryCoordinator.java

@@ -14,45 +14,43 @@
 
 package google.registry.reporting.icann;
 
+import static google.registry.reporting.icann.IcannReportingModule.ICANN_REPORTING_DATA_SET;
+import static google.registry.util.TypeUtils.getClassFromString;
+import static google.registry.util.TypeUtils.instantiate;
+
+import dagger.MembersInjector;
+import dagger.Module;
+import dagger.Provides;
 import google.registry.bigquery.BigqueryConnection;
+import google.registry.config.RegistryConfig.Config;
+import javax.inject.Inject;
+import javax.inject.Named;
 import org.joda.time.YearMonth;
 
 /**
  * Methods for preparing and querying DNS statistics.
  *
  * <p>DNS systems may have different ways of providing this information, so it's useful to
- * modularize this.
+ * modularize this, by providing defining the `registryPolicy.dnsCountQueryCoordinatorClass` in your
+ * config file.
  *
- * <p>Derived classes must provide a constructor that accepts a
- * {@link google.registry.reporting.icann.DnsCountQueryCoordinator.Params}.  To override this,
- * define dnsCountQueryCoordinatorClass in your config file.
+ * <p>Due to limitations of {@link MembersInjector}, any injectable field needs to be declared in
+ * the base class, even if it is only used in a derived class.
  */
-public interface DnsCountQueryCoordinator {
+public abstract class DnsCountQueryCoordinator {
 
-  /**
-   * Class to carry parameters for a new coordinator.
-   *
-   * <p>If your report query requires any additional parameters, add them here.
-   */
-  class Params {
+  @Inject BigqueryConnection bigquery;
 
-    public BigqueryConnection bigquery;
+  @Inject
+  @Config("projectId")
+  String projectId;
 
-    /** The Google Cloud project id. */
-    public String projectId;
-
-    /** The BigQuery dataset from which to query. */
-    public String icannReportingDataSet;
-
-    public Params(BigqueryConnection bigquery, String projectId, String icannReportingDataSet) {
-      this.bigquery = bigquery;
-      this.projectId = projectId;
-      this.icannReportingDataSet = icannReportingDataSet;
-    }
-  }
+  @Inject
+  @Named(ICANN_REPORTING_DATA_SET)
+  String icannReportingDataSet;
 
   /** Creates the string used to query bigtable for DNS count information. */
-  String createQuery(YearMonth yearMonth);
+  abstract String createQuery();
 
   /**
    * Do any necessary preparation for the DNS query.
@@ -61,5 +59,18 @@ public interface DnsCountQueryCoordinator {
    * interruptible futures to prepare the query (and the correct thing to do with such exceptions is
    * to handle them correctly or propagate them as-is, no {@link RuntimeException} wrapping).
    */
-  void prepareForQuery(YearMonth yearMonth) throws InterruptedException;
+  abstract void prepareForQuery(YearMonth yearMonth) throws InterruptedException;
+
+  @Module
+  public static class DnsCountQueryCoordinatorModule {
+    @Provides
+    static DnsCountQueryCoordinator provideDnsCountQueryCoordinator(
+        MembersInjector<DnsCountQueryCoordinator> injector,
+        @Config("dnsCountQueryCoordinatorClass") String customClass) {
+      DnsCountQueryCoordinator coordinator =
+          instantiate(getClassFromString(customClass, DnsCountQueryCoordinator.class));
+      injector.injectMembers(coordinator);
+      return coordinator;
+    }
+  }
 }

core/src/main/java/google/registry/reporting/icann/DnsCountQueryCoordinatorModule.java

@@ -1,41 +0,0 @@
-// Copyright 2018 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.
-
-package google.registry.reporting.icann;
-
-import static google.registry.reporting.icann.IcannReportingModule.ICANN_REPORTING_DATA_SET;
-import static google.registry.util.TypeUtils.getClassFromString;
-import static google.registry.util.TypeUtils.instantiate;
-
-import dagger.Module;
-import dagger.Provides;
-import google.registry.bigquery.BigqueryConnection;
-import google.registry.config.RegistryConfig.Config;
-import javax.inject.Named;
-
-/** Dagger module to provide the DnsCountQueryCoordinator. */
-@Module
-public class DnsCountQueryCoordinatorModule {
-
-  @Provides
-  static DnsCountQueryCoordinator provideDnsCountQueryCoordinator(
-      @Config("dnsCountQueryCoordinatorClass") String customClass,
-      BigqueryConnection bigquery,
-      @Config("projectId") String projectId,
-      @Named(ICANN_REPORTING_DATA_SET) String icannReportingDataSet) {
-    DnsCountQueryCoordinator.Params params =
-        new DnsCountQueryCoordinator.Params(bigquery, projectId, icannReportingDataSet);
-    return instantiate(getClassFromString(customClass, DnsCountQueryCoordinator.class), params);
-  }
-}

core/src/main/java/google/registry/reporting/icann/DummyDnsCountQueryCoordinator.java

@@ -18,16 +18,13 @@ import google.registry.util.ResourceUtils;
 import google.registry.util.SqlTemplate;
 import org.joda.time.YearMonth;
 
-/**
- * DNS Count query for the basic case.
- */
-public class BasicDnsCountQueryCoordinator implements DnsCountQueryCoordinator {
-
-  BasicDnsCountQueryCoordinator(DnsCountQueryCoordinator.Params params) {}
+/** DNS Count query where returned values are all -1. */
+public class DummyDnsCountQueryCoordinator extends DnsCountQueryCoordinator {
 
   @Override
-  public String createQuery(YearMonth yearMonth) {
-    return SqlTemplate.create(ResourceUtils.readResourceUtf8(this.getClass(), "sql/dns_counts.sql"))
+  public String createQuery() {
+    return SqlTemplate.create(
+            ResourceUtils.readResourceUtf8(this.getClass(), "sql/dns_counts_dummy.sql"))
         .build();
   }
 

core/src/main/resources/google/registry/reporting/icann/sql/dns_counts_cloud.sql

@@ -15,10 +15,5 @@
 
   -- Retrieve per-TLD DNS query counts.
 
-  -- This is a hack to enable using DNS counts from the internal-only #plx
-  -- workflow. See other references to b/67301320 in the codebase to see the
-  -- full extent of the hackery.
-  -- TODO(b/67301320): Delete this when we can make open-source DNS metrics.
-
 SELECT *
-FROM `domain-registry-alpha.icann_reporting.dns_counts_from_plx`
+FROM `%PROJECT_ID%.%ICANN_REPORTING_DATA_SET%.%DNS_TABLE_NAME%`

core/src/main/resources/google/registry/reporting/icann/sql/prepare_dns_counts_internal.sql

@@ -0,0 +1,31 @@
+#standardSQL
+  -- Copyright 2024 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.
+
+  -- Retrieve per-TLD DNS query counts.
+
+SELECT
+  tld,
+  CASE
+    WHEN transport = 'tcp' THEN 'dns-tcp-queries'
+    WHEN transport = 'udp' THEN 'dns-udp-queries'
+  END AS metricName,
+  SUM(query_count) AS count
+FROM
+  `%PROJECT_ID%.%DATASET_ID%.%TABLE_ID%`
+WHERE
+  STARTS_WITH(date_utc, '%YEAR_MONTH%')
+GROUP BY
+  tld,
+  metricName

core/src/test/java/google/registry/reporting/icann/ActivityReportingQueryBuilderTest.java

@@ -29,14 +29,11 @@ class ActivityReportingQueryBuilderTest {
   @SuppressWarnings("NonCanonicalType")
   private ActivityReportingQueryBuilder createQueryBuilder(String datasetName) {
     return new ActivityReportingQueryBuilder(
-        "domain-registry-alpha",
-        datasetName,
-        new BasicDnsCountQueryCoordinator(
-            new BasicDnsCountQueryCoordinator.Params(null, "domain-registry-alpha", datasetName)));
+        "domain-registry-alpha", datasetName, new DummyDnsCountQueryCoordinator());
   }
 
   @Test
-  void testAggregateQueryMatch_cloudSql() {
+  void testAggregateQueryMatch() {
     ActivityReportingQueryBuilder queryBuilder = createQueryBuilder("cloud_sql_icann_reporting");
     assertThat(queryBuilder.getReportQuery(yearMonth))
         .isEqualTo(
@@ -46,7 +43,7 @@ class ActivityReportingQueryBuilderTest {
   }
 
   @Test
-  void testIntermediaryQueryMatch_cloudSql() {
+  void testIntermediaryQueryMatch() {
     ImmutableList<String> expectedQueryNames =
         ImmutableList.of(
             ActivityReportingQueryBuilder.REGISTRAR_OPERATING_STATUS,
@@ -60,7 +57,7 @@ class ActivityReportingQueryBuilderTest {
     ImmutableMap<String, String> actualQueries = queryBuilder.getViewQueryMap(yearMonth);
     for (String queryName : expectedQueryNames) {
       String actualTableName = String.format("%s_201709", queryName);
-      String testFilename = String.format("%s_test_cloud_sql.sql", queryName);
+      String testFilename = String.format("%s_test.sql", queryName);
       assertThat(actualQueries.get(actualTableName))
           .isEqualTo(ReportingTestData.loadFile(testFilename));
     }

core/src/test/java/google/registry/reporting/icann/CloudDnsCountQueryCoordinatorTest.java

@@ -0,0 +1,46 @@
+// Copyright 2024 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.
+package google.registry.reporting.icann;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import org.joda.time.YearMonth;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+/** Unit tests for {@link CloudDnsCountQueryCoordinator}. */
+public class CloudDnsCountQueryCoordinatorTest {
+  public CloudDnsCountQueryCoordinatorTest() {}
+
+  private final YearMonth yearMonth = new YearMonth(2017, 9);
+  CloudDnsCountQueryCoordinator coordinator = new CloudDnsCountQueryCoordinator();
+
+  @BeforeEach
+  public void setUp() {
+    coordinator.projectId = "domain-registry-test";
+    coordinator.icannReportingDataSet = "icann_reporting";
+  }
+
+  @Test
+  public void testPreparatoryQueryConstruction() {
+    assertThat(coordinator.getPlxDnsTableQuery(yearMonth))
+        .isEqualTo(ReportingTestData.loadFile("prepare_dns_counts_cloud_test.sql"));
+  }
+
+  @Test
+  public void testQueryCreation() {
+    assertThat(coordinator.createQuery())
+        .isEqualTo(ReportingTestData.loadFile("dns_counts_cloud_test.sql"));
+  }
+}

core/src/test/java/google/registry/reporting/icann/IcannReportingStagerTest.java

@@ -50,7 +50,7 @@ class IcannReportingStagerTest {
     IcannReportingStager action = new IcannReportingStager();
     action.activityQueryBuilder =
         new ActivityReportingQueryBuilder(
-            "test-project", "icann_reporting", new BasicDnsCountQueryCoordinator(null));
+            "test-project", "icann_reporting", new DummyDnsCountQueryCoordinator());
     action.transactionsQueryBuilder =
         new TransactionsReportingQueryBuilder("test-project", "icann_reporting");
     action.reportingBucket = "test-bucket";

core/src/test/java/google/registry/reporting/icann/TransactionsReportingQueryBuilderTest.java

@@ -31,7 +31,7 @@ class TransactionsReportingQueryBuilderTest {
   }
 
   @Test
-  void testAggregateQueryMatch_cloud_sql() {
+  void testAggregateQueryMatch() {
     TransactionsReportingQueryBuilder queryBuilder =
         createQueryBuilder("cloud_sql_icann_reporting");
     assertThat(queryBuilder.getReportQuery(yearMonth))
@@ -43,7 +43,7 @@ class TransactionsReportingQueryBuilderTest {
   }
 
   @Test
-  void testIntermediaryQueryMatch_cloud_sql() {
+  void testIntermediaryQueryMatch() {
     ImmutableList<String> expectedQueryNames =
         ImmutableList.of(
             TransactionsReportingQueryBuilder.REGISTRAR_IANA_ID,
@@ -59,7 +59,7 @@ class TransactionsReportingQueryBuilderTest {
     ImmutableMap<String, String> actualQueries = queryBuilder.getViewQueryMap(yearMonth);
     for (String queryName : expectedQueryNames) {
       String actualTableName = String.format("%s_201709", queryName);
-      String testFilename = String.format("%s_test_cloud_sql.sql", queryName);
+      String testFilename = String.format("%s_test.sql", queryName);
       assertThat(actualQueries.get(actualTableName))
           .isEqualTo(ReportingTestData.loadFile(testFilename));
     }

core/src/test/resources/google/registry/reporting/icann/dns_counts_cloud_test.sql

@@ -0,0 +1,19 @@
+#standardSQL
+  -- Copyright 2017 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.
+
+  -- Retrieve per-TLD DNS query counts.
+
+SELECT *
+FROM `domain-registry-test.icann_reporting.dns_counts_from_plx`

core/src/test/resources/google/registry/reporting/icann/prepare_dns_counts_cloud_test.sql

@@ -0,0 +1,31 @@
+#standardSQL
+  -- Copyright 2024 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.
+
+  -- Retrieve per-TLD DNS query counts.
+
+SELECT
+  tld,
+  CASE
+    WHEN transport = 'tcp' THEN 'dns-tcp-queries'
+    WHEN transport = 'udp' THEN 'dns-udp-queries'
+  END AS metricName,
+  SUM(query_count) AS count
+FROM
+  `domain-registry-test.icann_reporting.zoneman_daily_query_counts`
+WHERE
+  STARTS_WITH(date_utc, '201709')
+GROUP BY
+  tld,
+  metricName