From db2e3dfb0b8d1abb801eefa470269d04c81bca3b Mon Sep 17 00:00:00 2001
From: haih-g <haih@google.com>
Date: Sat, 16 Sep 2017 12:15:42 -0700
Subject: [PATCH] [googledatastore] Google Datastore to use default
 credentials. (#984)

* Google Datastore to use default credentials.

* Google Datastore to skip Index.
---
 .../conf/googledatastore.properties           | 11 ++++--
 .../yahoo/ycsb/db/GoogleDatastoreClient.java  | 38 +++++++++++--------
 2 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/googledatastore/conf/googledatastore.properties b/googledatastore/conf/googledatastore.properties
index 408acf0d..1418b639 100644
--- a/googledatastore/conf/googledatastore.properties
+++ b/googledatastore/conf/googledatastore.properties
@@ -14,7 +14,7 @@
 # LICENSE file.
 
 #
-# Sample property file for Google Cloud Datastore DB client 
+# Sample property file for Google Cloud Datastore DB client
 
 ## Mandatory parameters
 #
@@ -39,12 +39,12 @@ readallfields = true
 #
 # googledatastore.readConsistency=STRONG
 
-# Decides how we group entities into entity groups. 
+# Decides how we group entities into entity groups.
 # (See the details section in README.md for documentation)
 #
 # googledatastore.entityGroupingMode=ONE_ENTITY_PER_GROUP
 
-# If you set the googledatastore.entityGroupingMode property to 
+# If you set the googledatastore.entityGroupingMode property to
 # MULTI_ENTITY_PER_GROUP, you can optionally specify the name of the root entity
 #
 # googledatastore.rootEntityName="YCSB_ROOT_ENTITY"
@@ -53,4 +53,7 @@ readallfields = true
 # requestdistribution = uniform
 
 # Enable/disable debug message, default is false.
-# googledatastore.debug = false
\ No newline at end of file
+# googledatastore.debug = false
+
+# Skip indexes, default is true.
+# googledatastore.skipIndex = true
\ No newline at end of file
diff --git a/googledatastore/src/main/java/com/yahoo/ycsb/db/GoogleDatastoreClient.java b/googledatastore/src/main/java/com/yahoo/ycsb/db/GoogleDatastoreClient.java
index 62db5833..03408cb8 100644
--- a/googledatastore/src/main/java/com/yahoo/ycsb/db/GoogleDatastoreClient.java
+++ b/googledatastore/src/main/java/com/yahoo/ycsb/db/GoogleDatastoreClient.java
@@ -18,6 +18,7 @@
 package com.yahoo.ycsb.db;
 
 import com.google.api.client.auth.oauth2.Credential;
+import com.google.api.client.googleapis.auth.oauth2.GoogleCredential;
 import com.google.datastore.v1.*;
 import com.google.datastore.v1.CommitRequest.Mode;
 import com.google.datastore.v1.ReadOptions.ReadConsistency;
@@ -82,6 +83,8 @@ public class GoogleDatastoreClient extends DB {
 
   private Datastore datastore = null;
 
+  private static boolean skipIndex = true;
+
   /**
    * Initialize any state for this DB. Called once per DB instance; there is
    * one DB instance per client thread.
@@ -93,6 +96,12 @@ public class GoogleDatastoreClient extends DB {
       logger.setLevel(Level.DEBUG);
     }
 
+    String skipIndexString = getProperties().getProperty(
+        "googledatastore.skipIndex", null);
+    if (null != skipIndexString && "false".equalsIgnoreCase(skipIndexString)) {
+      skipIndex = false;
+    }
+
     // We need the following 3 essential properties to initialize datastore:
     //
     // - DatasetId,
@@ -107,17 +116,8 @@ public class GoogleDatastoreClient extends DB {
 
     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.
 
@@ -157,11 +157,18 @@ public class GoogleDatastoreClient extends DB {
       // 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);
+      Credential credential = GoogleCredential.getApplicationDefault();
+      if (serviceAccountEmail != null && privateKeyFile != null) {
+        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);
+      } else {
+        logger.info("Using default gcloud credential.");
+        logger.info("DatasetID: " + datasetId
+            + ", Service Account Email: " + ((GoogleCredential) credential).getServiceAccountId());
+      }
 
       datastore = DatastoreFactory.get().create(
           options.credential(credential).projectId(datasetId).build());
@@ -298,7 +305,8 @@ public class GoogleDatastoreClient extends DB {
         entityBuilder.getMutableProperties()
             .put(val.getKey(),
                 Value.newBuilder()
-                .setStringValue(val.getValue().toString()).build());
+                .setStringValue(val.getValue().toString())
+                .setExcludeFromIndexes(skipIndex).build());
       }
       Entity entity = entityBuilder.build();
       logger.debug("entity built as: " + entity.toString());
-- 
GitLab