feat!: Making Datastore implement Closeable

README.md

@@ -50,20 +50,20 @@ If you are using Maven without the BOM, add this to your dependencies:
 If you are using Gradle 5.x or later, add this to your dependencies:
 
 ```Groovy
-implementation platform('com.google.cloud:libraries-bom:26.27.0')
+implementation platform('com.google.cloud:libraries-bom:26.29.0')
 
 implementation 'com.google.cloud:google-cloud-datastore'
 ```
 If you are using Gradle without BOM, add this to your dependencies:
 
 ```Groovy
-implementation 'com.google.cloud:google-cloud-datastore:2.17.5'
+implementation 'com.google.cloud:google-cloud-datastore:2.17.6'
 ```
 
 If you are using SBT, add this to your dependencies:
 
 ```Scala
-libraryDependencies += "com.google.cloud" % "google-cloud-datastore" % "2.17.5"
+libraryDependencies += "com.google.cloud" % "google-cloud-datastore" % "2.17.6"
 ```
 <!-- {x-version-update-end} -->
 
@@ -380,7 +380,7 @@ Java is a registered trademark of Oracle and/or its affiliates.
 [kokoro-badge-link-5]: http://storage.googleapis.com/cloud-devrel-public/java/badges/java-datastore/java11.html
 [stability-image]: https://img.shields.io/badge/stability-stable-green
 [maven-version-image]: https://img.shields.io/maven-central/v/com.google.cloud/google-cloud-datastore.svg
