Allow to reuse an existing global refdb database when running tests

Database creation and especially scheme initialization can take very
long time, up to a minute based on my observation. By allowing to reuse
and existing database we significantly reduce test run time. Use
SPANNER_DATABASE environment variable to specify a database.

Example:

  bazelisk test \
  --test_env='GOOGLE_APPLICATION_CREDENTIALS=/path/to/the/key.json' \
  --test_env='SPANNER_INSTANCE=test-instance' \
  --test_env='SPANNER_DATABASE=test-global-refdb' \
  --test_tag_filters=spanner-refdb //...

Note that the specified database doesn't even need to exist. It will be
created if necessary. However, it will not be dropped when the test
run finishes, leaving it ready for a next test run.

Change-Id: Ib4a46ab3f817a5063a5023f9bc89ad6b5c1eb747
diff --git a/src/main/resources/Documentation/build.md b/src/main/resources/Documentation/build.md
index 5f33214..45b93cc 100644
--- a/src/main/resources/Documentation/build.md
+++ b/src/main/resources/Documentation/build.md
@@ -40,13 +40,25 @@
 ### Running tests using a real spanner service
 
 In this case we have to provide an GCP service account key and the instance name
-under which the test refdb will be created. This is done by passing two
-environment variables `GOOGLE_APPLICATION_CREDENTIALS` and `SPANNER_INSTANCE`:
+under which the test refdb will be found or created. This is done by passing two
+environment variables `GOOGLE_APPLICATION_CREDENTIALS` and `SPANNER_INSTANCE`.
+
+If `SPANNER_DATABASE` environment variable is also specified, then a database
+with that name will be used for testing. If a database with that name does not
+exist it will be created. This database will not be dropped when the test run
+finishes, leaving it ready for a next test run.
+Using this option is recommended as database creation and especially schema
+initialization is known to take a very long time.
+
+If `SPANNER_DATABASE` environment variable is not specified then a new database
+will be created for this test run and will be dropped when the test run
+finishes.
 
 ```
 bazelisk test \
   --test_env='GOOGLE_APPLICATION_CREDENTIALS=/path/to/the/key.json' \
   --test_env='SPANNER_INSTANCE=test-instance' \
+  [--test_env='SPANNER_DATABASE=test-global-refdb'] \
   --test_tag_filters=@PLUGIN@ //...
 ```
 
@@ -57,6 +69,7 @@
 bazelisk test \
   --test_env='GOOGLE_APPLICATION_CREDENTIALS=/path/to/the/key.json' \
   --test_env='SPANNER_INSTANCE=test-instance' \
