<commit_before>/*
 * Copyright 2015 YCSB contributors. 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. See accompanying
 * LICENSE file.
 */

package com.yahoo.ycsb.db;

import com.google.api.client.auth.oauth2.Credential;
import com.google.api.services.datastore.DatastoreV1.*;
import com.google.api.services.datastore.DatastoreV1.CommitRequest.Mode;
import com.google.api.services.datastore.DatastoreV1.ReadOptions
  .ReadConsistency;
import com.google.api.services.datastore.client.Datastore;
import com.google.api.services.datastore.client.DatastoreException;
import com.google.api.services.datastore.client.DatastoreFactory;
import com.google.api.services.datastore.client.DatastoreHelper;
import com.google.api.services.datastore.client.DatastoreOptions;

import com.yahoo.ycsb.ByteIterator;
import com.yahoo.ycsb.DB;
import com.yahoo.ycsb.DBException;
import com.yahoo.ycsb.Status;
import com.yahoo.ycsb.StringByteIterator;

import org.apache.log4j.Level;
import org.apache.log4j.Logger;

import java.io.IOException;
import java.security.GeneralSecurityException;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.Vector;

import javax.annotation.Nullable;

/**
 * Google Cloud Datastore Client for YCSB.
 */

public class GoogleDatastoreClient extends DB {
  /**
   * Defines a MutationType used in this class.
   */
  private enum MutationType {
    UPSERT,
    UPDATE,
    DELETE
  }

  /**
   * Defines a EntityGroupingMode enum used in this class.
   */
  private enum EntityGroupingMode {
    ONE_ENTITY_PER_GROUP,
    MULTI_ENTITY_PER_GROUP
  }

  private static Logger logger =
      Logger.getLogger(GoogleDatastoreClient.class);

  // Read consistency defaults to "STRONG" per YCSB guidance.
  // User can override this via configure.
  private ReadConsistency readConsistency = ReadConsistency.STRONG;

  private EntityGroupingMode entityGroupingMode =
      EntityGroupingMode.ONE_ENTITY_PER_GROUP;

  private String rootEntityName;

  private Datastore datastore = null;

  public GoogleDatastoreClient() {}