-[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-datastore/2.17.5
+[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-datastore/2.17.6
 [authentication]: https://github.com/googleapis/google-cloud-java#authentication
 [auth-scopes]: https://developers.google.com/identity/protocols/oauth2/scopes
 [predefined-iam-roles]: https://cloud.google.com/iam/docs/understanding-roles#predefined_roles

google-cloud-datastore/clirr-ignored-differences.xml

@@ -38,4 +38,25 @@
         <differenceType>5001</differenceType>
         <to>com/google/cloud/http/BaseHttpServiceException</to>
     </difference>
+    <difference>
+        <className>com/google/cloud/datastore/Datastore</className>
+        <method>void close()</method>
+        <differenceType>7012</differenceType>
+    </difference>
+    <difference>
+        <className>com/google/cloud/datastore/spi/v1/DatastoreRpc</className>
+        <method>void close()</method>
+        <differenceType>7012</differenceType>
+    </difference>
+    <difference>
+        <className>com/google/cloud/datastore/Datastore</className>
+        <method>boolean isClosed()</method>
+        <differenceType>7012</differenceType>
+    </difference>
+    <difference>
+        <className>com/google/cloud/datastore/spi/v1/DatastoreRpc</className>
+        <method>boolean isClosed()</method>
+        <differenceType>7012</differenceType>
+    </difference>
+
 </differences>

google-cloud-datastore/src/main/java/com/google/cloud/datastore/Datastore.java

@@ -22,7 +22,7 @@
 import java.util.List;
 
 /** An interface for Google Cloud Datastore. */
-public interface Datastore extends Service<DatastoreOptions>, DatastoreReaderWriter {
+public interface Datastore extends Service<DatastoreOptions>, DatastoreReaderWriter, AutoCloseable {
 
   /**
    * Returns a new Datastore transaction.
@@ -49,9 +49,9 @@ public interface Datastore extends Service<DatastoreOptions>, DatastoreReaderWri
    * @param <T> the type of the return value
    */
   interface TransactionCallable<T> {
+
     T run(DatastoreReaderWriter readerWriter) throws Exception;
   }
-
   /**
    * Invokes the callback's {@link Datastore.TransactionCallable#run} method with a {@link
    * DatastoreReaderWriter} that is associated with a new transaction. The transaction will be
@@ -508,4 +508,15 @@ interface TransactionCallable<T> {
   default AggregationResults runAggregation(AggregationQuery query, ReadOption... options) {
     throw new UnsupportedOperationException("Not implemented.");
   }
+
+  /**
+   * Closes the gRPC channels associated with this instance and frees up their resources. This
+   * method blocks until all channels are closed. Once this method is called, this Datastore client
+   * is no longer usable.
+   */
+  @Override
+  void close() throws Exception;
+
+  /** Returns true if this background resource has been shut down. */
+  boolean isClosed();
 }

google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java

@@ -48,9 +48,12 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Callable;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 
 final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datastore {
 
+  Logger logger = Logger.getLogger(Datastore.class.getName());
   private final DatastoreRpc datastoreRpc;
   private final RetrySettings retrySettings;
   private static final ExceptionHandler TRANSACTION_EXCEPTION_HANDLER =
@@ -90,6 +93,20 @@ public Transaction newTransaction() {
     return new TransactionImpl(this);
   }
 
+  @Override
+  public void close() throws Exception {
+    try {
+      datastoreRpc.close();
+    } catch (Exception e) {
+      logger.log(Level.WARNING, "Failed to close channels", e);
+    }
+  }
+
+  @Override
+  public boolean isClosed() {
+    return datastoreRpc.isClosed();
+  }
+
   static class ReadWriteTransactionCallable<T> implements Callable<T> {
 
     private final Datastore datastore;

google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryAndTraceDatastoreRpcDecorator.java

@@ -109,6 +109,16 @@ public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryReques
         () -> datastoreRpc.runAggregationQuery(request), SPAN_NAME_RUN_AGGREGATION_QUERY);
   }
 
+  @Override
+  public void close() throws Exception {
+    datastoreRpc.close();
+  }
+
+  @Override
+  public boolean isClosed() {
+    return datastoreRpc.isClosed();
+  }
+
   public <O> O invokeRpc(Callable<O> block, String startSpan) {
     Span span = traceUtil.startSpan(startSpan);
     try (Scope scope = traceUtil.getTracer().withSpan(span)) {

google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/DatastoreRpc.java

@@ -36,7 +36,7 @@
 import com.google.datastore.v1.RunQueryResponse;
 
 /** Provides access to the remote Datastore service. */
-public interface DatastoreRpc extends ServiceRpc {
+public interface DatastoreRpc extends ServiceRpc, AutoCloseable {
 
   /**
    * Sends an allocate IDs request.
@@ -96,4 +96,10 @@ BeginTransactionResponse beginTransaction(BeginTransactionRequest request)
   default RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryRequest request) {
     throw new UnsupportedOperationException("Not implemented.");
   }
+
+  @Override
+  void close() throws Exception;
+
+  /** Returns true if this background resource has been shut down. */
+  boolean isClosed();
 }

google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/GrpcDatastoreRpc.java

@@ -63,7 +63,7 @@
 import java.util.Collections;
 
 @InternalApi
-public class GrpcDatastoreRpc implements AutoCloseable, DatastoreRpc {
+public class GrpcDatastoreRpc implements DatastoreRpc {
 
   private final GrpcDatastoreStub datastoreStub;
   private final ClientContext clientContext;
@@ -146,6 +146,11 @@ public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryReques
     return datastoreStub.runAggregationQueryCallable().call(request);
   }
 
+  @Override
+  public boolean isClosed() {
+    return closed && datastoreStub.isShutdown();
+  }
+
   private boolean isEmulator(DatastoreOptions datastoreOptions) {
     return isLocalHost(datastoreOptions.getHost())
         || NoCredentials.getInstance().equals(datastoreOptions.getCredentials());

google-cloud-datastore/src/main/java/com/google/cloud/datastore/spi/v1/HttpDatastoreRpc.java

@@ -211,4 +211,14 @@ public RunAggregationQueryResponse runAggregationQuery(RunAggregationQueryReques
       throw translate(ex);
     }
   }
+
+  @Override
+  public void close() throws Exception {
+    throw new UnsupportedOperationException("close() is not supported");
+  }
+
+  @Override
+  public boolean isClosed() {
+    throw new UnsupportedOperationException("isClosed() is not supported");
+  }
 }

google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/LocalDatastoreHelper.java

@@ -57,12 +57,12 @@ public class LocalDatastoreHelper extends BaseEmulatorHelper<DatastoreOptions> {
   private static final String GCLOUD_CMD_TEXT = "gcloud beta emulators datastore start";
   private static final String GCLOUD_CMD_PORT_FLAG = "--host-port=";
   private static final String VERSION_PREFIX = "cloud-datastore-emulator ";
-  private static final String MIN_VERSION = "1.2.0";
+  private static final String MIN_VERSION = "2.0.2"; // latest version compatible with java 8
 
   // Downloadable emulator settings
   private static final String BIN_NAME = "cloud-datastore-emulator/cloud_datastore_emulator";
   private static final String FILENAME = "cloud-datastore-emulator-" + MIN_VERSION + ".zip";
-  private static final String MD5_CHECKSUM = "ec2237a0f0ac54964c6bd95e12c73720";
+  private static final String MD5_CHECKSUM = "e0d1170519cf52e2e5f9f93892cdf70c";
   private static final String BIN_CMD_PORT_FLAG = "--port=";
   private static final URL EMULATOR_URL;
   private static final String EMULATOR_URL_ENV_VAR = "DATASTORE_EMULATOR_URL";

google-cloud-datastore/src/test/java/com/google/cloud/datastore/DatastoreTest.java

@@ -31,6 +31,7 @@
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
@@ -39,7 +40,6 @@
 import com.google.cloud.datastore.Query.ResultType;
 import com.google.cloud.datastore.StructuredQuery.OrderBy;
 import com.google.cloud.datastore.StructuredQuery.PropertyFilter;
-import com.google.cloud.datastore.it.MultipleAttemptsRule;
 import com.google.cloud.datastore.spi.DatastoreRpcFactory;
 import com.google.cloud.datastore.spi.v1.DatastoreRpc;
 import com.google.cloud.datastore.testing.LocalDatastoreHelper;
@@ -79,29 +79,22 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.TimeoutException;
 import java.util.function.Predicate;
 import org.easymock.EasyMock;
 import org.junit.AfterClass;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.BeforeClass;
-import org.junit.ClassRule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.threeten.bp.Duration;
 
 @RunWith(JUnit4.class)
 public class DatastoreTest {
-  private static final int NUMBER_OF_ATTEMPTS = 5;
-
-  @ClassRule
-  public static MultipleAttemptsRule rr = new MultipleAttemptsRule(NUMBER_OF_ATTEMPTS, 10);
-
-  private static LocalDatastoreHelper helper = LocalDatastoreHelper.create(1.0);
-  private static final DatastoreOptions options = helper.getOptions();
-  private static final Datastore datastore = options.getService();
+  private static final LocalDatastoreHelper helper = LocalDatastoreHelper.create(1.0, 9090);
+  private static DatastoreOptions options = helper.getOptions();
+  private static Datastore datastore;
   private static final String PROJECT_ID = options.getProjectId();
   private static final String KIND1 = "kind1";
   private static final String KIND2 = "kind2";
@@ -177,6 +170,8 @@ public class DatastoreTest {
   @BeforeClass
   public static void beforeClass() throws IOException, InterruptedException {
     helper.start();
+    options = helper.getOptions();
+    datastore = options.getService();
   }
 
   @Before
@@ -197,7 +192,8 @@ public void setUp() {
   }
 
   @AfterClass
-  public static void afterClass() throws IOException, InterruptedException, TimeoutException {
+  public static void afterClass() throws Exception {
+    datastore.close();
     helper.stop(Duration.ofMinutes(1));
   }
 
@@ -1386,6 +1382,21 @@ public void testDatabaseIdKeyFactory() {
     checkKeyProperties(incompleteKey);
   }
 
+  @Test
+  public void testDatastoreClose() throws Exception {
+    Datastore datastore = options.toBuilder().build().getService();
+    Entity entity = datastore.get(KEY3);
+    assertNull(entity);
+
+    datastore.close();
+    assertTrue(datastore.isClosed());
+
+    assertThrows(
+        "io.grpc.StatusRuntimeException: UNAVAILABLE: Channel shutdown invoked",
+        DatastoreException.class,
+        () -> datastore.get(KEY3));
+  }
+
   private void checkKeyProperties(BaseKey key) {
     assertEquals(options.getDatabaseId(), key.getDatabaseId());
     assertEquals(options.getProjectId(), key.getProjectId());

google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreAggregationsTest.java

@@ -43,6 +43,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import org.junit.After;
+import org.junit.AfterClass;
 import org.junit.Test;
 
 // TODO(jainsahab) Move all the aggregation related tests from ITDatastoreTest to this file
@@ -63,6 +64,11 @@ public void tearDown() {
     DATASTORE.delete(keysToDelete);
   }
 
+  @AfterClass
+  public static void afterClass() throws Exception {
+    DATASTORE.close();
+  }
+
   Key key1 = DATASTORE.newKeyFactory().setKind(KIND).newKey(1);
   Key key2 = DATASTORE.newKeyFactory().setKind(KIND).newKey(2);
   Key key3 = DATASTORE.newKeyFactory().setKind(KIND).newKey(3);

google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreConceptsTest.java

@@ -67,7 +67,9 @@
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import org.junit.After;
+import org.junit.AfterClass;
 import org.junit.Before;
+import org.junit.BeforeClass;
 import org.junit.Test;
 
 /*
@@ -77,7 +79,7 @@ public class ITDatastoreConceptsTest {
   private static final RemoteDatastoreHelper HELPER = RemoteDatastoreHelper.create();
   private static final DatastoreOptions OPTIONS = HELPER.getOptions();
   private static final FullEntity<IncompleteKey> TEST_FULL_ENTITY = FullEntity.newBuilder().build();
-  private Datastore datastore;
+  private static Datastore datastore;
   private KeyFactory keyFactory;
   private Key taskKey;
   private Entity testEntity;
@@ -87,13 +89,15 @@ public class ITDatastoreConceptsTest {
 
   private static final String TASK_CONCEPTS = "TaskConcepts";
 
-  /**
-   * Initializes Datastore and cleans out any residual values. Also initializes global variables
-   * used for testing.
-   */
+  /** Initializes Datastore for testing. */
+  @BeforeClass
+  public static void beforeClass() throws Exception {
+    datastore = OPTIONS.getService();
+  }
+
+  /** Cleans out any residual values. Also initializes global variables used for testing. */
   @Before
   public void setUp() {
-    datastore = OPTIONS.getService();
     StructuredQuery<Key> query = Query.newKeyQueryBuilder().build();
     QueryResults<Key> result = datastore.run(query);
     datastore.delete(Iterators.toArray(result, Key.class));
@@ -128,6 +132,11 @@ public void tearDown() {
     datastore.delete(taskKeysToDelete);
   }
 
+  @AfterClass
+  public static void afterClass() throws Exception {
+    datastore.close();
+  }
+
   private void assertValidKey(Key taskKey) {
     datastore.put(Entity.newBuilder(taskKey, TEST_FULL_ENTITY).build());
   }

google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITDatastoreTest.java

@@ -144,8 +144,10 @@ public class ITDatastoreTest {
   @Rule public MultipleAttemptsRule multipleAttemptsRule = new MultipleAttemptsRule(3);
 
   @AfterClass
-  public static void afterClass() {
+  public static void afterClass() throws Exception {
     HELPER.deleteNamespace();
+    DATASTORE_1.close();
+    DATASTORE_2.close();
   }
 
   public ITDatastoreTest(