Use @DefaultCredential for Cloud API access in GAE

This change completes the switch to @DefaultCredential for
all use cases in GAE.

Impacted modules:
- IcannReporting
- CreateCdnsTld command
- LoadSnapshot command.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=213511730

java/google/registry/bigquery/BigqueryConnection.java

@@ -26,12 +26,8 @@ import static google.registry.bigquery.BigqueryUtils.toJobReferenceString;
 import static google.registry.config.RegistryConfig.getProjectId;
 import static org.joda.time.DateTimeZone.UTC;
 
-import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
 import com.google.api.client.googleapis.json.GoogleJsonResponseException;
 import com.google.api.client.http.AbstractInputStreamContent;
-import com.google.api.client.http.HttpRequestInitializer;
-import com.google.api.client.http.HttpTransport;
-import com.google.api.client.json.JsonFactory;
 import com.google.api.services.bigquery.Bigquery;
 import com.google.api.services.bigquery.model.Dataset;
 import com.google.api.services.bigquery.model.DatasetReference;
@@ -51,6 +47,7 @@ import com.google.api.services.bigquery.model.TableFieldSchema;
 import com.google.api.services.bigquery.model.TableReference;
 import com.google.api.services.bigquery.model.TableRow;
 import com.google.api.services.bigquery.model.ViewDefinition;
+import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableTable;
 import com.google.common.flogger.FluentLogger;
@@ -73,6 +70,7 @@ import java.util.List;
 import java.util.Random;
 import java.util.concurrent.ExecutorService;
 import javax.annotation.Nullable;
+import javax.inject.Inject;
 import org.joda.time.DateTime;
 import org.joda.time.Duration;
 
