fix retry helper for datastore and change the way we get a default in…

…stance of the service
googleapis · Apr 21, 2015 · 7e932b7 · 7e932b7
1 parent 8e0c9d3
commit 7e932b7
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 7 deletions.
diff --git a/src/main/java/com/google/gcloud/datastore/DatastoreServiceFactory.java b/src/main/java/com/google/gcloud/datastore/DatastoreServiceFactory.java
@@ -27,9 +27,15 @@ public DatastoreService get(DatastoreServiceOptions options) {
       }
     };
 
-  public static DatastoreService getDefault(DatastoreServiceOptions options) {
-    return INSTANCE.get(options);
+  /**
+   * Returns the default factory instance.
+   */
+  public static DatastoreServiceFactory instance() {
+    return INSTANCE;
   }
 
+  /**
+   * Returns a {@code DatastoreService} for the given options.
+   */
   public abstract DatastoreService get(DatastoreServiceOptions options);
 }
diff --git a/src/main/java/com/google/gcloud/datastore/DatastoreServiceImpl.java b/src/main/java/com/google/gcloud/datastore/DatastoreServiceImpl.java
@@ -60,8 +60,8 @@ public RetryResult afterEval(Exception exception, RetryResult retryResult) {
 
         @Override
         public RetryResult beforeEval(Exception exception) {
-          if (exception instanceof DatastoreServiceException) {
-            boolean retriable = ((DatastoreServiceException) exception).code().retriable();
+          if (exception instanceof DatastoreRpcException) {
+            boolean retriable = ((DatastoreRpcException) exception).retryable();
             return retriable ? Interceptor.RetryResult.RETRY : Interceptor.RetryResult.ABORT;
           }
           return null;

diff --git a/src/main/java/com/google/gcloud/datastore/package-info.java b/src/main/java/com/google/gcloud/datastore/package-info.java
@@ -20,7 +20,7 @@
  * <p>A simple usage example:
  * <pre> {@code
  * DatastoreServiceOptions options = DatastoreServiceOptions.builder().dataset(DATASET).build();
- * DatastoreService datastore = DatastoreServiceFactory.getDefault(options);
+ * DatastoreService datastore = DatastoreServiceFactory.instance().get(options);
  * KeyFactory keyFactory = datastore.newKeyFactory().kind(kind);
  * Key key = keyFactory.newKey(keyName);
  * Entity entity = datastore.get(key);

diff --git a/src/main/java/com/google/gcloud/examples/DatastoreExample.java b/src/main/java/com/google/gcloud/examples/DatastoreExample.java
@@ -187,7 +187,7 @@ public static void main(String... args) {
           .namespace(NAMESPACE)
           .build();
       String name = args.length > 1 ? args[1] : System.getProperty("user.name");
-      datastore = DatastoreServiceFactory.getDefault(options);
+      datastore = DatastoreServiceFactory.instance().get(options);
       KeyFactory keyFactory = datastore.newKeyFactory().kind(USER_KIND);
       key = keyFactory.newKey(name);
       String actionName = args.length > 2 ? args[2].toLowerCase() : DEFAULT_ACTION;

diff --git a/src/test/java/com/google/gcloud/datastore/DatastoreServiceTest.java b/src/test/java/com/google/gcloud/datastore/DatastoreServiceTest.java
@@ -25,12 +25,18 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
+import com.google.api.services.datastore.DatastoreV1;
+import com.google.api.services.datastore.DatastoreV1.EntityResult;
 import com.google.common.collect.Iterators;
+import com.google.gcloud.RetryParams;
 import com.google.gcloud.datastore.Query.ResultType;
 import com.google.gcloud.datastore.StructuredQuery.OrderBy;
 import com.google.gcloud.datastore.StructuredQuery.Projection;
 import com.google.gcloud.datastore.StructuredQuery.PropertyFilter;
+import com.google.gcloud.spi.DatastoreRpc;
+import com.google.gcloud.spi.DatastoreRpc.DatastoreRpcException.Reason;
 
+import org.easymock.EasyMock;
 import org.junit.AfterClass;
 import org.junit.Before;
 import org.junit.BeforeClass;
@@ -106,7 +112,7 @@ public void setUp() throws IOException, InterruptedException {
         .dataset(DATASET)
         .host("http://localhost:" + LocalGcdHelper.PORT)
         .build();
-    datastore = DatastoreServiceFactory.getDefault(options);
+    datastore = DatastoreServiceFactory.instance().get(options);
     StructuredQuery<Key> query = Query.keyQueryBuilder().build();
     QueryResults<Key> result = datastore.run(query);
     datastore.delete(Iterators.toArray(result, Key.class));
@@ -624,4 +630,25 @@ public void testKeyFactory() {
     assertEquals(KEY1, keyFactory.newKey("name"));
     assertEquals(Key.builder(KEY1).id(2).build(), keyFactory.newKey(2));
   }
+
+  @Test
+  public void testRetires() throws Exception {
+    DatastoreV1.LookupRequest requestPb =
+        DatastoreV1.LookupRequest.newBuilder().addKey(KEY1.toPb()).build();
+    DatastoreV1.LookupResponse responsePb = DatastoreV1.LookupResponse.newBuilder()
+        .addFound(EntityResult.newBuilder().setEntity(ENTITY1.toPb())).build();
+    DatastoreRpc mock = EasyMock.createStrictMock(DatastoreRpc.class);
+    EasyMock.expect(mock.lookup(requestPb))
+        .andThrow(new DatastoreRpc.DatastoreRpcException(Reason.UNAVAILABLE))
+        .andReturn(responsePb);
+    EasyMock.replay(mock);
+    DatastoreServiceOptions options = this.options.toBuilder()
+        .retryParams(RetryParams.getDefaultInstance())
+        .datastoreRpc(mock)
+        .build();
+    DatastoreService datastore = DatastoreServiceFactory.instance().get(options);
+    Entity entity = datastore.get(KEY1);
+    assertEquals(ENTITY1, entity);
+    EasyMock.verify(mock);
+  }
 }