  /**
   * Initialize any state for this DB. Called once per DB instance; there is
   * one DB instance per client thread.
   */
  @Override
  public void init() throws DBException {
    String debug = getProperties().getProperty("googledatastore.debug", null);
    if (null != debug && "true".equalsIgnoreCase(debug)) {
      logger.setLevel(Level.DEBUG);
    }

    // We need the following 3 essential properties to initialize datastore:
    //
    // - DatasetId,
    // - Path to private key file,
    // - Service account email address.
    String datasetId = getProperties().getProperty(
        "googledatastore.datasetId", null);
    if (datasetId == null) {
      throw new DBException(
          "Required property \"datasetId\" missing.");
    }

    String privateKeyFile = getProperties().getProperty(
        "googledatastore.privateKeyFile", null);
    if (privateKeyFile == null) {
      throw new DBException(
          "Required property \"privateKeyFile\" missing.");
    }

    String serviceAccountEmail = getProperties().getProperty(
        "googledatastore.serviceAccountEmail", null);
    if (serviceAccountEmail == null) {
      throw new DBException(
          "Required property \"serviceAccountEmail\" missing.");
    }

    // Below are properties related to benchmarking.

    String readConsistencyConfig = getProperties().getProperty(
        "googledatastore.readConsistency", null);
    if (readConsistencyConfig != null) {
      try {
        this.readConsistency = ReadConsistency.valueOf(
            readConsistencyConfig.trim().toUpperCase());
      } catch (IllegalArgumentException e) {
        throw new DBException("Invalid read consistency specified: " +
            readConsistencyConfig + ". Expecting STRONG or EVENTUAL.");
      }
    }

    //
    // Entity Grouping Mode (googledatastore.entitygroupingmode), see
    // documentation in conf/googledatastore.properties.
    //
    String entityGroupingConfig = getProperties().getProperty(
        "googledatastore.entityGroupingMode", null);
    if (entityGroupingConfig != null) {
      try {
        this.entityGroupingMode = EntityGroupingMode.valueOf(
            entityGroupingConfig.trim().toUpperCase());
      } catch (IllegalArgumentException e) {
        throw new DBException("Invalid entity grouping mode specified: " +
            entityGroupingConfig + ". Expecting ONE_ENTITY_PER_GROUP or " +
            "MULTI_ENTITY_PER_GROUP.");
      }
    }

    this.rootEntityName = getProperties().getProperty(
        "googledatastore.rootEntityName", "YCSB_ROOT_ENTITY");

    try {
      // Setup the connection to Google Cloud Datastore with the credentials
      // obtained from the configure.
      DatastoreOptions.Builder options = new DatastoreOptions.Builder();
      Credential credential = DatastoreHelper.getServiceAccountCredential(
          serviceAccountEmail, privateKeyFile);
      logger.info("Using JWT Service Account credential.");
      logger.info("DatasetID: " + datasetId + ", Service Account Email: " +
          serviceAccountEmail + ", Private Key File Path: " + privateKeyFile);

      datastore = DatastoreFactory.get().create(
          options.credential(credential).dataset(datasetId).build());

    } catch (GeneralSecurityException exception) {
      throw new DBException("Security error connecting to the datastore: " +
            "cause: " + exception.getCause() +
            " details: " + exception.getMessage());

    } catch (IOException exception) {
      throw new DBException("I/O error connecting to the datastore: " +
            "cause: " + exception.getCause() +
            " details: " + exception.getMessage());
    }

    logger.info("Datastore client instance created: " +
        datastore.toString());
  }

  @Override<commit_msg>[googledatastore] Add cause to DBException<commit_after>/*
 * Copyright 2015 YCSB contributors. 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. See accompanying
 * LICENSE file.
 */

package com.yahoo.ycsb.db;

import com.google.api.client.auth.oauth2.Credential;
import com.google.api.services.datastore.DatastoreV1.*;
import com.google.api.services.datastore.DatastoreV1.CommitRequest.Mode;
import com.google.api.services.datastore.DatastoreV1.ReadOptions
  .ReadConsistency;
import com.google.api.services.datastore.client.Datastore;
import com.google.api.services.datastore.client.DatastoreException;
import com.google.api.services.datastore.client.DatastoreFactory;
import com.google.api.services.datastore.client.DatastoreHelper;
import com.google.api.services.datastore.client.DatastoreOptions;

import com.yahoo.ycsb.ByteIterator;
import com.yahoo.ycsb.DB;
import com.yahoo.ycsb.DBException;
import com.yahoo.ycsb.Status;
import com.yahoo.ycsb.StringByteIterator;

import org.apache.log4j.Level;
import org.apache.log4j.Logger;

import java.io.IOException;
import java.security.GeneralSecurityException;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.Vector;

import javax.annotation.Nullable;

/**
 * Google Cloud Datastore Client for YCSB.
 */

public class GoogleDatastoreClient extends DB {
  /**
   * Defines a MutationType used in this class.
   */
  private enum MutationType {
    UPSERT,
    UPDATE,
    DELETE
  }

  /**
   * Defines a EntityGroupingMode enum used in this class.
   */
  private enum EntityGroupingMode {
    ONE_ENTITY_PER_GROUP,
    MULTI_ENTITY_PER_GROUP
  }

  private static Logger logger =
      Logger.getLogger(GoogleDatastoreClient.class);

  // Read consistency defaults to "STRONG" per YCSB guidance.
  // User can override this via configure.
  private ReadConsistency readConsistency = ReadConsistency.STRONG;

  private EntityGroupingMode entityGroupingMode =
      EntityGroupingMode.ONE_ENTITY_PER_GROUP;

