feat: Multi-db support in Datastore.

apis/Google.Cloud.Datastore.V1/Google.Cloud.Datastore.V1.GeneratedSnippets/Google.Cloud.Datastore.V1.GeneratedSnippets.csproj

@@ -10,6 +10,7 @@
     <PackageReference Include="Google.Api.Gax.Grpc.Testing" Version="[4.4.0, 5.0.0)" />
     <ProjectReference Include="..\..\..\tools\Google.Cloud.ClientTesting\Google.Cloud.ClientTesting.csproj" />
     <ProjectReference Include="..\Google.Cloud.Datastore.V1\Google.Cloud.Datastore.V1.csproj" />
+    <PackageReference Include="Google.Cloud.Firestore.Admin.V1" Version="[3.3.0, 4.0.0)" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.8.0" />
     <PackageReference Include="NSubstitute" Version="5.1.0" />
     <PackageReference Include="System.Linq.Async" Version="6.0.1" />

apis/Google.Cloud.Datastore.V1/Google.Cloud.Datastore.V1.IntegrationTests/DatastoreDbTest.cs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+using Google.Cloud.ClientTesting;
 using Grpc.Core;
 using System;
 using System.Collections.Generic;
@@ -68,6 +69,39 @@ public void Lookup_NamespaceOnly()
             Assert.Equal("bar", (string)entity["foo"]);
         }
 
+        [Fact(Skip = "Multiple databases are available only in preview right now.")]
+        public async Task MultiDb_InsertLookupDelete()
+        {
+            await _fixture.RunWithTemporaryDatabaseAsync(databaseId =>
+            {
+                var db = _fixture.CreateDatastoreDbWithDatabase(databaseId);
+                var keyFactory = db.CreateKeyFactory("test_dbid");
+                var entities = new[]
+                {
+                    new Entity { Key = keyFactory.CreateKey("x"), ["description"] = "predefined_key" },
+                    new Entity { Key = keyFactory.CreateIncompleteKey(), ["description"] = "incomplete_key" }
+                };
+
+                var insertedKeys = db.Insert(entities);
+
+                Assert.Null(insertedKeys[0]); // Insert with predefined key
+                Assert.NotNull(insertedKeys[1]); // Insert with incomplete key
+                Assert.Equal(insertedKeys[1], entities[1].Key); // Inserted key is propagated into entity
+
+                var lookupKey = new Key
+                {
+                    PartitionId = new() { ProjectId = _fixture.ProjectId, NamespaceId = _fixture.NamespaceId, DatabaseId = databaseId },
+                    Path = { insertedKeys[1].Path }
+                };
+                var fetchedEntity = db.Lookup(lookupKey);
+                Assert.NotNull(fetchedEntity);
+                Assert.Equal("incomplete_key", fetchedEntity["description"]);
+
+                db.Delete(lookupKey);
+                Assert.Null(db.Lookup(lookupKey));
+            });
+        }
+
         [Fact]
         public async Task Lookup_NoPartition()
         {
@@ -85,7 +119,7 @@ public async Task Lookup_NoPartition()
                 var lookupKey = new Key { Path = { insertedKey.Path } };
                 var result = db.Lookup(lookupKey);
                 Assert.NotNull(result);
-                Assert.Equal("bar", (string)entity["foo"]);
+                Assert.Equal("bar", (string) entity["foo"]);
 
                 // And the same lookup asynchronously...
                 Assert.Equal(result, await db.LookupAsync(lookupKey));

apis/Google.Cloud.Datastore.V1/Google.Cloud.Datastore.V1.IntegrationTests/DatastoreFixture.cs

@@ -1,4 +1,4 @@
-// Copyright 2016 Google Inc. All Rights Reserved.
+// Copyright 2016 Google Inc. All Rights Reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -12,12 +12,18 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 using Google.Api.Gax;
+using Google.Api.Gax.ResourceNames;
+using Google.Apis.Auth.OAuth2;
+using Google.Apis.Http;
 using Google.Cloud.ClientTesting;
+using Google.Cloud.Firestore.Admin.V1;
 using System;
 using System.Linq;
+using System.Net.Http;
 using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
+using static Google.Cloud.Firestore.Admin.V1.Database.Types;
 
 namespace Google.Cloud.Datastore.V1.IntegrationTests
 {
@@ -31,12 +37,35 @@ public sealed class DatastoreFixture : CloudProjectFixtureBase, ICollectionFixtu
         private const int RetryCount = 10;
         private static readonly TimeSpan RetryDelay = TimeSpan.FromSeconds(3);
 
+        private readonly HttpClient _firestoreRestApiHttpClient = new();
+        private readonly FirestoreAdminClient _firestoreAdminClient;
+
+        internal bool RunningOnEmulator { get; }
+
+        // Creating databases and indexes can take a while.
+        // We don't want to wait *forever* (which would be the behavior of default poll settings)
+        // but we need to have a timeout of more than a minute.
+        private static readonly PollSettings AdminOperationPollSettings =
+            new PollSettings(expiration: Expiration.FromTimeout(TimeSpan.FromMinutes(5)), delay: TimeSpan.FromSeconds(5));
+
         public string NamespaceId { get; }
         public PartitionId PartitionId => new PartitionId { ProjectId = ProjectId, NamespaceId = NamespaceId };
 
         public DatastoreFixture()
         {
             NamespaceId = IdGenerator.FromDateTime(prefix: "test-");
+            RunningOnEmulator = !string.IsNullOrEmpty(Environment.GetEnvironmentVariable("DATASTORE_EMULATOR_HOST"));
+            _firestoreAdminClient = FirestoreAdminClient.Create();
+
+            // Scope used for the REST API to delete databases.
+            string scope = "https://www.googleapis.com/auth/datastore";
+            string credentialsPath = Environment.GetEnvironmentVariable("GOOGLE_APPLICATION_CREDENTIALS");
+
+            // Initalize credentials to be used with REST API calls. 
+            GoogleCredential googleCredential = GoogleCredential.FromFile(credentialsPath).CreateScoped(scope);
+            _firestoreRestApiHttpClient = new HttpClientFactory()
+                .CreateHttpClient(new CreateHttpClientArgs { Initializers = { googleCredential } });
+            _firestoreRestApiHttpClient.BaseAddress = new Uri("https://firestore.googleapis.com");
         }
 
         public override void Dispose()
@@ -99,5 +128,70 @@ public DatastoreDb CreateDatastoreDb(string namespaceId = null)
             };
             return builder.Build();
         }
