From 984a7b3136d306ca49117e0a44e50e0f85dfb9f7 Mon Sep 17 00:00:00 2001
From: Alex Viscreanu <alexviscreanu@gmail.com>
Date: Tue, 11 Jun 2024 22:23:57 +0200
Subject: [PATCH] feat(app): allow configuring firestore instance id

---
 firebase.go      | 10 +++++++++-
 firebase_test.go | 44 ++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/firebase.go b/firebase.go
index 47a1006e..ed300592 100644
--- a/firebase.go
+++ b/firebase.go
@@ -51,6 +51,7 @@ type App struct {
 	projectID        string
 	serviceAccountID string
 	storageBucket    string
+	firestoreID      string
 	opts             []option.ClientOption
 }
 
@@ -61,6 +62,7 @@ type Config struct {
 	ProjectID        string                  `json:"projectId"`
 	ServiceAccountID string                  `json:"serviceAccountId"`
 	StorageBucket    string                  `json:"storageBucket"`
+	FirestoreID      string                  `json:"firestoreId"`
 }
 
 // Auth returns an instance of auth.Client.
@@ -107,7 +109,7 @@ func (a *App) Firestore(ctx context.Context) (*firestore.Client, error) {
 	if a.projectID == "" {
 		return nil, errors.New("project id is required to access Firestore")
 	}
-	return firestore.NewClient(ctx, a.projectID, a.opts...)
+	return firestore.NewClientWithDatabase(ctx, a.projectID, a.firestoreID, a.opts...)
 }
 
 // InstanceID returns an instance of iid.Client.
@@ -161,12 +163,18 @@ func NewApp(ctx context.Context, config *Config, opts ...option.ClientOption) (*
 		ao = *config.AuthOverride
 	}
 
+	fid := firestore.DefaultDatabaseID
+	if config.FirestoreID != "" {
+		fid = config.FirestoreID
+	}
+
 	return &App{
 		authOverride:     ao,
 		dbURL:            config.DatabaseURL,
 		projectID:        pid,
 		serviceAccountID: config.ServiceAccountID,
 		storageBucket:    config.StorageBucket,
+		firestoreID:      fid,
 		opts:             o,
 	}, nil
 }
diff --git a/firebase_test.go b/firebase_test.go
index 7830225e..4e503f09 100644
--- a/firebase_test.go
+++ b/firebase_test.go
@@ -29,6 +29,7 @@ import (
 	"testing"
 	"time"
 
+	"cloud.google.com/go/firestore"
 	"firebase.google.com/go/v4/messaging"
 	"golang.org/x/oauth2"
 	"golang.org/x/oauth2/google"
@@ -282,8 +283,20 @@ func TestFirestore(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	if c, err := app.Firestore(ctx); c == nil || err != nil {
-		t.Errorf("Firestore() = (%v, %v); want (auth, nil)", c, err)
+	c, err := app.Firestore(ctx)
+	if c == nil || err != nil {
+		t.Fatalf("Firestore() = (%v, %v); want (auth, nil)", c, err)
+	}
+
+	// verify that the client is using the overridden database ID
+
+	doc := c.Doc("foo/bar")
+	if doc == nil {
+		t.Fatalf("Doc() = nil; want doc")
+	}
+
+	if !strings.Contains(doc.Path, firestore.DefaultDatabaseID) {
+		t.Fatalf("Doc().Path = %q; want to contain %q", doc.Path, firestore.DefaultDatabaseID)
 	}
 }
 
@@ -338,6 +351,33 @@ func TestFirestoreWithNoProjectID(t *testing.T) {
 	}
 }
 
+func TestFirestoreWithDatabaseID(t *testing.T) {
+	firestoreid := "my-awesome-firestore-db"
+
+	ctx := context.Background()
+	config := &Config{FirestoreID: firestoreid}
+	app, err := NewApp(ctx, config, option.WithCredentialsFile("testdata/service_account.json"))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	c, err := app.Firestore(ctx)
+	if c == nil || err != nil {
+		t.Fatalf("Firestore() = (%v, %v); want (auth, nil)", c, err)
+	}
+
+	// verify that the client is using the default database ID
+
+	doc := c.Doc("foo/bar")
+	if doc == nil {
+		t.Fatalf("Doc() = nil; want doc")
+	}
+
+	if !strings.Contains(doc.Path, firestoreid) {
+		t.Fatalf("Doc().Path = %q; want to contain %q", doc.Path, firestore.DefaultDatabaseID)
+	}
+}
+
 func TestInstanceID(t *testing.T) {
 	ctx := context.Background()
 	app, err := NewApp(ctx, nil, option.WithCredentialsFile("testdata/service_account.json"))