  private String rootEntityName;

  private Datastore datastore = null;

  public GoogleDatastoreClient() {}

  /**
   * Initialize any state for this DB. Called once per DB instance; there is
   * one DB instance per client thread.
   */
  @Override
  public void init() throws DBException {
    String debug = getProperties().getProperty("googledatastore.debug", null);
    if (null != debug && "true".equalsIgnoreCase(debug)) {
      logger.setLevel(Level.DEBUG);
    }

    // We need the following 3 essential properties to initialize datastore:
    //
    // - DatasetId,
    // - Path to private key file,
    // - Service account email address.
    String datasetId = getProperties().getProperty(
        "googledatastore.datasetId", null);
    if (datasetId == null) {
      throw new DBException(
          "Required property \"datasetId\" missing.");
    }

    String privateKeyFile = getProperties().getProperty(
        "googledatastore.privateKeyFile", null);
    if (privateKeyFile == null) {
      throw new DBException(
          "Required property \"privateKeyFile\" missing.");
    }

    String serviceAccountEmail = getProperties().getProperty(
        "googledatastore.serviceAccountEmail", null);
    if (serviceAccountEmail == null) {
      throw new DBException(
          "Required property \"serviceAccountEmail\" missing.");
    }

    // Below are properties related to benchmarking.

    String readConsistencyConfig = getProperties().getProperty(
        "googledatastore.readConsistency", null);
    if (readConsistencyConfig != null) {
      try {
        this.readConsistency = ReadConsistency.valueOf(
            readConsistencyConfig.trim().toUpperCase());
      } catch (IllegalArgumentException e) {
        throw new DBException("Invalid read consistency specified: " +
            readConsistencyConfig + ". Expecting STRONG or EVENTUAL.");
      }
    }

    //
    // Entity Grouping Mode (googledatastore.entitygroupingmode), see
    // documentation in conf/googledatastore.properties.
    //
    String entityGroupingConfig = getProperties().getProperty(
        "googledatastore.entityGroupingMode", null);
    if (entityGroupingConfig != null) {
      try {
        this.entityGroupingMode = EntityGroupingMode.valueOf(
            entityGroupingConfig.trim().toUpperCase());
      } catch (IllegalArgumentException e) {
        throw new DBException("Invalid entity grouping mode specified: " +
            entityGroupingConfig + ". Expecting ONE_ENTITY_PER_GROUP or " +
            "MULTI_ENTITY_PER_GROUP.");
      }
    }

    this.rootEntityName = getProperties().getProperty(
        "googledatastore.rootEntityName", "YCSB_ROOT_ENTITY");

    try {
      // Setup the connection to Google Cloud Datastore with the credentials
      // obtained from the configure.
      DatastoreOptions.Builder options = new DatastoreOptions.Builder();
      Credential credential = DatastoreHelper.getServiceAccountCredential(
          serviceAccountEmail, privateKeyFile);
      logger.info("Using JWT Service Account credential.");
      logger.info("DatasetID: " + datasetId + ", Service Account Email: " +
          serviceAccountEmail + ", Private Key File Path: " + privateKeyFile);

      datastore = DatastoreFactory.get().create(
          options.credential(credential).dataset(datasetId).build());

    } catch (GeneralSecurityException exception) {
      throw new DBException("Security error connecting to the datastore: " +
          exception.getMessage(), exception);

    } catch (IOException exception) {
      throw new DBException("I/O error connecting to the datastore: " +
          exception.getMessage(), exception);
    }

    logger.info("Datastore client instance created: " +
        datastore.toString());
  }

  @Override
{'fraction_non_alphanumeric': 0.09649878593248218, 'fraction_numerical': 0.002349808099005248, 'mean_word_length': 3.389137160536267, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '260798', 'n_tokens_mistral': 3611, 'n_tokens_neox': 3526, 'n_words': 925}