+  [--test_env='SPANNER_DATABASE=test-global-refdb'] \
   //plugins/@PLUGIN@/...
 ```
 
diff --git a/src/test/java/com/googlesource/gerrit/plugins/spannerrefdb/RealSpannerRefDb.java b/src/test/java/com/googlesource/gerrit/plugins/spannerrefdb/RealSpannerRefDb.java
index 1712196..e09620e 100644
--- a/src/test/java/com/googlesource/gerrit/plugins/spannerrefdb/RealSpannerRefDb.java
+++ b/src/test/java/com/googlesource/gerrit/plugins/spannerrefdb/RealSpannerRefDb.java
@@ -14,11 +14,13 @@
 
 package com.googlesource.gerrit.plugins.spannerrefdb;
 
+import com.google.api.client.util.Strings;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.spanner.Database;
 import com.google.cloud.spanner.DatabaseAdminClient;
 import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.DatabaseId;
+import com.google.cloud.spanner.DatabaseNotFoundException;
 import com.google.cloud.spanner.KeySet;
 import com.google.cloud.spanner.Mutation;
 import com.google.cloud.spanner.SpannerOptions;
@@ -36,7 +38,12 @@
     if (INSTANCE == null) {
       String keyPath = System.getenv("GOOGLE_APPLICATION_CREDENTIALS");
       String instance = System.getenv("SPANNER_INSTANCE");
-      INSTANCE = new RealSpannerRefDb(keyPath, instance);
+      String database = System.getenv("SPANNER_DATABASE");
+      if (Strings.isNullOrEmpty(database)) {
+        INSTANCE = new RealSpannerRefDb(keyPath, instance);
+      } else {
+        INSTANCE = new RealSpannerRefDb(keyPath, instance, database);
+      }
     }
     return INSTANCE;
   }
@@ -44,6 +51,7 @@
   private final String keyPath;
   private final String instance;
   private final String dbName;
+  private final boolean dropDbOnShutdown;
 
   private boolean dbInitialized;
   private SpannerRefDatabase refdb;
@@ -51,9 +59,19 @@
   private ScheduledThreadPoolExecutor heartbeatExecutor;
 
   private RealSpannerRefDb(String keyPath, String instance) {
+    this(keyPath, instance, "global-refdb-" + System.currentTimeMillis(), true);
+  }
+
+  private RealSpannerRefDb(String keyPath, String instance, String dbName) {
+    this(keyPath, instance, dbName, false);
+  }
+
+  private RealSpannerRefDb(
+      String keyPath, String instance, String dbName, boolean dropDbOnShutdown) {
     this.keyPath = keyPath;
     this.instance = instance;
-    this.dbName = "global-refdb-" + System.currentTimeMillis();
+    this.dbName = dbName;
+    this.dropDbOnShutdown = dropDbOnShutdown;
   }
 
   @Override
@@ -61,9 +79,8 @@
     if (!dbInitialized) {
       initialize();
       dbInitialized = true;
-    } else {
-      deleteAllData();
     }
+    deleteAllData();
   }
 
   @Override
@@ -87,22 +104,40 @@
   }
 
   private void initialize() throws Exception {
+
     GoogleCredentials credentials = GoogleCredentials.fromStream(new FileInputStream(keyPath));
 
     SpannerOptions options = SpannerOptions.newBuilder().setCredentials(credentials).build();
     DatabaseId dbId = DatabaseId.of(options.getProjectId(), instance, dbName);
     DatabaseAdminClient dbAdminClient = options.getService().getDatabaseAdminClient();
 
-    // create empty database
-    Database database =
-        dbAdminClient
-            .createDatabase(dbId.getInstanceId().getInstance(), dbId.getDatabase(), List.of())
-            .get();
+    Database database;
+    try {
+      database = dbAdminClient.getDatabase(dbId.getInstanceId().getInstance(), dbId.getDatabase());
+    } catch (DatabaseNotFoundException e) {
+      // create new database and schema
+      database =
+          dbAdminClient
+              .createDatabase(dbId.getInstanceId().getInstance(), dbId.getDatabase(), List.of())
+              .get();
+      DatabaseSchemaCreator databaseSchemaCreator = new DatabaseSchemaCreator(dbAdminClient, dbId);
+      databaseSchemaCreator.start();
+    }
 
-    DatabaseSchemaCreator databaseSchemaCreator = new DatabaseSchemaCreator(dbAdminClient, dbId);
-    databaseSchemaCreator.start();
+    if (dropDbOnShutdown) {
+      Database toBeDropped = database;
+      Runtime.getRuntime()
+          .addShutdownHook(
+              new Thread(
+                  () -> {
+                    toBeDropped.drop();
+                  }));
+    }
+
     dbClient = options.getService().getDatabaseClient(dbId);
-    heartbeatExecutor = new ScheduledThreadPoolExecutor(2);
+
+    heartbeatExecutor = createHeartbeatExecutor();
+
     Lock.Factory lockFactory =
         new Lock.Factory() {
           @Override
@@ -111,8 +146,17 @@
           }
         };
     refdb = new SpannerRefDatabase(dbClient, lockFactory);
+  }
 
-    // schedule heartbeat stop and database drop when JVM gets shutdown
+  private void deleteAllData() {
+    List<Mutation> mutations = new ArrayList<>();
+    mutations.add(Mutation.delete("refs", KeySet.all()));
+    mutations.add(Mutation.delete("locks", KeySet.all()));
+    dbClient.write(mutations);
+  }
+
+  private ScheduledThreadPoolExecutor createHeartbeatExecutor() {
+    ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(2);
     Runtime.getRuntime()
         .addShutdownHook(
             new Thread(
@@ -123,14 +167,7 @@
                   } catch (InterruptedException e) {
                     throw new RuntimeException(e);
                   }
-                  database.drop();
                 }));
-  }
-
-  private void deleteAllData() {
-    List<Mutation> mutations = new ArrayList<>();
-    mutations.add(Mutation.delete("refs", KeySet.all()));
-    mutations.add(Mutation.delete("locks", KeySet.all()));
-    dbClient.write(mutations);
+    return executor;
   }
 }