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

// 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.

package google.registry.monitoring.whitebox;

import static com.google.common.base.Predicates.in;
import static com.google.common.base.Predicates.not;
import static com.google.common.collect.Multimaps.filterKeys;
import static google.registry.request.Action.Method.POST;
import static google.registry.util.FormattingLogger.getLoggerForCallerClass;
import static java.util.stream.Collectors.joining;

import com.google.api.services.bigquery.Bigquery;
import com.google.api.services.bigquery.model.TableDataInsertAllRequest;
import com.google.api.services.bigquery.model.TableDataInsertAllResponse;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableListMultimap;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import google.registry.bigquery.BigqueryFactory;
import google.registry.config.RegistryConfig.Config;
import google.registry.request.Action;
import google.registry.request.Parameter;
import google.registry.request.ParameterMap;
import google.registry.request.auth.Auth;
import google.registry.util.FormattingLogger;
import java.io.IOException;
import java.util.Map;
import javax.inject.Inject;

/** Action for exporting metrics to BigQuery. */
@Action(
  path = MetricsExportAction.PATH,
  method = POST,
  auth = Auth.AUTH_INTERNAL_ONLY
)
public class MetricsExportAction implements Runnable {

  public static final String PATH = "/_dr/task/metrics";
  private static final FormattingLogger logger = getLoggerForCallerClass();
  private static final String DATASET_ID = "metrics";
  private static final ImmutableSet<String> SPECIAL_PARAMS = ImmutableSet.of("tableId", "insertId");

  @Inject @Parameter("tableId") String tableId;
  @Inject @Parameter("insertId") String insertId;
  @Inject @Config("projectId") String projectId;
  @Inject BigqueryFactory bigqueryFactory;
  @Inject @ParameterMap ImmutableListMultimap<String, String> parameters;
  @Inject MetricsExportAction() {}

  /** Exports metrics to BigQuery. */
  @Override
  public void run() {
    try {
      Bigquery bigquery = bigqueryFactory.create(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 =
          ImmutableMap.copyOf(
              filterKeys(parameters, not(in(SPECIAL_PARAMS))).entries());
      TableDataInsertAllResponse response = bigquery.tabledata()
          .insertAll(
              projectId,
              DATASET_ID,
              tableId,
              new TableDataInsertAllRequest()
                  .setRows(
                      ImmutableList.of(new TableDataInsertAllRequest.Rows()
                          .setInsertId(insertId)
                          .setJson(jsonRows))))
          .execute();

      if (response.getInsertErrors() != null && !response.getInsertErrors().isEmpty()) {
        throw new RuntimeException(
            response
                .getInsertErrors()
                .stream()
                .map(
                    error -> {
                      try {
                        return error.toPrettyString();
                      } catch (IOException e) {
                        return error.toString();
                      }
                    })
                .collect(joining("\n")));
      }
    } catch (Throwable e) {
      logger.warningfmt("Caught Unknown Exception: %s", e);
    }
  }
}