+
+        public DatastoreDb CreateDatastoreDbWithDatabase(string databaseId)
+        {
+            var builder = new DatastoreDbBuilder
+            {
+                ProjectId = ProjectId,
+                NamespaceId = NamespaceId,
+                DatabaseId = databaseId ?? "",
+                EmulatorDetection = EmulatorDetection.EmulatorOrProduction
+            };
+            return builder.Build();
+        }
+
+        public async Task RunWithTemporaryDatabaseAsync(Action<string> testFunction)
+        {
+            var databaseId = IdGenerator.FromDateTime(prefix: "test-db-");
+            await CreateDatabaseAsync(databaseId);
+
+            try
+            {
+                testFunction(databaseId);
+            }
+            finally
+            {
+                // Cleanup the test database.
+                try
+                {
+                    await DeleteDatabaseAsync(databaseId);
+                }
+                catch (Exception)
+                {
+                    // Silently ignore errors to prevent tests from failing.
+                }
+            }
+        }
+
+        /// <summary>
+        /// Creates a new Datastore database using Firestore Admin Client.
+        /// </summary>
+        public async Task CreateDatabaseAsync(string databaseId)
+        {
+            var pr = new ProjectName(ProjectId);
+            // Creating a new database is not supported on Datastore emulator.
+            Assert.False(RunningOnEmulator);
+            var operation = await _firestoreAdminClient.CreateDatabaseAsync(
+                new ProjectName(ProjectId),
+                new Database { LocationId = "us-east1", Type = DatabaseType.DatastoreMode },
+                databaseId);
+            await operation.PollUntilCompletedAsync(AdminOperationPollSettings);
+            Console.WriteLine($"Success creating database {databaseId}");
+        }
+
+        private async Task DeleteDatabaseAsync(string databaseId)
+        {
+            var deleteDbUrl = $"/v1/projects/{ProjectId}/databases/{databaseId}";
+            try
+            {
+                var response = await _firestoreRestApiHttpClient.DeleteAsync(deleteDbUrl).ConfigureAwait(false);
+                response.EnsureSuccessStatusCode();
+            }
+            catch (Exception)
+            {
+                // Silently ignore errors to prevent tests from failing.
+            }
+        }
     }
 }

