/*
 * Copyright © 2019 Cask Data, Inc.
 *
 * 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 io.cdap.plugin.gcp.datastore.sink;

import com.google.common.collect.MoreCollectors;
import com.google.datastore.v1.AllocateIdsRequest;
import com.google.datastore.v1.AllocateIdsResponse;
import com.google.datastore.v1.BeginTransactionRequest;
import com.google.datastore.v1.BeginTransactionResponse;
import com.google.datastore.v1.CommitRequest;
import com.google.datastore.v1.Entity;
import com.google.datastore.v1.client.Datastore;
import com.google.datastore.v1.client.DatastoreException;
import com.google.datastore.v1.client.DatastoreHelper;
import com.google.protobuf.ByteString;
import io.cdap.plugin.gcp.datastore.sink.util.DatastoreSinkConstants;
import io.cdap.plugin.gcp.datastore.util.DatastoreUtil;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.mapreduce.Counter;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormatCounter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;

/**
 * {@link DatastoreRecordWriter} writes the job outputs to the Datastore. Accepts <code>null</code> key, FullEntity
 * pairs but writes only FullEntities to the Datastore.
 */
public class DatastoreRecordWriter extends RecordWriter<NullWritable, Entity> {

  private static final Logger LOG = LoggerFactory.getLogger(DatastoreRecordWriter.class);

  private final Datastore datastore;
  private final int batchSize;
  private final boolean useAutogeneratedKey;
  private CommitRequest.Builder builder;
  private int totalCount;
  private int numberOfRecordsInBatch;
  private String projectId;
  private Counter counter;

  public DatastoreRecordWriter(TaskAttemptContext taskAttemptContext) throws IOException {
    Configuration config = taskAttemptContext.getConfiguration();
    this.projectId = config.get(DatastoreSinkConstants.CONFIG_PROJECT);
    String serviceAccount = config.get(DatastoreSinkConstants.CONFIG_SERVICE_ACCOUNT);
    Boolean isServiceAccountFilePath = config.getBoolean(DatastoreSinkConstants.CONFIG_SERVICE_ACCOUNT_IS_FILE_PATH,
                                                         true);
    this.batchSize = config.getInt(DatastoreSinkConstants.CONFIG_BATCH_SIZE, 25);
    this.useAutogeneratedKey = config.getBoolean(DatastoreSinkConstants.CONFIG_USE_AUTOGENERATED_KEY, false);
    LOG.debug("Initialize RecordWriter(projectId={}, batchSize={}, useAutogeneratedKey={}, "
      + "serviceAccount={})", projectId, batchSize, useAutogeneratedKey, serviceAccount);

    this.datastore = DatastoreUtil.getDatastoreV1(serviceAccount,
                                                  isServiceAccountFilePath,
                                                  projectId);

    this.totalCount = 0;
    this.numberOfRecordsInBatch = 0;
    this.builder = newCommitRequest();
    this.counter = taskAttemptContext.getCounter(FileOutputFormatCounter.BYTES_WRITTEN);
  }

  private CommitRequest.Builder newCommitRequest() throws IOException {
    // Create an RPC request to begin a new transaction.
    BeginTransactionRequest.Builder treq = BeginTransactionRequest.newBuilder();
    // Execute the RPC synchronously.
    CommitRequest.Builder builder = CommitRequest.newBuilder();
    BeginTransactionResponse tres;
    try {
      tres = datastore.beginTransaction(treq.build());
    } catch (DatastoreException e) {
      throw new IOException("Failed to begin datastore transaction", e);
    }
    // Get the transaction handle from the response.
    ByteString tx = tres.getTransaction();

    return builder
      .setProjectId(projectId)
      .setTransaction(tx);
  }

  @Override
  public void write(NullWritable key, Entity entity) throws IOException {
    LOG.trace("RecordWriter write({})", entity);
    if (useAutogeneratedKey) {
      AllocateIdsRequest request =
        AllocateIdsRequest.newBuilder()
          .setProjectId(projectId)
          .addKeys(entity.getKey())
          .build();
      AllocateIdsResponse response;
      try {
         response = datastore.allocateIds(request);
      } catch (DatastoreException e) {
        throw new IOException("Failed to allocate id", e);
      }

      Entity fullEntity = Entity.newBuilder()
        .setKey(response.getKeysList().stream().collect(MoreCollectors.onlyElement()))
        .putAllProperties(entity.getPropertiesMap())
        .build();
      builder.addMutations(DatastoreHelper.makeInsert(fullEntity).build());
    } else {
      builder.addMutations(DatastoreHelper.makeInsert(entity).build());
    }
    ++totalCount;
    ++numberOfRecordsInBatch;
    if (totalCount % batchSize == 0) {
      flush();
    }
  }

  @Override
  public void close(TaskAttemptContext taskAttemptContext) throws IOException {
    flush();
    LOG.debug("Total number of values written to Cloud Datastore: {}", totalCount);
  }

  private void flush() throws IOException {
    if (numberOfRecordsInBatch > 0) {
      LOG.debug("Writing a batch of {} values to Cloud Datastore.", numberOfRecordsInBatch);
      try {
        CommitRequest request = builder.build();
        datastore.commit(request);
        counter.increment(request.getSerializedSize());
      } catch (DatastoreException e) {
        throw new IOException("Datastore commit failed", e);
      }
      builder = newCommitRequest();
      numberOfRecordsInBatch = 0;
    }
  }
}

@@ -5185,26 +5185,26 @@
 eHelper.make
-In
+Up
 sert(entity)

Use upsert when writing to datastore