feat: add reset to ConnectorRegistry (#304)

GoogleCloudPlatform · Nov 14, 2023 · e6812d0 · e6812d0
1 parent 07ae5ee
commit e6812d0
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectorRegistry.java b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/ConnectorRegistry.java
@@ -37,4 +37,12 @@ public static void register(String name, ConnectorConfig config) {
   public static void close(String name) {
     InternalConnectorRegistry.INSTANCE.close(name);
   }
+
+  /**
+   * Resets the entire AlloyDB JDBC Connector. This will stop all background threads. The next
+   * attempt to open a connection or register a configuration will start a new ConnectorRegistry.
+   */
+  public static void reset() {
+    InternalConnectorRegistry.INSTANCE.resetInstance();
+  }
 }
diff --git a/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/InternalConnectorRegistry.java b/alloydb-jdbc-connector/src/main/java/com/google/cloud/alloydb/InternalConnectorRegistry.java
@@ -95,17 +95,22 @@ public void close(String name) {
   }
 
   /** Shutdown all connectors and remove the singleton instance. */
-  public void shutdown() {
+  private void shutdown() {
     this.unnamedConnectors.forEach((key, c) -> c.close());
     this.unnamedConnectors.clear();
     this.namedConnectors.forEach((key, c) -> c.close());
     this.namedConnectors.clear();
-    this.executor.shutdown();
   }
 
   @Override
   public void close() {
     shutdown();
+    this.executor.shutdown();
+  }
+
+  /** Calls shutdown on the singleton. */
+  public void resetInstance() {
+    shutdown();
   }
 
   private Connector getConnector(ConnectionConfig config) {