apis/Google.Cloud.Datastore.V1/Google.Cloud.Datastore.V1.IntegrationTests/Google.Cloud.Datastore.V1.IntegrationTests.csproj

@@ -10,6 +10,7 @@
     <PackageReference Include="Google.Api.Gax.Grpc.Testing" Version="[4.4.0, 5.0.0)" />
     <ProjectReference Include="..\..\..\tools\Google.Cloud.ClientTesting\Google.Cloud.ClientTesting.csproj" />
     <ProjectReference Include="..\Google.Cloud.Datastore.V1\Google.Cloud.Datastore.V1.csproj" />
+    <PackageReference Include="Google.Cloud.Firestore.Admin.V1" Version="[3.3.0, 4.0.0)" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.8.0" />
     <PackageReference Include="NSubstitute" Version="5.1.0" />
     <PackageReference Include="System.Linq.Async" Version="6.0.1" />

apis/Google.Cloud.Datastore.V1/Google.Cloud.Datastore.V1.Snippets/Google.Cloud.Datastore.V1.Snippets.csproj

@@ -10,6 +10,7 @@
     <PackageReference Include="Google.Api.Gax.Grpc.Testing" Version="[4.4.0, 5.0.0)" />
     <ProjectReference Include="..\..\..\tools\Google.Cloud.ClientTesting\Google.Cloud.ClientTesting.csproj" />
     <ProjectReference Include="..\Google.Cloud.Datastore.V1\Google.Cloud.Datastore.V1.csproj" />
+    <PackageReference Include="Google.Cloud.Firestore.Admin.V1" Version="[3.3.0, 4.0.0)" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.8.0" />
     <PackageReference Include="NSubstitute" Version="5.1.0" />
     <PackageReference Include="System.Linq.Async" Version="6.0.1" />

apis/Google.Cloud.Datastore.V1/Google.Cloud.Datastore.V1.Tests/Google.Cloud.Datastore.V1.Tests.csproj

@@ -10,6 +10,7 @@
     <PackageReference Include="Google.Api.Gax.Grpc.Testing" Version="[4.4.0, 5.0.0)" />
     <ProjectReference Include="..\..\..\tools\Google.Cloud.ClientTesting\Google.Cloud.ClientTesting.csproj" />
     <ProjectReference Include="..\Google.Cloud.Datastore.V1\Google.Cloud.Datastore.V1.csproj" />
+    <PackageReference Include="Google.Cloud.Firestore.Admin.V1" Version="[3.3.0, 4.0.0)" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.8.0" />
     <PackageReference Include="NSubstitute" Version="5.1.0" />
     <PackageReference Include="System.Linq.Async" Version="6.0.1" />

apis/Google.Cloud.Datastore.V1/Google.Cloud.Datastore.V1.Tests/KeyFactoryTest.cs