@@ -96,20 +94,11 @@ public class BigqueryConnection implements AutoCloseable {
   private static final Duration TEMP_TABLE_TTL = Duration.standardHours(24);
 
   /** Bigquery client instance wrapped by this class. */
-  private Bigquery bigquery;
+  private final Bigquery bigquery;
 
   /** Executor service for bigquery jobs. */
   private ListeningExecutorService service;
 
-  /** Credential object to use for initializing HTTP requests to the bigquery API. */
-  private HttpRequestInitializer credential;
-
-  /** HTTP transport object to use for accessing bigquery API. */
-  private HttpTransport httpTransport;
-
-  /** JSON factory object to use for accessing bigquery API. */
-  private JsonFactory jsonFactory;
-
   /** Pseudo-randomness source to use for creating random table names. */
   private Random random = new Random();
 
@@ -122,12 +111,17 @@ public class BigqueryConnection implements AutoCloseable {
   /** Duration to wait between polls for job status. */
   private Duration pollInterval = Duration.millis(1000);
 
+  BigqueryConnection(Bigquery bigquery) {
+    this.bigquery = bigquery;
+  }
+
   /** Builder for a {@link BigqueryConnection}, since the latter is immutable once created. */
   public static class Builder {
     private BigqueryConnection instance;
 
-    public Builder() {
-      instance = new BigqueryConnection();
+    @Inject
+    Builder(Bigquery bigquery) {
+      instance = new BigqueryConnection(bigquery);
     }
 
     /**
@@ -139,13 +133,6 @@ public class BigqueryConnection implements AutoCloseable {
       return this;
     }
 
-    public Builder setCredential(GoogleCredential credential) {
-      instance.credential = checkNotNull(credential);
-      instance.httpTransport = credential.getTransport();
-      instance.jsonFactory = credential.getJsonFactory();
-      return this;
-    }
-
     public Builder setDatasetId(String datasetId) {
       instance.datasetId = checkNotNull(datasetId);
       return this;
@@ -167,7 +154,11 @@ public class BigqueryConnection implements AutoCloseable {
     public BigqueryConnection build() {
       try {
         checkNotNull(instance.service, "Must provide executor service");
+        instance.initialize();
         return instance;
+      } catch (Throwable e) {
+        Throwables.throwIfUnchecked(e);
+        throw new RuntimeException("Cannot initialize BigqueryConnection", e);
       } finally {
         // Clear the internal instance so you can't accidentally mutate it through this builder.
         instance = null;
@@ -306,13 +297,10 @@ public class BigqueryConnection implements AutoCloseable {
   }
 
   /**
-   * Initializes the BigqueryConnection object by setting up the API client and creating the
-   * default dataset if it doesn't exist.
+   * Initializes the BigqueryConnection object by setting up the API client and creating the default
+   * dataset if it doesn't exist.
    */
-  public BigqueryConnection initialize() throws Exception {
-    bigquery = new Bigquery.Builder(httpTransport, jsonFactory, credential)
-        .setApplicationName(getClass().getSimpleName())
-        .build();
+  private BigqueryConnection initialize() throws Exception {
     createDatasetIfNeeded(datasetId);
     createDatasetIfNeeded(TEMP_DATASET_NAME);
     return this;

java/google/registry/bigquery/BigqueryModule.java

@@ -14,20 +14,16 @@
 
 package google.registry.bigquery;
 
-import com.google.api.client.http.HttpRequestInitializer;
-import com.google.api.client.http.HttpTransport;
-import com.google.api.client.json.JsonFactory;
+import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
 import com.google.api.services.bigquery.Bigquery;
-import com.google.api.services.bigquery.BigqueryScopes;
 import com.google.api.services.bigquery.model.TableFieldSchema;
 import com.google.common.collect.ImmutableList;
 import dagger.Module;
 import dagger.Provides;
 import dagger.multibindings.Multibinds;
+import google.registry.config.CredentialModule.DefaultCredential;
 import google.registry.config.RegistryConfig.Config;
 import java.util.Map;
-import java.util.Set;
-import java.util.function.Function;
 
 /** Dagger module for Google {@link Bigquery} connection objects. */
 @Module
@@ -39,11 +35,8 @@ public abstract class BigqueryModule {
 
   @Provides
   static Bigquery provideBigquery(
-      HttpTransport transport,
-      JsonFactory jsonFactory,
-      Function<Set<String>, ? extends HttpRequestInitializer> credential,
-      @Config("projectId") String projectId) {
-    return new Bigquery.Builder(transport, jsonFactory, credential.apply(BigqueryScopes.all()))
+      @DefaultCredential GoogleCredential credential, @Config("projectId") String projectId) {
+    return new Bigquery.Builder(credential.getTransport(), credential.getJsonFactory(), credential)
         .setApplicationName(projectId)
         .build();
   }

java/google/registry/bigquery/CheckedBigquery.java

@@ -1,4 +1,4 @@
-// Copyright 2017 The Nomulus Authors. All Rights Reserved.
+// 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.
@@ -17,14 +17,7 @@ package google.registry.bigquery;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.Sets.newConcurrentHashSet;
 
-import com.google.api.client.extensions.appengine.http.UrlFetchTransport;
-import com.google.api.client.googleapis.extensions.appengine.auth.oauth2.AppIdentityCredential;
-import com.google.api.client.http.HttpRequestInitializer;
-import com.google.api.client.http.HttpTransport;
-import com.google.api.client.json.JsonFactory;
-import com.google.api.client.json.jackson2.JacksonFactory;
 import com.google.api.services.bigquery.Bigquery;
-import com.google.api.services.bigquery.BigqueryScopes;
 import com.google.api.services.bigquery.model.Dataset;
 import com.google.api.services.bigquery.model.DatasetReference;
 import com.google.api.services.bigquery.model.Table;
@@ -39,8 +32,8 @@ import java.util.Map;
 import java.util.Set;
 import javax.inject.Inject;
 
-/** Factory for creating {@link Bigquery} connections. */
-public class BigqueryFactory {
+/** Wrapper of {@link Bigquery} with validation helpers. */
+public class CheckedBigquery {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
@@ -49,45 +42,16 @@ public class BigqueryFactory {
   private static Set<String> knownExistingTables = newConcurrentHashSet();
 
   @Inject Map<String, ImmutableList<TableFieldSchema>> bigquerySchemas;
-  @Inject Subfactory subfactory;
-  @Inject BigqueryFactory() {}
+  @Inject Bigquery bigquery;
 
-  /** This class is broken out solely so that it can be mocked inside of tests. */
-  static class Subfactory {
-
-    @Inject Subfactory() {}
-
-    public Bigquery create(
-        String applicationName,
-        HttpTransport transport,
-        JsonFactory jsonFactory,
-        HttpRequestInitializer httpRequestInitializer) {
-      return new Bigquery.Builder(transport, jsonFactory, httpRequestInitializer)
-          .setApplicationName(applicationName)
-          .build();
-    }
-  }
-
-  /** Returns a new connection to BigQuery. */
-  public Bigquery create(
-      String applicationName,
-      HttpTransport transport,
-      JsonFactory jsonFactory,
-      HttpRequestInitializer httpRequestInitializer) {
-    return subfactory.create(applicationName, transport, jsonFactory, httpRequestInitializer);
-  }
+  @Inject
+  CheckedBigquery() {}
 
   /**
    * Returns a new connection to Bigquery, first ensuring that the given dataset exists in the
    * project with the given id, creating it if required.
    */
-  public Bigquery create(String projectId, String datasetId) throws IOException {
-    Bigquery bigquery = create(
-        getClass().getSimpleName(),
-        new UrlFetchTransport(),
-        new JacksonFactory(),
-        new AppIdentityCredential(BigqueryScopes.all()));
-
+  public Bigquery ensureDataSetExists(String projectId, String datasetId) throws IOException {
     // Note: it's safe for multiple threads to call this as the dataset will only be created once.
     if (!knownExistingDatasets.contains(datasetId)) {
       ensureDataset(bigquery, projectId, datasetId);
@@ -101,9 +65,9 @@ public class BigqueryFactory {
    * Returns a new connection to Bigquery, first ensuring that the given dataset and table exist in
    * project with the given id, creating them if required.
    */
-  public Bigquery create(String projectId, String datasetId, String tableId)
+  public Bigquery ensureDataSetAndTableExist(String projectId, String datasetId, String tableId)
       throws IOException {
-    Bigquery bigquery = create(projectId, datasetId);
+    ensureDataSetExists(projectId, datasetId);
     checkArgument(bigquerySchemas.containsKey(tableId), "Unknown table ID: %s", tableId);
 
     if (!knownExistingTables.contains(tableId)) {

java/google/registry/config/files/default-config.yaml

@@ -274,6 +274,11 @@ misc:
   # hosts from being used on domains.
   asyncDeleteDelaySeconds: 90
 
+cloudDns:
+  # CloudDns testing config. Set both properties to null in Production.
+  rootUrl: https://staging-www.sandbox.googleapis.com
+  servicePath: dns/v2beta1_staging/projects/
+
 beam:
   # The default zone to run Apache Beam (Cloud Dataflow) jobs in.
   defaultJobZone: us-east1-c

java/google/registry/config/files/nomulus-config-production-sample.yaml

@@ -57,5 +57,9 @@ registrarConsole:
 misc:
   sheetExportId: placeholder
 
+cloudDns:
+  rootUrl: null
+  servicePath: null
+
 kms:
   projectId: placeholder

java/google/registry/export/LoadSnapshotAction.java

@@ -34,9 +34,9 @@ import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
-import google.registry.bigquery.BigqueryFactory;
 import google.registry.bigquery.BigqueryUtils.SourceFormat;
 import google.registry.bigquery.BigqueryUtils.WriteDisposition;
+import google.registry.bigquery.CheckedBigquery;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.export.BigqueryPollJobAction.BigqueryPollJobEnqueuer;
 import google.registry.request.Action;
@@ -70,7 +70,7 @@ public class LoadSnapshotAction implements Runnable {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
-  @Inject BigqueryFactory bigqueryFactory;
+  @Inject CheckedBigquery checkedBigquery;
   @Inject BigqueryPollJobEnqueuer bigqueryPollEnqueuer;
   @Inject Clock clock;
   @Inject @Config("projectId") String projectId;
@@ -109,7 +109,7 @@ public class LoadSnapshotAction implements Runnable {
 
   private String loadSnapshot(String snapshotId, String gcsFilename, Iterable<String> kinds)
       throws IOException {
-    Bigquery bigquery = bigqueryFactory.create(projectId, SNAPSHOTS_DATASET);
+    Bigquery bigquery = checkedBigquery.ensureDataSetExists(projectId, SNAPSHOTS_DATASET);
     DateTime now = clock.nowUtc();
     String loadMessage =
         String.format("Loading Datastore snapshot %s from %s...", snapshotId, gcsFilename);

java/google/registry/export/UpdateSnapshotViewAction.java

@@ -24,7 +24,7 @@ import com.google.api.services.bigquery.model.ViewDefinition;
 import com.google.appengine.api.taskqueue.TaskOptions;
 import com.google.appengine.api.taskqueue.TaskOptions.Method;
 import com.google.common.flogger.FluentLogger;
-import google.registry.bigquery.BigqueryFactory;
+import google.registry.bigquery.CheckedBigquery;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.request.Action;
 import google.registry.request.HttpException.InternalServerErrorException;
@@ -69,7 +69,7 @@ public class UpdateSnapshotViewAction implements Runnable {
   @Config("projectId")
   String projectId;
 
-  @Inject BigqueryFactory bigqueryFactory;
+  @Inject CheckedBigquery checkedBigquery;
 
   @Inject
   UpdateSnapshotViewAction() {}
@@ -106,7 +106,7 @@ public class UpdateSnapshotViewAction implements Runnable {
       SqlTemplate viewQueryTemplate)
       throws IOException {
 
-    Bigquery bigquery = bigqueryFactory.create(projectId, viewDataset);
+    Bigquery bigquery = checkedBigquery.ensureDataSetExists(projectId, viewDataset);
     updateTable(
         bigquery,
         new Table()

java/google/registry/monitoring/whitebox/MetricsExportAction.java

@@ -28,7 +28,7 @@ import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
-import google.registry.bigquery.BigqueryFactory;
+import google.registry.bigquery.CheckedBigquery;
 import google.registry.config.RegistryConfig.Config;
 import google.registry.request.Action;
 import google.registry.request.Parameter;
@@ -54,7 +54,8 @@ public class MetricsExportAction implements Runnable {
   @Inject @Parameter("tableId") String tableId;
   @Inject @Parameter("insertId") String insertId;
   @Inject @Config("projectId") String projectId;
-  @Inject BigqueryFactory bigqueryFactory;
+
+  @Inject CheckedBigquery checkedBigquery;
   @Inject @ParameterMap ImmutableListMultimap<String, String> parameters;
   @Inject MetricsExportAction() {}
 
@@ -62,7 +63,8 @@ public class MetricsExportAction implements Runnable {
   @Override
   public void run() {
     try {
-      Bigquery bigquery = bigqueryFactory.create(projectId, DATASET_ID, tableId);
+      Bigquery bigquery =
+          checkedBigquery.ensureDataSetAndTableExist(projectId, DATASET_ID, tableId);
       // Filter out the special parameters that the Action is called with.  Everything that's left
       // is returned in a Map that is suitable to pass to Bigquery as row data.
       Map<String, Object> jsonRows =

java/google/registry/reporting/icann/IcannReportingModule.java

@@ -18,9 +18,6 @@ import static google.registry.request.RequestParameters.extractOptionalEnumParam
 import static google.registry.request.RequestParameters.extractOptionalParameter;
 import static java.lang.annotation.RetentionPolicy.RUNTIME;
 
-import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
-import com.google.api.client.http.HttpTransport;
-import com.google.api.client.json.jackson2.JacksonFactory;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.MoreExecutors;
 import dagger.Module;
@@ -54,7 +51,6 @@ public final class IcannReportingModule {
   static final String DATASTORE_EXPORT_DATA_SET = "latest_datastore_export";
   static final String MANIFEST_FILE_NAME = "MANIFEST.txt";
   private static final String DEFAULT_SUBDIR = "icann/monthly";
-  private static final String BIGQUERY_SCOPE = "https://www.googleapis.com/auth/cloud-platform";
 
   /** Provides an optional subdirectory to store/upload reports to, extracted from the request. */
   @Provides
@@ -103,19 +99,16 @@ public final class IcannReportingModule {
    * @see google.registry.tools.BigqueryParameters for justifications of defaults.
    */
   @Provides
-  static BigqueryConnection provideBigqueryConnection(HttpTransport transport) {
+  static BigqueryConnection provideBigqueryConnection(
+      BigqueryConnection.Builder bigQueryConnectionBuilder) {
     try {
-      GoogleCredential credential = GoogleCredential
-          .getApplicationDefault(transport, new JacksonFactory());
       BigqueryConnection connection =
-          new BigqueryConnection.Builder()
+          bigQueryConnectionBuilder
               .setExecutorService(MoreExecutors.newDirectExecutorService())
-              .setCredential(credential.createScoped(ImmutableList.of(BIGQUERY_SCOPE)))
               .setDatasetId(ICANN_REPORTING_DATA_SET)
               .setOverwrite(true)
               .setPollInterval(Duration.standardSeconds(1))
               .build();
-      connection.initialize();
       return connection;
     } catch (Throwable e) {
       throw new RuntimeException("Could not initialize BigqueryConnection!", e);

java/google/registry/tools/BigqueryCommand.java

@@ -16,6 +16,8 @@ package google.registry.tools;
 
 import com.beust.jcommander.ParametersDelegate;
 import google.registry.bigquery.BigqueryConnection;
+import javax.inject.Inject;
+import javax.inject.Provider;
 
 /** A {@link Command} that uses the bigquery client API. */
 abstract class BigqueryCommand implements Command {
@@ -28,9 +30,12 @@ abstract class BigqueryCommand implements Command {
   /** Connection object for interacting with the Bigquery API. */
   private BigqueryConnection bigquery;
 
+  @Inject Provider<BigqueryConnection.Builder> bigQueryConnectionBuilderProvider;
+
   @Override
   public void run() throws Exception {
-    try (BigqueryConnection autoClosingBigquery = bigqueryParameters.newConnection()) {
+    try (BigqueryConnection autoClosingBigquery =
+        bigqueryParameters.newConnection(bigQueryConnectionBuilderProvider.get())) {
       bigquery = autoClosingBigquery;
       runWithBigquery();
     }

java/google/registry/tools/BigqueryParameters.java

@@ -16,13 +16,7 @@ package google.registry.tools;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
-import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
-import com.google.api.client.http.HttpTransport;
-import com.google.api.client.http.javanet.NetHttpTransport;
-import com.google.api.client.json.JsonFactory;
-import com.google.api.client.json.jackson2.JacksonFactory;
 import google.registry.bigquery.BigqueryConnection;
-import java.io.IOException;
 import java.util.concurrent.Executors;
 import org.joda.time.Duration;
 
@@ -57,31 +51,15 @@ final class BigqueryParameters {
       description = "Number of threads for running simultaneous BigQuery operations.")
   private int bigqueryNumThreads = DEFAULT_NUM_THREADS;
 
-  private static final HttpTransport HTTP_TRANSPORT = new NetHttpTransport();
-  private static final JsonFactory JSON_FACTORY = new JacksonFactory();
-
   /** Returns a new BigqueryConnection constructed according to the delegate's flag settings. */
-  BigqueryConnection newConnection() throws Exception {
-    BigqueryConnection connection = new BigqueryConnection.Builder()
-        .setExecutorService(Executors.newFixedThreadPool(bigqueryNumThreads))
-        .setCredential(newCredential())
-        .setDatasetId(bigqueryDataset)
-        .setOverwrite(bigqueryOverwrite)
-        .setPollInterval(bigqueryPollInterval)
-        .build();
-    connection.initialize();
+  BigqueryConnection newConnection(BigqueryConnection.Builder connectionBuilder) throws Exception {
+    BigqueryConnection connection =
+        connectionBuilder
+            .setExecutorService(Executors.newFixedThreadPool(bigqueryNumThreads))
+            .setDatasetId(bigqueryDataset)
+            .setOverwrite(bigqueryOverwrite)
+            .setPollInterval(bigqueryPollInterval)
+            .build();
     return connection;
   }
-
-  /** Creates a credential object for the Bigquery client using application default credentials. */
-  private GoogleCredential newCredential() {
-    try {
-      return GoogleCredential.getApplicationDefault(HTTP_TRANSPORT, JSON_FACTORY);
-    } catch (IOException e) {
-      throw new RuntimeException(
-          "Could not obtain application default credentials - "
-              + "did you remember to run 'gcloud auth application-default login'?",
-          e);
-    }
-  }
 }

java/google/registry/tools/CreateCdnsTld.java

@@ -14,27 +14,21 @@
 
 package google.registry.tools;
 
+import static com.google.common.base.Verify.verify;
+
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
-import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
-import com.google.api.client.googleapis.javanet.GoogleNetHttpTransport;
-import com.google.api.client.http.HttpTransport;
-import com.google.api.client.json.JsonFactory;
-import com.google.api.client.json.jackson2.JacksonFactory;
 import com.google.api.services.dns.Dns;
 import com.google.api.services.dns.model.ManagedZone;
 import com.google.api.services.dns.model.ManagedZoneDnsSecConfig;
-import com.google.common.annotations.VisibleForTesting;
 import google.registry.config.RegistryConfig.Config;
 import java.io.IOException;
-import java.security.GeneralSecurityException;
-import java.util.Arrays;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import javax.inject.Inject;
 
 @Parameters(separators = " =", commandDescription = "Create a Managed Zone for a TLD in Cloud DNS.")
-class CreateCdnsTld extends ConfirmingCommand {
+final class CreateCdnsTld extends ConfirmingCommand {
 
   @Parameter(names = "--description", description = "Description of the new TLD.")
   String description;
@@ -57,6 +51,8 @@ class CreateCdnsTld extends ConfirmingCommand {
   @Config("projectId")
   String projectId;
 
+  @Inject Dns dnsService;
+
   private static final String KEY_VALUE_FORMAT = "  %s = %s";
 
   private ManagedZone managedZone;
@@ -96,41 +92,19 @@ class CreateCdnsTld extends ConfirmingCommand {
   }
 
   @Override
-  public String execute() throws IOException, GeneralSecurityException {
-    Dns dnsService = createDnsService();
+  public String execute() throws IOException {
+    validateDnsService();
     Dns.ManagedZones.Create request = dnsService.managedZones().create(projectId, managedZone);
     ManagedZone response = request.execute();
     return String.format("Created managed zone: %s", response);
   }
 
-  @VisibleForTesting
-  Dns createDnsService() throws IOException, GeneralSecurityException {
-    // TODO(b/67367533): We should be obtaining the Dns instance from CloudDnsWriter module.  But
-    // to do this cleanly we need to refactor everything down to the credential object.  Having
-    // done that, this method will go away and this class will become final.
-    HttpTransport httpTransport = GoogleNetHttpTransport.newTrustedTransport();
-    JsonFactory jsonFactory = JacksonFactory.getDefaultInstance();
-
-    GoogleCredential credential = GoogleCredential.getApplicationDefault();
-    if (credential.createScopedRequired()) {
-      credential =
-          credential.createScoped(
-              Arrays.asList(
-                  "https://www.googleapis.com/auth/cloud-platform",
-                  "https://www.googleapis.com/auth/cloud-platform.read-only",
-                  "https://www.googleapis.com/auth/ndev.clouddns.readonly",
-                  "https://www.googleapis.com/auth/ndev.clouddns.readwrite"));
-    }
-
-    Dns.Builder builder =
-        new Dns.Builder(httpTransport, jsonFactory, credential).setApplicationName(projectId);
+  private void validateDnsService() {
+    // Sanity check to ensure only Production and Sandbox points to the CloudDns prod site.
     if (RegistryToolEnvironment.get() != RegistryToolEnvironment.PRODUCTION
         && RegistryToolEnvironment.get() != RegistryToolEnvironment.SANDBOX) {
-      builder
-          .setRootUrl("https://staging-www.sandbox.googleapis.com")
-          .setServicePath("dns/v2beta1_staging/projects/");
+      verify(!Dns.DEFAULT_ROOT_URL.equals(dnsService.getRootUrl()));
+      verify(!Dns.DEFAULT_SERVICE_PATH.equals(dnsService.getServicePath()));
     }
-
-    return builder.build();
   }
 }

java/google/registry/tools/RegistryToolComponent.java

@@ -16,6 +16,7 @@ package google.registry.tools;
 
 import com.google.monitoring.metrics.MetricWriter;
 import dagger.Component;
+import google.registry.bigquery.BigqueryModule;
 import google.registry.config.CredentialModule;
 import google.registry.config.RegistryConfig.ConfigModule;
 import google.registry.dns.writer.VoidDnsWriterModule;
@@ -53,6 +54,7 @@ import javax.inject.Singleton;
       // TODO(b/36866706): Find a way to replace this with a command-line friendly version
       AppIdentityCredentialModule.class,
       AuthModule.class,
+      BigqueryModule.class,
       ConfigModule.class,
       CredentialModule.class,
       DatastoreServiceModule.class,
@@ -99,6 +101,7 @@ interface RegistryToolComponent {
   void inject(GetKeyringSecretCommand command);
   void inject(GhostrydeCommand command);
   void inject(ListCursorsCommand command);
+  void inject(LoadSnapshotCommand command);
   void inject(LockDomainCommand command);
   void inject(LoginCommand command);
   void inject(LogoutCommand command);