Allow users to configure wait options for new Neo4j databases (#31129)

* Allow users to configure wait options for new Neo4j databases

Since v5 of Neo4j, databases are created in an asynchronous manner.
Before this commit, `Neo4jResourceManager` did not allow users
to specify whether to wait or not for the database creation (via
the corresponding Cypher syntax element).

This commit fixes that, while keeping the behavior the same as
before, by default (which is to not wait).

* Format

* Fold back intermediate classes

* Format again

it/neo4j/src/main/java/org/apache/beam/it/neo4j/DatabaseWaitOption.java

@@ -0,0 +1,20 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+package org.apache.beam.it.neo4j;
+
+public interface DatabaseWaitOption {}

it/neo4j/src/main/java/org/apache/beam/it/neo4j/DatabaseWaitOptions.java

@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+package org.apache.beam.it.neo4j;
+
+public class DatabaseWaitOptions {
+
+  public static DatabaseWaitOption waitDatabase() {
+    return DatabaseWait.WAIT;
+  }
+
+  public static DatabaseWaitOption waitDatabase(int seconds) {
+    return new DatabaseWaitInSeconds(seconds);
+  }
+
+  public static DatabaseWaitOption noWaitDatabase() {
+    return DatabaseNoWait.NO_WAIT;
+  }
+
+  static String asCypher(DatabaseWaitOption option) {
+    if (option == null || option == DatabaseNoWait.NO_WAIT) {
+      return "NOWAIT";
+    }
+    if (option == DatabaseWait.WAIT) {
+      return "WAIT";
+    }
+    if (option instanceof DatabaseWaitInSeconds) {
+      DatabaseWaitInSeconds wait = (DatabaseWaitInSeconds) option;
+      return String.format("WAIT %s SECONDS", wait.getSeconds());
+    }
+    throw new Neo4jResourceManagerException(
+        String.format("Unsupported wait option type %s", option.getClass()));
+  }
+
+  private enum DatabaseNoWait implements DatabaseWaitOption {
+    NO_WAIT;
+  }
+
+  private enum DatabaseWait implements DatabaseWaitOption {
+    WAIT;
+  }
+
+  private static class DatabaseWaitInSeconds implements DatabaseWaitOption {
+    private final int seconds;
+
+    public DatabaseWaitInSeconds(int seconds) {
+      this.seconds = seconds;
+    }
+
+    public int getSeconds() {
+      return seconds;
+    }
+  }
+}

it/neo4j/src/main/java/org/apache/beam/it/neo4j/Neo4jResourceManager.java

@@ -61,6 +61,7 @@ public class Neo4jResourceManager extends TestContainerResourceManager<Neo4jCont
 
   private final Driver neo4jDriver;
   private final String databaseName;
+  private final DatabaseWaitOption waitOption;
   private final String connectionString;
   private final boolean usingStaticDatabase;
 
@@ -92,9 +93,11 @@ private Neo4jResourceManager(Builder builder) {
     this.usingStaticDatabase = builder.databaseName != null;
     if (usingStaticDatabase) {
       this.databaseName = builder.databaseName;
+      this.waitOption = null;
     } else {
       this.databaseName = generateDatabaseName(builder.testId);
-      createDatabase(databaseName);
+      this.waitOption = builder.waitOption;
+      createDatabase(databaseName, waitOption);
     }
   }
 
@@ -138,7 +141,7 @@ public synchronized void cleanupAll() {
     // First, delete the database if it was not given as a static argument
     try {
       if (!usingStaticDatabase) {
-        dropDatabase(databaseName);
+        dropDatabase(databaseName, waitOption);
       }
     } catch (Exception e) {
       LOG.error("Failed to delete Neo4j database {}.", databaseName, e);
@@ -164,21 +167,25 @@ public synchronized void cleanupAll() {
     LOG.info("Neo4j manager successfully cleaned up.");
   }
 
-  private void createDatabase(String databaseName) {
+  private void createDatabase(String databaseName, DatabaseWaitOption waitOption) {
     try (Session session =
         neo4jDriver.session(SessionConfig.builder().withDatabase("system").build())) {
-      session.run("CREATE DATABASE $db", Collections.singletonMap("db", databaseName)).consume();
+      String query =
+          String.format("CREATE DATABASE $db %s", DatabaseWaitOptions.asCypher(waitOption));
+      session.run(query, Collections.singletonMap("db", databaseName)).consume();
     } catch (Exception e) {
       throw new Neo4jResourceManagerException(
           String.format("Error dropping database %s.", databaseName), e);
     }
   }
 
   @VisibleForTesting
-  void dropDatabase(String databaseName) {
+  void dropDatabase(String databaseName, DatabaseWaitOption waitOption) {
     try (Session session =
         neo4jDriver.session(SessionConfig.builder().withDatabase("system").build())) {
-      session.run("DROP DATABASE $db", Collections.singletonMap("db", databaseName)).consume();
+      String query =
+          String.format("DROP DATABASE $db %s", DatabaseWaitOptions.asCypher(waitOption));
+      session.run(query, Collections.singletonMap("db", databaseName)).consume();
     } catch (Exception e) {
       throw new Neo4jResourceManagerException(
           String.format("Error dropping database %s.", databaseName), e);
@@ -194,6 +201,7 @@ public static final class Builder
       extends TestContainerResourceManager.Builder<Neo4jResourceManager> {
 
     private @Nullable String databaseName;
+    private @Nullable DatabaseWaitOption waitOption;
 
     private String adminPassword;
 
@@ -203,6 +211,7 @@ private Builder(String testId) {
       super(testId, DEFAULT_NEO4J_CONTAINER_NAME, DEFAULT_NEO4J_CONTAINER_TAG);
       this.adminPassword = generatePassword(4, 10, 2, 2, 0, Collections.emptyList());
       this.databaseName = null;
+      this.waitOption = null;
       this.driver = null;
     }
 
@@ -219,7 +228,27 @@ private Builder(String testId) {
      * @return this builder object with the database name set.
      */
     public Builder setDatabaseName(String databaseName) {
+      return setDatabaseName(databaseName, DatabaseWaitOptions.noWaitDatabase());
+    }
+
+    /**
+     * Sets the database name to that of a static database instance and sets the wait policy. Use
+     * this method only when attempting to operate on a pre-existing Neo4j database.
+     *
+     * <p>Note: if a database name is set, and a static Neo4j server is being used
+     * (useStaticContainer() is also called on the builder), then a database will be created on the
+     * static server if it does not exist, and it will not be removed when cleanupAll() is called on
+     * the Neo4jResourceManager.
+     *
+     * <p>{@link DatabaseWaitOptions} exposes all configurable wait options
+     *
+     * @param databaseName The database name.
+     * @param waitOption The database wait policy.
+     * @return this builder object with the database name set.
+     */
+    public Builder setDatabaseName(String databaseName, DatabaseWaitOption waitOption) {
       this.databaseName = databaseName;
+      this.waitOption = waitOption;
       return this;
     }
 

it/neo4j/src/test/java/org/apache/beam/it/neo4j/Neo4jResourceManagerIT.java

@@ -42,7 +42,7 @@ public class Neo4jResourceManagerIT {
   public void setUp() {
     neo4jResourceManager =
         Neo4jResourceManager.builder("placeholder")
-            .setDatabaseName("neo4j")
+            .setDatabaseName("neo4j", DatabaseWaitOptions.waitDatabase())
             .setAdminPassword("password")
             .build();
   }

it/neo4j/src/test/java/org/apache/beam/it/neo4j/Neo4jResourceManagerTest.java

@@ -19,9 +19,11 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
+import static org.mockito.AdditionalMatchers.and;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyMap;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.endsWith;
 import static org.mockito.ArgumentMatchers.startsWith;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
@@ -88,6 +90,16 @@ public void testCreateResourceManagerBuilderReturnsDefaultNeo4jResourceManager()
         .isInstanceOf(Neo4jResourceManager.class);
   }
 
+  @Test
+  public void testDatabaseIsCreatedWithNoWaitOptions() {
+    Neo4jResourceManager.Builder builder =
+        Neo4jResourceManager.builder(TEST_ID)
+            .setDatabaseName(STATIC_DATABASE_NAME, DatabaseWaitOptions.noWaitDatabase());
+    new Neo4jResourceManager(neo4jDriver, container, builder);
+
+    verify(session).run(and(startsWith("CREATE DATABASE"), endsWith("NOWAIT")), anyMap());
+  }
+
   @Test
   public void testGetUriShouldReturnCorrectValue() {
     assertThat(testManager.getUri()).matches("neo4j://" + HOST + ":" + MAPPED_PORT);
@@ -103,7 +115,8 @@ public void testDropDatabaseShouldThrowErrorIfDriverFailsToRunQuery() {
     doThrow(ClientException.class).when(session).run(anyString(), anyMap());
 
     assertThrows(
-        Neo4jResourceManagerException.class, () -> testManager.dropDatabase(STATIC_DATABASE_NAME));
+        Neo4jResourceManagerException.class,
+        () -> testManager.dropDatabase(STATIC_DATABASE_NAME, DatabaseWaitOptions.noWaitDatabase()));
   }
 
   @Test