feat(firestore): Support for multiple databases (#5331)

* feat: adds Firestore multi-DB support * fix: per reviewer * fix: actually fixed this time * added integration test * added test skip for new database * feat(firestore): Multi DB support * feat(firestore): Handling metadata header error * feat(firestore): Removed mock server setup * feat(firestore): URL escaping for routing header * feat(firestore): Resolving merge conflicts * feat(firestore): Refactoring code * feat(firestore): Run tests on (default) by default --------- Co-authored-by: Baha Aiman <[email protected]> Co-authored-by: meredithslota <[email protected]>
googleapis · Sep 18, 2023 · 94d4b1b · 94d4b1b
1 parent e29ba0c
commit 94d4b1b
Show file tree

Hide file tree

Showing 5 changed files with 162 additions and 25 deletions.
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -128,6 +128,7 @@ project's service account.
 - `GCLOUD_TESTS_GOLANG_DATASTORE_DATABASES`: Comma separated list of developer's Datastore databases. If not provided, default database i.e. empty string is used.
 - `GCLOUD_TESTS_GOLANG_FIRESTORE_PROJECT_ID`: Developers Console project's ID
 (e.g. doorway-cliff-677) for the Firestore project.
+- `GCLOUD_TESTS_GOLANG_FIRESTORE_DATABASES` : Comma separated list of developer's Firestore databases. If not provided, default database is used.
 - `GCLOUD_TESTS_GOLANG_FIRESTORE_KEY`: The path to the JSON key file of the
 Firestore project's service account.
 - `GCLOUD_TESTS_API_KEY`: API key for using the Translate API created above.
@@ -228,6 +229,9 @@ export GCLOUD_TESTS_GOLANG_DATASTORE_DATABASES=your-database-1,your-database-2
 # Developers Console project's ID (e.g. doorway-cliff-677) for the Firestore project.
 export GCLOUD_TESTS_GOLANG_FIRESTORE_PROJECT_ID=your-firestore-project
 
+# Comma separated list of developer's Firestore databases. If not provided, default database is used.
+export GCLOUD_TESTS_GOLANG_FIRESTORE_DATABASES=your-database-1,your-database-2
+
 # The path to the JSON key file of the Firestore project's service account.
 export GCLOUD_TESTS_GOLANG_FIRESTORE_KEY=~/directory/your-firestore-project-abcd1234.json
 

diff --git a/firestore/client.go b/firestore/client.go
@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net/url"
 	"os"
 	"strings"
 	"time"
@@ -43,6 +44,18 @@ import (
 // the resource being operated on.
 const resourcePrefixHeader = "google-cloud-resource-prefix"
 
+// requestParamsHeader is routing header required to access named databases
+const reqParamsHeader = "x-goog-request-params"
+
+// reqParamsHeaderVal constructs header from dbPath
+// dbPath is of the form projects/{project_id}/databases/{database_id}
+func reqParamsHeaderVal(dbPath string) string {
+	splitPath := strings.Split(dbPath, "/")
+	projectID := splitPath[1]
+	databaseID := splitPath[3]
+	return fmt.Sprintf("project_id=%s&database_id=%s", url.QueryEscape(projectID), url.QueryEscape(databaseID))
+}
+
 // DetectProjectID is a sentinel value that instructs NewClient to detect the
 // project ID. It is given in place of the projectID argument. NewClient will
 // use the project ID from the given credentials or the default credentials
@@ -52,6 +65,9 @@ const resourcePrefixHeader = "google-cloud-resource-prefix"
 // does not have the project ID encoded.
 const DetectProjectID = "*detect-project-id*"
 
+// DefaultDatabaseID is name of the default database
+const DefaultDatabaseID = "(default)"
+
 // A Client provides access to the Firestore service.
 type Client struct {
 	c            *vkit.Client
@@ -98,12 +114,28 @@ func NewClient(ctx context.Context, projectID string, opts ...option.ClientOptio
 	c := &Client{
 		c:            vc,
 		projectID:    projectID,
-		databaseID:   "(default)", // always "(default)", for now
+		databaseID:   DefaultDatabaseID,
 		readSettings: &readSettings{},
 	}
 	return c, nil
 }
 
+// NewClientWithDatabase creates a new Firestore client that accesses the
+// specified database.
+func NewClientWithDatabase(ctx context.Context, projectID string, databaseID string, opts ...option.ClientOption) (*Client, error) {
+	if databaseID == "" {
+		return nil, fmt.Errorf("firestore: To create a client using the %s database, please use NewClient", DefaultDatabaseID)
+	}
+
+	client, err := NewClient(ctx, projectID, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	client.databaseID = databaseID
+	return client, nil
+}
+
 func detectProjectID(ctx context.Context, opts ...option.ClientOption) (string, error) {
 	creds, err := transport.Creds(ctx, opts...)
 	if err != nil {
@@ -127,12 +159,25 @@ func (c *Client) path() string {
 }
 
 func withResourceHeader(ctx context.Context, resource string) context.Context {
-	md, _ := metadata.FromOutgoingContext(ctx)
+	md, ok := metadata.FromOutgoingContext(ctx)
+	if !ok {
+		md = metadata.New(nil)
+	}
 	md = md.Copy()
 	md[resourcePrefixHeader] = []string{resource}
 	return metadata.NewOutgoingContext(ctx, md)
 }
 
+func withRequestParamsHeader(ctx context.Context, requestParams string) context.Context {
+	md, ok := metadata.FromOutgoingContext(ctx)
+	if !ok {
+		md = metadata.New(nil)
+	}
+	md = md.Copy()
+	md[reqParamsHeader] = []string{requestParams}
+	return metadata.NewOutgoingContext(ctx, md)
+}
+
 // Collection creates a reference to a collection with the given path.
 // A path is a sequence of IDs separated by slashes.
 //
@@ -234,7 +279,9 @@ func (c *Client) getAll(ctx context.Context, docRefs []*DocumentRef, tid []byte,
 		req.ConsistencySelector = &pb.BatchGetDocumentsRequest_Transaction{Transaction: tid}
 	}
 
-	streamClient, err := c.c.BatchGetDocuments(withResourceHeader(ctx, req.Database), req)
+	batchGetDocsCtx := withResourceHeader(ctx, req.Database)
+	batchGetDocsCtx = withRequestParamsHeader(batchGetDocsCtx, reqParamsHeaderVal(c.path()))
+	streamClient, err := c.c.BatchGetDocuments(batchGetDocsCtx, req)
 	if err != nil {
 		return nil, err
 	}

diff --git a/firestore/client_test.go b/firestore/client_test.go
@@ -32,6 +32,45 @@ var testClient = &Client{
 	readSettings: &readSettings{},
 }
 
+func TestNewClientWithDatabase(t *testing.T) {
+	for _, tc := range []struct {
+		desc       string
+		databaseID string
+		projectID  string
+		wantErr    bool
+	}{
+		{
+			desc:       "Empty databaseID",
+			databaseID: "",
+			projectID:  "p1",
+			wantErr:    true,
+		},
+		{
+			desc:       "Error from NewClient bubbled to NewClientWithDatabase",
+			databaseID: "db1",
+			projectID:  "",
+			wantErr:    true,
+		},
+		{
+			desc:       "Valid databaseID",
+			databaseID: "db1",
+			projectID:  "p1",
+			wantErr:    false,
+		},
+	} {
+		client, err := NewClientWithDatabase(context.Background(), tc.projectID, tc.databaseID)
+
+		if err != nil && !tc.wantErr {
+			t.Errorf("NewClientWithDatabase: %s got %v want nil", tc.desc, err)
+		} else if err == nil && tc.wantErr {
+			t.Errorf("NewClientWithDatabase: %s got %v wanted error", tc.desc, err)
+		} else if err == nil && tc.databaseID != client.databaseID {
+			t.Errorf("NewClientWithDatabase: %s got %v want %v", tc.desc, client.databaseID, tc.databaseID)
+		}
+
+	}
+}
+
 func TestClientCollectionAndDoc(t *testing.T) {
 	coll1 := testClient.Collection("X")
 	db := "projects/projectID/databases/(default)"

diff --git a/firestore/integration_test.go b/firestore/integration_test.go
@@ -26,7 +26,7 @@ import (
 	"reflect"
 	"runtime"
 	"sort"
-	"sync"
+	"strings"
 	"testing"
 	"time"
 
@@ -46,15 +46,30 @@ import (
 )
 
 func TestMain(m *testing.M) {
-	initIntegrationTest()
-	status := m.Run()
-	cleanupIntegrationTest()
-	os.Exit(status)
+	databaseIDs := []string{DefaultDatabaseID}
+	databasesStr, ok := os.LookupEnv(envDatabases)
+	if ok {
+		databaseIDs = append(databaseIDs, strings.Split(databasesStr, ",")...)
+	}
+
+	testParams = make(map[string]interface{})
+	for _, databaseID := range databaseIDs {
+		testParams["databaseID"] = databaseID
+		initIntegrationTest()
+		status := m.Run()
+		if status != 0 {
+			os.Exit(status)
+		}
+		cleanupIntegrationTest()
+	}
+
+	os.Exit(0)
 }
 
 const (
 	envProjID     = "GCLOUD_TESTS_GOLANG_FIRESTORE_PROJECT_ID"
 	envPrivateKey = "GCLOUD_TESTS_GOLANG_FIRESTORE_KEY"
+	envDatabases  = "GCLOUD_TESTS_GOLANG_FIRESTORE_DATABASES"
 )
 
 var (
@@ -64,9 +79,12 @@ var (
 	collectionIDs = uid.NewSpace("go-integration-test", nil)
 	wantDBPath    string
 	indexNames    []string
+	testParams    map[string]interface{}
 )
 
 func initIntegrationTest() {
+	databaseID := testParams["databaseID"].(string)
+	log.Printf("Setting up tests to run on databaseID: %q\n", databaseID)
 	flag.Parse() // needed for testing.Short()
 	if testing.Short() {
 		return
@@ -84,7 +102,7 @@ func initIntegrationTest() {
 		log.Fatal("The project key must be set. See CONTRIBUTING.md for details")
 	}
 	projectPath := "projects/" + testProjectID
-	wantDBPath = projectPath + "/databases/(default)"
+	wantDBPath = projectPath + "/databases/" + databaseID
 
 	ti := &testutil.HeadersEnforcer{
 		Checkers: []*testutil.HeaderChecker{
@@ -105,7 +123,7 @@ func initIntegrationTest() {
 		},
 	}
 	copts := append(ti.CallOptions(), option.WithTokenSource(ts))
-	c, err := NewClient(ctx, testProjectID, copts...)
+	c, err := NewClientWithDatabase(ctx, testProjectID, databaseID, copts...)
 	if err != nil {
 		log.Fatalf("NewClient: %v", err)
 	}
@@ -131,13 +149,11 @@ func initIntegrationTest() {
 // Without indexes, FailedPrecondition rpc error is seen with
 // desc 'The query requires multiple indexes'.
 func createIndexes(ctx context.Context, dbPath string) {
-	var createIndexWg sync.WaitGroup
 
 	indexFields := [][]string{{"updatedAt", "weight", "height"}, {"weight", "height"}}
 	indexNames = make([]string, len(indexFields))
 	indexParent := fmt.Sprintf("%s/collectionGroups/%s", dbPath, iColl.ID)
 
-	createIndexWg.Add(len(indexFields))
 	for i, fields := range indexFields {
 		var adminPbIndexFields []*adminpb.Index_IndexField
 		for _, field := range fields {
@@ -155,21 +171,17 @@ func createIndexes(ctx context.Context, dbPath string) {
 				Fields:     adminPbIndexFields,
 			},
 		}
-		go func(req *adminpb.CreateIndexRequest, i int) {
-			op, createErr := iAdminClient.CreateIndex(ctx, req)
-			if createErr != nil {
-				log.Fatalf("CreateIndex: %v", createErr)
-			}
+		op, createErr := iAdminClient.CreateIndex(ctx, req)
+		if createErr != nil {
+			log.Fatalf("CreateIndex: %v", createErr)
+		}
 
-			createdIndex, waitErr := op.Wait(ctx)
-			if waitErr != nil {
-				log.Fatalf("Wait: %v", waitErr)
-			}
-			indexNames[i] = createdIndex.Name
-			createIndexWg.Done()
-		}(req, i)
+		createdIndex, waitErr := op.Wait(ctx)
+		if waitErr != nil {
+			log.Fatalf("Wait: %v", waitErr)
+		}
+		indexNames[i] = createdIndex.Name
 	}
-	createIndexWg.Wait()
 }
 
 // deleteIndexes deletes composite indexes created in createIndexes function
@@ -2232,6 +2244,40 @@ func TestIntegration_ColGroupRefPartitionsLarge(t *testing.T) {
 	}
 }
 
+func TestIntegration_NewClientWithDatabase(t *testing.T) {
+	if testing.Short() {
+		t.Skip("Integration tests skipped in short mode")
+	}
+	for _, tc := range []struct {
+		desc    string
+		dbName  string
+		wantErr bool
+		opt     []option.ClientOption
+	}{
+		{
+			desc:    "Success",
+			dbName:  testParams["databaseID"].(string),
+			wantErr: false,
+		},
+		{
+			desc:    "Error from NewClient bubbled to NewClientWithDatabase",
+			dbName:  testParams["databaseID"].(string),
+			wantErr: true,
+			opt:     []option.ClientOption{option.WithCredentialsFile("non existent filepath")},
+		},
+	} {
+		ctx := context.Background()
+		c, err := NewClientWithDatabase(ctx, iClient.projectID, tc.dbName, tc.opt...)
+		if err != nil && !tc.wantErr {
+			t.Errorf("NewClientWithDatabase: %s got %v want nil", tc.desc, err)
+		} else if err == nil && tc.wantErr {
+			t.Errorf("NewClientWithDatabase: %s got %v wanted error", tc.desc, err)
+		} else if err == nil && c.databaseID != tc.dbName {
+			t.Errorf("NewClientWithDatabase: %s got %v want %v", tc.desc, c.databaseID, tc.dbName)
+		}
+	}
+}
+
 // TestIntegration_BulkWriter_Set tests setting values and serverTimeStamp in single write.
 func TestIntegration_BulkWriter_Set(t *testing.T) {
 	doc := iColl.NewDoc()

diff --git a/internal/kokoro/continuous.sh b/internal/kokoro/continuous.sh
@@ -34,6 +34,7 @@ export GCLOUD_TESTS_GOLANG_KEY=$GOOGLE_APPLICATION_CREDENTIALS
 export GCLOUD_TESTS_GOLANG_DATASTORE_DATABASES=database-01
 export GCLOUD_TESTS_GOLANG_FIRESTORE_PROJECT_ID=gcloud-golang-firestore-tests
 export GCLOUD_TESTS_GOLANG_FIRESTORE_KEY=$KOKORO_KEYSTORE_DIR/72523_go_firestore_integration_service_account
+export GCLOUD_TESTS_GOLANG_FIRESTORE_DATABASES=database-02
 export GCLOUD_TESTS_API_KEY=`cat $KOKORO_KEYSTORE_DIR/72523_go_gcloud_tests_api_key`
 export GCLOUD_TESTS_GOLANG_KEYRING=projects/dulcet-port-762/locations/us/keyRings/go-integration-test
 export GCLOUD_TESTS_GOLANG_PROFILER_ZONE="us-west1-b"