@@ -1,4 +1,4 @@
-// Copyright 2016 Google Inc. All Rights Reserved.
+// Copyright 2016 Google Inc. All Rights Reserved.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -22,11 +22,11 @@ public class KeyFactoryTest
         [Fact]
         public void PartitionOnlyFactory()
         {
-            var factory = new KeyFactory("project", "ns", "book");
+            var factory = new KeyFactory("project", "ns", "db", "book");
             var actual = factory.CreateKey(10L);
             var expected = new Key
             {
-                PartitionId = new PartitionId { ProjectId = "project", NamespaceId = "ns" },
+                PartitionId = new PartitionId { ProjectId = "project", NamespaceId = "ns", DatabaseId = "db" },
                 Path = { new PathElement { Id = 10L, Kind = "book" } }
             };
             Assert.Equal(expected, actual);
@@ -37,14 +37,14 @@ public void FactoryFromKey()
         {
             var parentKey = new Key
             {
-                PartitionId = new PartitionId { ProjectId = "project", NamespaceId = "ns" },
+                PartitionId = new PartitionId { ProjectId = "project", NamespaceId = "ns", DatabaseId = "db" },
                 Path = { new PathElement { Id = 10L, Kind = "author" } }
             };
             var factory = new KeyFactory(parentKey, "book");
             var actual = factory.CreateKey("subkey-name");
             var expected = new Key
             {
-                PartitionId = new PartitionId { ProjectId = "project", NamespaceId = "ns" },
+                PartitionId = new PartitionId { ProjectId = "project", NamespaceId = "ns", DatabaseId = "db" },
                 Path = {
                     new PathElement { Id = 10L, Kind = "author" },
                     new PathElement { Name = "subkey-name", Kind = "book" }
@@ -58,15 +58,15 @@ public void FactoryFromEntity()
         {
             var parentKey = new Key
             {
-                PartitionId = new PartitionId { ProjectId = "project", NamespaceId = "ns" },
+                PartitionId = new PartitionId { ProjectId = "project", NamespaceId = "ns", DatabaseId = "db" },
                 Path = { new PathElement { Id = 10L, Kind = "author" } }
             };
             var book = new Entity { Key = parentKey };
             var factory = new KeyFactory(book, "book");
             var actual = factory.CreateKey(20L);
             var expected = new Key
             {
-                PartitionId = new PartitionId { ProjectId = "project", NamespaceId = "ns" },
+                PartitionId = new PartitionId { ProjectId = "project", NamespaceId = "ns", DatabaseId = "db" },
                 Path = {
                     new PathElement { Id = 10L, Kind = "author" },
                     new PathElement { Id = 20L, Kind = "book" }

apis/Google.Cloud.Datastore.V1/Google.Cloud.Datastore.V1/DatastoreDb.cs

@@ -58,6 +58,8 @@ public abstract class DatastoreDb
     {
         internal const string DefaultNamespaceId = "";
 
+        internal const string DefaultDatabaseId = "";
+
         /// <summary>
         /// The <see cref="DatastoreClient"/> used for all remote operations.
         /// </summary>
@@ -73,6 +75,11 @@ public abstract class DatastoreDb
         /// </summary>
         public virtual string NamespaceId { get { throw new NotImplementedException(); } }
 
+        /// <summary>
+        /// The ID of the database against which the request is to be made.
+        /// </summary>
+        public virtual string DatabaseId { get { throw new NotImplementedException(); } }
+
         /// <summary>
         /// Creates a <see cref="DatastoreDb"/> to operate on the partition identified by <paramref name="projectId"/>
         /// and <paramref name="namespaceId"/>, using the <paramref name="client"/> to perform remote operations.
@@ -509,6 +516,7 @@ public virtual Task<IReadOnlyList<Entity>> LookupAsync(IEnumerable<Key> keys, Re
         internal static IReadOnlyList<Entity> LookupImpl(
             DatastoreClient client,
             string projectId,
+            string databaseId,
             ReadOptions readOptions,
             IEnumerable<Key> keys,
             CallSettings callSettings)
@@ -522,7 +530,15 @@ internal static IReadOnlyList<Entity> LookupImpl(
             // TODO: Limit how many times we go round? Ensure that we make progress on each iteration?
             while (keysToFetch.Count() > 0)
             {
-                var response = client.Lookup(projectId, readOptions, keysToFetch, callSettings);
+                var lookupRequest = new LookupRequest
+                {
+                    ProjectId = projectId,
+                    DatabaseId = databaseId,
+                    ReadOptions = readOptions,
+                    Keys = { keysToFetch },
+                };
+
+                var response = client.Lookup(lookupRequest, callSettings);
                 foreach (var found in response.Found)
                 {
                     foreach (var index in keyToIndex[found.Entity.Key])
@@ -539,6 +555,7 @@ internal static IReadOnlyList<Entity> LookupImpl(
         internal static async Task<IReadOnlyList<Entity>> LookupImplAsync(
             DatastoreClient client,
             string projectId,
+            string databaseId,
             ReadOptions readOptions,
             IEnumerable<Key> keys,
             CallSettings callSettings)
@@ -552,7 +569,15 @@ internal static async Task<IReadOnlyList<Entity>> LookupImplAsync(
             // TODO: Limit how many times we go round? Ensure that we make progress on each iteration?
             while (keysToFetch.Count() > 0)
             {
-                var response = await client.LookupAsync(projectId, readOptions, keysToFetch, callSettings).ConfigureAwait(false);
+                var lookupRequest = new LookupRequest
+                {
+                    ProjectId = projectId,
+                    DatabaseId = databaseId,
+                    ReadOptions = readOptions,
+                    Keys = { keysToFetch },
+                };
+
+                var response = await client.LookupAsync(lookupRequest, callSettings).ConfigureAwait(false);
                 foreach (var found in response.Found)
                 {
                     foreach (var index in keyToIndex[found.Entity.Key])