From 94d4b1b58d2c8f3dac18e7efb0be641b6311c775 Mon Sep 17 00:00:00 2001 From: Eric Schmidt Date: Mon, 18 Sep 2023 09:52:29 -0700 Subject: [PATCH] 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 Co-authored-by: meredithslota --- CONTRIBUTING.md | 4 ++ firestore/client.go | 53 +++++++++++++++++++-- firestore/client_test.go | 39 +++++++++++++++ firestore/integration_test.go | 90 ++++++++++++++++++++++++++--------- internal/kokoro/continuous.sh | 1 + 5 files changed, 162 insertions(+), 25 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 3a391131aa84..18f175dac99b 100644 --- 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 index dfee95b21de3..9fad2c592c76 100644 --- 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 index 0f6fb3ba913d..b39bc95267e0 100644 --- 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 index 7f881a855f1c..0cc7ffb901b6 100644 --- 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 index b6e397a4dfb5..21b718b3d31e 100755 --- 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"