Repository: firebase/firebase-admin-java
Pull request #723 (2022-09-27 18:42:55)

Implement MultiDB Support


Patch:
diff --git a/src/main/java/com/google/firebase/cloud/FirestoreClient.java b/src/main/java/com/google/firebase/cloud/FirestoreClient.java
index fddf332..1012017 100644
--- a/src/main/java/com/google/firebase/cloud/FirestoreClient.java
+++ b/src/main/java/com/google/firebase/cloud/FirestoreClient.java
@@ -11,6 +11,9 @@ import com.google.firebase.FirebaseApp;
 import com.google.firebase.ImplFirebaseTrampolines;
 import com.google.firebase.internal.FirebaseService;
 import com.google.firebase.internal.NonNull;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -32,7 +35,7 @@ public class FirestoreClient {
 
   private final Firestore firestore;
 
-  private FirestoreClient(FirebaseApp app) {
+  private FirestoreClient(FirebaseApp app, String databaseId) {
     checkNotNull(app, "FirebaseApp must not be null");
     String projectId = ImplFirebaseTrampolines.getProjectId(app);
     checkArgument(!Strings.isNullOrEmpty(projectId),
@@ -47,13 +50,14 @@ public class FirestoreClient {
         .setCredentialsProvider(
             FixedCredentialsProvider.create(ImplFirebaseTrampolines.getCredentials(app)))
         .setProjectId(projectId)
+        .setDatabaseId(databaseId)
         .build()
         .getService();
   }
 
   /**
-   * Returns the Firestore instance associated with the default Firebase app. Returns the same
-   * instance for all invocations. The Firestore instance and all references obtained from it
+   * Returns the default Firestore instance associated with the default Firebase app. Returns the
+   * same instance for all invocations. The Firestore instance and all references obtained from it
    * becomes unusable, once the default app is deleted.
    *
    * @return A non-null <a href="https://googlecloudplatform.github.io/google-cloud-java/google-cloud-clients/apidocs/com/google/cloud/firestore/Firestore.html">{@code Firestore}</a>
@@ -65,9 +69,9 @@ public class FirestoreClient {
   }
 
   /**
-   * Returns the Firestore instance associated with the specified Firebase app. For a given app,
-   * always returns the same instance. The Firestore instance and all references obtained from it
-   * becomes unusable, once the specified app is deleted.
+   * Returns the default Firestore instance associated with the specified Firebase app. For a given
+   * app, invocation always returns the same instance. The Firestore instance and all references
+   * obtained from it becomes unusable, once the specified app is deleted.
    *
    * @param app A non-null {@link FirebaseApp}.
    * @return A non-null <a href="https://googlecloudplatform.github.io/google-cloud-java/google-cloud-clients/apidocs/com/google/cloud/firestore/Firestore.html">{@code Firestore}</a>
@@ -75,32 +79,86 @@ public class FirestoreClient {
    */
   @NonNull
   public static Firestore getFirestore(FirebaseApp app) {
-    return getInstance(app).firestore;
+    return getFirestore(app, ImplFirebaseTrampolines.getFirestoreOptions(app).getDatabaseId());
+  }
+
+  /**
+   * Returns the Firestore instance associated with the specified Firebase app. Returns the same
+   * instance for all invocations given the same app and database parameter. The Firestore instance
+   * and all references obtained from it becomes unusable, once the specified app is deleted.
+   *
+   * @param app      A non-null {@link FirebaseApp}.
+   * @param database - The name of database.
+   * @return A non-null <a href="https://googlecloudplatform.github.io/google-cloud-java/google-cloud-clients/apidocs/com/google/cloud/firestore/Firestore.html">{@code Firestore}</a>
+   *     instance.
+   */
+  @NonNull
+  static Firestore getFirestore(FirebaseApp app, String database) {
+    return getInstance(app, database).firestore;
+  }
+
+  /**
+   * Returns the Firestore instance associated with the default Firebase app. Returns the same
+   * instance for all invocations given the same database parameter. The Firestore instance and all
+   * references obtained from it becomes unusable, once the default app is deleted.
+   *
+   * @param database - The name of database.
+   * @return A non-null <a href="https://googlecloudplatform.github.io/google-cloud-java/google-cloud-clients/apidocs/com/google/cloud/firestore/Firestore.html">{@code Firestore}</a>
+   *     instance.
+   */
+  @NonNull
+  static Firestore getFirestore(String database) {
+    return getFirestore(FirebaseApp.getInstance(), database);
   }
 
-  private static synchronized FirestoreClient getInstance(FirebaseApp app) {
+  private static synchronized FirestoreClient getInstance(FirebaseApp app, String database) {
     FirestoreClientService service = ImplFirebaseTrampolines.getService(app,
         SERVICE_ID, FirestoreClientService.class);
     if (service == null) {
       service = ImplFirebaseTrampolines.addService(app, new FirestoreClientService(app));
     }
-    return service.getInstance();
+    return service.getInstance().get(database);
   }
 
   private static final String SERVICE_ID = FirestoreClient.class.getName();
 
-  private static class FirestoreClientService extends FirebaseService<FirestoreClient> {
+  private static class FirestoreClientService extends FirebaseService<FirestoreInstances> {
 
     FirestoreClientService(FirebaseApp app) {
-      super(SERVICE_ID, new FirestoreClient(app));
+      super(SERVICE_ID, new FirestoreInstances(app));
     }
 
     @Override
     public void destroy() {
-      try {
-        instance.firestore.close();
-      } catch (Exception e) {
-        logger.warn("Error while closing the Firestore instance", e);
+      instance.destroy();
+    }
+  }
+
+  private static class FirestoreInstances {
+
+    private final FirebaseApp app;
+
+    private final Map<String, FirestoreClient> clients =
+        Collections.synchronizedMap(new HashMap<>());
+
+    private FirestoreInstances(FirebaseApp app) {
+      this.app = app;
+    }
+
+    FirestoreClient get(String databaseId) {
+      return clients.computeIfAbsent(databaseId, id -> new FirestoreClient(app, id));
+    }
+
+    void destroy() {
+      synchronized (clients) {
+        for (FirestoreClient client : clients.values()) {
+          try {
+            client.firestore.close();
+          } catch (Exception e) {
+            logger.warn("Error while closing the Firestore instance", e);
+          }
+        }
+        clients.clear();
       }
     }
   }


Test patch:
diff --git a/src/test/java/com/google/firebase/cloud/FirestoreClientTest.java b/src/test/java/com/google/firebase/cloud/FirestoreClientTest.java
index 018d0e4..a5f617c 100644
--- a/src/test/java/com/google/firebase/cloud/FirestoreClientTest.java
+++ b/src/test/java/com/google/firebase/cloud/FirestoreClientTest.java
@@ -2,9 +2,9 @@ package com.google.firebase.cloud;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertSame;
-import static org.junit.Assert.assertTrue;
-import static org.junit.Assert.fail;
+import static org.junit.Assert.assertThrows;
 
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.firestore.DocumentReference;
@@ -26,6 +26,7 @@ public class FirestoreClientTest {
       // Setting credentials is not required (they get overridden by Admin SDK), but without
       // this Firestore logs an ugly warning during tests.
       .setCredentials(new MockGoogleCredentials("test-token"))
+      .setDatabaseId("differedDefaultDatabaseId")
       .build();
 
   @After
@@ -35,47 +36,75 @@ public class FirestoreClientTest {
 
   @Test
   public void testExplicitProjectId() throws IOException {
+    final String databaseId = "databaseIdInTestExplicitProjectId";
     FirebaseApp app = FirebaseApp.initializeApp(FirebaseOptions.builder()
         .setCredentials(GoogleCredentials.fromStream(ServiceAccount.EDITOR.asStream()))
         .setProjectId("explicit-project-id")
         .setFirestoreOptions(FIRESTORE_OPTIONS)
         .build());
-    Firestore firestore = FirestoreClient.getFirestore(app);
-    assertEquals("explicit-project-id", firestore.getOptions().getProjectId());
+    Firestore firestore1 = FirestoreClient.getFirestore(app);
+    assertEquals("explicit-project-id", firestore1.getOptions().getProjectId());
+    assertEquals(FIRESTORE_OPTIONS.getDatabaseId(), firestore1.getOptions().getDatabaseId());
 
-    firestore = FirestoreClient.getFirestore();
-    assertEquals("explicit-project-id", firestore.getOptions().getProjectId());
+    assertSame(firestore1, FirestoreClient.getFirestore());
+
+    Firestore firestore2 = FirestoreClient.getFirestore(app, databaseId);
+    assertEquals("explicit-project-id", firestore2.getOptions().getProjectId());
+    assertEquals(databaseId, firestore2.getOptions().getDatabaseId());
+
+    assertSame(firestore2, FirestoreClient.getFirestore(databaseId));
+
+    assertNotSame(firestore1, firestore2);
   }
 
   @Test
   public void testServiceAccountProjectId() throws IOException {
+    final String databaseId = "databaseIdInTestServiceAccountProjectId";
     FirebaseApp app = FirebaseApp.initializeApp(FirebaseOptions.builder()
         .setCredentials(GoogleCredentials.fromStream(ServiceAccount.EDITOR.asStream()))
         .setFirestoreOptions(FIRESTORE_OPTIONS)
         .build());
-    Firestore firestore = FirestoreClient.getFirestore(app);
-    assertEquals("mock-project-id", firestore.getOptions().getProjectId());
+    Firestore firestore1 = FirestoreClient.getFirestore(app);
+    assertEquals("mock-project-id", firestore1.getOptions().getProjectId());
+    assertEquals(FIRESTORE_OPTIONS.getDatabaseId(), firestore1.getOptions().getDatabaseId());
+
+    assertSame(firestore1, FirestoreClient.getFirestore());
 
-    firestore = FirestoreClient.getFirestore();
-    assertEquals("mock-project-id", firestore.getOptions().getProjectId());
+    Firestore firestore2 = FirestoreClient.getFirestore(app, databaseId);
+    assertEquals("mock-project-id", firestore2.getOptions().getProjectId());
+    assertEquals(databaseId, firestore2.getOptions().getDatabaseId());
+
+    assertSame(firestore2, FirestoreClient.getFirestore(databaseId));
+
+    assertNotSame(firestore1, firestore2);
   }
 
   @Test
   public void testFirestoreOptions() throws IOException {
+    final String databaseId = "databaseIdInTestFirestoreOptions";
     FirebaseApp app = FirebaseApp.initializeApp(FirebaseOptions.builder()
         .setCredentials(GoogleCredentials.fromStream(ServiceAccount.EDITOR.asStream()))
         .setProjectId("explicit-project-id")
         .setFirestoreOptions(FIRESTORE_OPTIONS)
         .build());
-    Firestore firestore = FirestoreClient.getFirestore(app);
-    assertEquals("explicit-project-id", firestore.getOptions().getProjectId());
+    Firestore firestore1 = FirestoreClient.getFirestore(app);
+    assertEquals("explicit-project-id", firestore1.getOptions().getProjectId());
+    assertEquals(FIRESTORE_OPTIONS.getDatabaseId(), firestore1.getOptions().getDatabaseId());
+
+    assertSame(firestore1, FirestoreClient.getFirestore());
+
+    Firestore firestore2 = FirestoreClient.getFirestore(app, databaseId);
+    assertEquals("explicit-project-id", firestore2.getOptions().getProjectId());
+    assertEquals(databaseId, firestore2.getOptions().getDatabaseId());
 
-    firestore = FirestoreClient.getFirestore();
-    assertEquals("explicit-project-id", firestore.getOptions().getProjectId());
+    assertSame(firestore2, FirestoreClient.getFirestore(databaseId));
+
+    assertNotSame(firestore1, firestore2);
   }
 
   @Test
   public void testFirestoreOptionsOverride() throws IOException {
+    final String databaseId = "databaseIdInTestFirestoreOptions";
     FirebaseApp app = FirebaseApp.initializeApp(FirebaseOptions.builder()
         .setCredentials(GoogleCredentials.fromStream(ServiceAccount.EDITOR.asStream()))
         .setProjectId("explicit-project-id")
@@ -84,48 +113,51 @@ public class FirestoreClientTest {
             .setCredentials(GoogleCredentials.fromStream(ServiceAccount.EDITOR.asStream()))
             .build())
         .build());
-    Firestore firestore = FirestoreClient.getFirestore(app);
-    assertEquals("explicit-project-id", firestore.getOptions().getProjectId());
+    Firestore firestore1 = FirestoreClient.getFirestore(app);
+    assertEquals("explicit-project-id", firestore1.getOptions().getProjectId());
     assertSame(ImplFirebaseTrampolines.getCredentials(app),
-        firestore.getOptions().getCredentialsProvider().getCredentials());
+        firestore1.getOptions().getCredentialsProvider().getCredentials());
+    assertEquals("(default)", firestore1.getOptions().getDatabaseId());
+
+    assertSame(firestore1, FirestoreClient.getFirestore());
 
-    firestore = FirestoreClient.getFirestore();
-    assertEquals("explicit-project-id", firestore.getOptions().getProjectId());
+    Firestore firestore2 = FirestoreClient.getFirestore(app, databaseId);
+    assertEquals("explicit-project-id", firestore2.getOptions().getProjectId());
     assertSame(ImplFirebaseTrampolines.getCredentials(app),
-        firestore.getOptions().getCredentialsProvider().getCredentials());
+        firestore2.getOptions().getCredentialsProvider().getCredentials());
+    assertEquals(databaseId, firestore2.getOptions().getDatabaseId());
+
+    assertSame(firestore2, FirestoreClient.getFirestore(databaseId));
+
+    assertNotSame(firestore1, firestore2);
   }
 
   @Test
   public void testAppDelete() throws IOException {
+    final String databaseId = "databaseIdInTestAppDelete";
     FirebaseApp app = FirebaseApp.initializeApp(FirebaseOptions.builder()
         .setCredentials(GoogleCredentials.fromStream(ServiceAccount.EDITOR.asStream()))
         .setProjectId("mock-project-id")
         .setFirestoreOptions(FIRESTORE_OPTIONS)
         .build());
 
-    Firestore firestore = FirestoreClient.getFirestore(app);
-    assertNotNull(firestore);
-    DocumentReference document = firestore.collection("collection").document("doc");
+    Firestore firestore1 = FirestoreClient.getFirestore(app);
+    assertNotNull(firestore1);
+    assertSame(firestore1, FirestoreClient.getFirestore());
+
+    Firestore firestore2 = FirestoreClient.getFirestore(app, databaseId);
+    assertNotNull(firestore2);
+    assertSame(firestore2, FirestoreClient.getFirestore(databaseId));
+
+    assertNotSame(firestore1, firestore2);
+
+    DocumentReference document = firestore1.collection("collection").document("doc");
     app.delete();
-    try {
-      FirestoreClient.getFirestore(app);
-      fail("No error thrown for deleted app");
-    } catch (IllegalStateException expected) {
-      // ignore
-    }
-
-    try {
-      document.get();
-      fail("No error thrown for deleted app");
-    } catch (IllegalStateException expected) {
-      // ignore
-    }
-
-    try {
-      FirestoreClient.getFirestore();
-      fail("No error thrown for deleted app");
-    } catch (IllegalStateException expected) {
-      // ignore
-    }
+
+    assertThrows(IllegalStateException.class, () -> FirestoreClient.getFirestore(app));
+    assertThrows(IllegalStateException.class, () -> document.get());
+    assertThrows(IllegalStateException.class, () -> FirestoreClient.getFirestore());
+    assertThrows(IllegalStateException.class, () -> FirestoreClient.getFirestore(app, databaseId));
+    assertThrows(IllegalStateException.class, () -> FirestoreClient.getFirestore(databaseId));
   }
 }