CosmosClient: Adds CreateAndInitializeAsync Method

Microsoft.Azure.Cosmos/src/CosmosClient.cs

@@ -5,6 +5,7 @@
 namespace Microsoft.Azure.Cosmos
 {
     using System;
+    using System.Collections.Generic;
     using System.Net;
     using System.Text;
     using System.Threading;
@@ -257,6 +258,154 @@ public CosmosClient(
                 clientOptions);
         }
 
+        /// <summary>
+        /// Creates a new CosmosClient with the account endpoint URI string and TokenCredential.
+        /// In addition to that it initializes the client with containers provided i.e The SDK warms up the caches and 
+        /// connections before the first call to the service is made. Use this to obtain lower latency while startup of your application.
+        /// CosmosClient is thread-safe. Its recommended to maintain a single instance of CosmosClient per lifetime 
+        /// of the application which enables efficient connection management and performance. Please refer to the
+        /// <see href="https://docs.microsoft.com/azure/cosmos-db/performance-tips">performance guide</see>.
+        /// </summary>
+        /// <param name="accountEndpoint">The cosmos service endpoint to use</param>
+        /// <param name="authKeyOrResourceToken">The cosmos account key or resource token to use to create the client.</param>
+        /// <param name="containers">Containers to be initialized identified by it's database name and container name.</param>
+        /// <param name="cosmosClientOptions">(Optional) client options</param>
+        /// <param name="cancellationToken">(Optional) Cancellation Token</param>
+        /// <returns>
+        /// A CosmosClient object.
+        /// </returns>
+        /// <example>
+        /// The CosmosClient is created with the AccountEndpoint, AccountKey or ResourceToken and 2 containers in the account are initialized
+        /// <code language="c#">
+        /// <![CDATA[
+        /// using Microsoft.Azure.Cosmos;
+        /// List<(string, string)> containersToInitialize = new List<(string, string)>
+        /// { ("DatabaseName1", "ContainerName1"), ("DatabaseName2", "ContainerName2") };
+        /// 
+        /// CosmosClient cosmosClient = await CosmosClient.CreateAndInitializeAsync("account-endpoint-from-portal", 
+        ///                                                                         "account-key-from-portal",
+        ///                                                                         containersToInitialize)
+        /// 
+        /// // Dispose cosmosClient at application exit
+        /// ]]>
+        /// </code>
+        /// </example>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+        static async Task<CosmosClient> CreateAndInitializeAsync(string accountEndpoint, 
+                                                                        string authKeyOrResourceToken,
+                                                                        IReadOnlyList<(string databaseId, string containerId)> containers,
+                                                                        CosmosClientOptions cosmosClientOptions = null,
+                                                                        CancellationToken cancellationToken = default)
+        {
+            if (containers == null)
+            {
+                throw new ArgumentNullException(nameof(containers));
+            }
+
+            CosmosClient cosmosClient = new CosmosClient(accountEndpoint, 
+                                                         authKeyOrResourceToken, 
+                                                         cosmosClientOptions);
+
+            await cosmosClient.InitializeContainersAsync(containers, cancellationToken);
+            return cosmosClient;
+        }
+
+        /// <summary>
+        /// Creates a new CosmosClient with the account endpoint URI string and TokenCredential.
+        /// In addition to that it initializes the client with containers provided i.e The SDK warms up the caches and 
+        /// connections before the first call to the service is made. Use this to obtain lower latency while startup of your application.
+        /// CosmosClient is thread-safe. Its recommended to maintain a single instance of CosmosClient per lifetime 
+        /// of the application which enables efficient connection management and performance. Please refer to the
+        /// <see href="https://docs.microsoft.com/azure/cosmos-db/performance-tips">performance guide</see>.
+        /// </summary>
+        /// <param name="connectionString">The connection string to the cosmos account. ex: https://mycosmosaccount.documents.azure.com:443/;AccountKey=SuperSecretKey; </param>
+        /// <param name="containers">Containers to be initialized identified by it's database name and container name.</param>
+        /// <param name="cosmosClientOptions">(Optional) client options</param>
+        /// <param name="cancellationToken">(Optional) Cancellation Token</param>
+        /// <returns>
+        /// A CosmosClient object.
+        /// </returns>
+        /// <example>
+        /// The CosmosClient is created with the ConnectionString and 2 containers in the account are initialized
+        /// <code language="c#">
+        /// <![CDATA[
+        /// using Microsoft.Azure.Cosmos;
+        /// List<(string, string)> containersToInitialize = new List<(string, string)>
+        /// { ("DatabaseName1", "ContainerName1"), ("DatabaseName2", "ContainerName2") };
+        /// 
+        /// CosmosClient cosmosClient = await CosmosClient.CreateAndInitializeAsync("connection-string-from-portal",
+        ///                                                                         containersToInitialize)
+        /// 
+        /// // Dispose cosmosClient at application exit
+        /// ]]>
+        /// </code>
+        /// </example>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+        static async Task<CosmosClient> CreateAndInitializeAsync(string connectionString,
+                                                                        IReadOnlyList<(string databaseId, string containerId)> containers,
+                                                                        CosmosClientOptions cosmosClientOptions = null,
+                                                                        CancellationToken cancellationToken = default)
+        {
+            if (containers == null)
+            {
+                throw new ArgumentNullException(nameof(containers));
+            }
+
+            CosmosClient cosmosClient = new CosmosClient(connectionString,
+                                                         cosmosClientOptions);
+
+            await cosmosClient.InitializeContainersAsync(containers, cancellationToken);
+            return cosmosClient;
+        }
+
+        /// <summary>
+        /// Creates a new CosmosClient with the account endpoint URI string and TokenCredential.
+        /// In addition to that it initializes the client with containers provided i.e The SDK warms up the caches and 
+        /// connections before the first call to the service is made. Use this to obtain lower latency while startup of your application.
+        /// CosmosClient is thread-safe. Its recommended to maintain a single instance of CosmosClient per lifetime 
+        /// of the application which enables efficient connection management and performance. Please refer to the
+        /// <see href="https://docs.microsoft.com/azure/cosmos-db/performance-tips">performance guide</see>.
+        /// </summary>
+        /// <param name="accountEndpoint">The cosmos service endpoint to use.</param>
+        /// <param name="tokenCredential"><see cref="TokenCredential"/>The token to provide AAD token for authorization.</param>
+        /// <param name="containers">Containers to be initialized identified by it's database name and container name.</param>
+        /// <param name="cosmosClientOptions">(Optional) client options</param>
+        /// <param name="cancellationToken">(Optional) Cancellation Token</param>
+        /// <returns>
+        /// A CosmosClient object.
+        /// </returns>
+#if PREVIEW
+        public
+#else
+        internal
+#endif
+        static async Task<CosmosClient> CreateAndInitializeAsync(string accountEndpoint,
+                                                                        TokenCredential tokenCredential,
+                                                                        IReadOnlyList<(string databaseId, string containerId)> containers,
+                                                                        CosmosClientOptions cosmosClientOptions = null,
+                                                                        CancellationToken cancellationToken = default)
+        {
+            if (containers == null)
+            {
+                throw new ArgumentNullException(nameof(containers));
+            }
+
+            CosmosClient cosmosClient = new CosmosClient(accountEndpoint,
+                                                         tokenCredential,
+                                                         cosmosClientOptions);
+
+            await cosmosClient.InitializeContainersAsync(containers, cancellationToken);
+            return cosmosClient;
+        }
+
         /// <summary>
         /// Used for unit testing only.
         /// </summary>
@@ -1042,6 +1191,57 @@ private FeedIteratorInternal GetDatabaseQueryStreamIteratorHelper(
                options: requestOptions);
         }
 
+        private Task InitializeContainersAsync(IReadOnlyList<(string databaseId, string containerId)> containers, 
+                                          CancellationToken cancellationToken)
+        {
+            try
+            {
+                List<Task> tasks = new List<Task>();
+                foreach ((string databaseId, string containerId) in containers)
+                {
+                    tasks.Add(this.InitializeContainerAsync(databaseId, containerId, cancellationToken));
+                }
+
+                return Task.WhenAll(tasks);
+            }
+            catch
+            {
+                this.Dispose();
+                throw;
+            }
+        }
+
+        private async Task InitializeContainerAsync(string databaseId, string containerId, CancellationToken cancellationToken = default)
+        {
+            ContainerInternal container = (ContainerInternal)this.GetContainer(databaseId, containerId);
+            IReadOnlyList<FeedRange> feedRanges = await container.GetFeedRangesAsync(cancellationToken);
+            List<Task> tasks = new List<Task>();
+            foreach (FeedRange feedRange in feedRanges)
+            {
+                tasks.Add(CosmosClient.InitializeFeedRangeAsync(container, feedRange, cancellationToken));
+            }
+
+            await Task.WhenAll(tasks);
+        }
+
+        private static async Task InitializeFeedRangeAsync(ContainerInternal container, FeedRange feedRange, CancellationToken cancellationToken = default)
+        {
+            // Do a dummy querry for each Partition Key Range to warm up the caches and connections
+            string guidToCheck = Guid.NewGuid().ToString();
+            QueryDefinition queryDefinition = new QueryDefinition($"select * from c where c.id = '{guidToCheck}'");
+            using (FeedIterator feedIterator = container.GetItemQueryStreamIterator(feedRange,
+                                                                                    queryDefinition,
+                                                                                    continuationToken: null,
+                                                                                    requestOptions: new QueryRequestOptions() { }))
+            {
+                while (feedIterator.HasMoreResults)
+                {
+                    using ResponseMessage response = await feedIterator.ReadNextAsync(cancellationToken);
+                    response.EnsureSuccessStatusCode();
+                }
+            }
+        }
+
         /// <summary>
         /// Dispose of cosmos client
         /// </summary>

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/ClientCreateAndInitializeTest.cs

@@ -0,0 +1,130 @@
+namespace Microsoft.Azure.Cosmos.SDK.EmulatorTests
+{
+    using System;
+    using System.Collections.Generic;
+    using System.Net;
+    using System.Net.Http;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Routing;
+    using Microsoft.Azure.Documents;
+    using Microsoft.VisualStudio.TestTools.UnitTesting;
+
+    [TestClass]
+    public class ClientCreateAndInitializeTest : BaseCosmosClientHelper
+    {
+        private ContainerInternal Container = null;
+        private const string PartitionKey = "/pk";
+
+        [TestInitialize]
+        public async Task TestInitialize()
+        {
+            this.cosmosClient = TestCommon.CreateCosmosClient(useGateway: false);
+            this.database = await this.cosmosClient.CreateDatabaseAsync(
+                   id: "ClientCreateAndInitializeDatabase");
+            ContainerResponse response = await this.database.CreateContainerAsync(
+                        new ContainerProperties(id: "ClientCreateAndInitializeContainer", partitionKeyPath: PartitionKey),
+                        throughput: 20000,
+                        cancellationToken: this.cancellationToken);
+            Assert.IsNotNull(response);
+            Assert.IsNotNull(response.Container);
+            Assert.IsNotNull(response.Resource);
+            this.Container = (ContainerInlineCore)response;
+
+            // Create items with different
+            for (int i = 0; i < 500; i++)
+            {
+                ToDoActivity item = ToDoActivity.CreateRandomToDoActivity();
+                item.pk = "Status" + i.ToString();
+                item.id = i.ToString();
+                ItemResponse<ToDoActivity> itemResponse = await this.Container.CreateItemAsync(item);
+                Assert.AreEqual(HttpStatusCode.Created, itemResponse.StatusCode);
+            }
+        }
+
+        [TestCleanup]
+        public async Task Cleanup()
+        {
+            await base.TestCleanup();
+        }
+
+        [TestMethod]
+        public async Task CreateAndInitializeTest()
+        {
+            int httpCallsMade = 0;
+            HttpClientHandlerHelper httpClientHandlerHelper = new HttpClientHandlerHelper
+            {
+                RequestCallBack = (request, cancellationToken) =>
+                {
+                    httpCallsMade++;
+                    return null;
+                }
+            };
+
+            (string endpoint, string authKey) = TestCommon.GetAccountInfo();
+            List<(string, string)> containers = new List<(string, string)> 
+            { ("ClientCreateAndInitializeDatabase", "ClientCreateAndInitializeContainer")};
+
+            CosmosClientOptions cosmosClientOptions = new CosmosClientOptions
+            {
+                HttpClientFactory = () => new HttpClient(httpClientHandlerHelper)
+            };
+
+            CosmosClient cosmosClient = await CosmosClient.CreateAndInitializeAsync(endpoint, authKey, containers, cosmosClientOptions);
+            Assert.IsNotNull(cosmosClient);
+            int httpCallsMadeAfterCreation = httpCallsMade;
+
+            ContainerInternal container = (ContainerInternal)cosmosClient.GetContainer("ClientCreateAndInitializeDatabase", "ClientCreateAndInitializeContainer");
+            ItemResponse<ToDoActivity> readResponse = await container.ReadItemAsync<ToDoActivity>("1", new Cosmos.PartitionKey("Status1"));
+            Assert.AreEqual(httpCallsMade, httpCallsMadeAfterCreation);
+            cosmosClient.Dispose();
+        }
+
+        [TestMethod]
+        [ExpectedException(typeof(HttpRequestException))]
+        public async Task AuthIncorrectTest()
+        {
+            List<(string databaseId, string containerId)> containers = new List<(string databaseId, string containerId)>
+            { ("ClientCreateAndInitializeDatabase", "ClientCreateAndInitializeContainer")};
+            string authKey = TestCommon.GetAccountInfo().authKey;
+            CosmosClient cosmosClient = await CosmosClient.CreateAndInitializeAsync("https://127.0.0.1:0000/", authKey, containers);
+            cosmosClient.Dispose();
+        }
+
+        [TestMethod]
+        [ExpectedException(typeof(CosmosException))]
+        public async Task DatabaseIncorrectTest()
+        {
+            List<(string databaseId, string containerId)> containers = new List<(string databaseId, string containerId)>
+            { ("IncorrectDatabase", "ClientCreateAndInitializeContainer")};
+            (string endpoint, string authKey) = TestCommon.GetAccountInfo();
+            try
+            {
+                CosmosClient cosmosClient = await CosmosClient.CreateAndInitializeAsync(endpoint, authKey, containers);
+            }
+            catch (CosmosException ex)
+            {
+                Assert.IsTrue(ex.StatusCode == HttpStatusCode.NotFound);
+                throw ex;
+            }
+        }
+
+        [TestMethod]
+        [ExpectedException(typeof(CosmosException))]
+        public async Task ContainerIncorrectTest()
+        {
+            List<(string databaseId, string containerId)> containers = new List<(string databaseId, string containerId)>
+            { ("ClientCreateAndInitializeDatabase", "IncorrectContainer")};
+            (string endpoint, string authKey) = TestCommon.GetAccountInfo();
+            try
+            {
+                CosmosClient cosmosClient = await CosmosClient.CreateAndInitializeAsync(endpoint, authKey, containers);
+            }
+            catch (CosmosException ex)
+            {
+                Assert.IsTrue(ex.StatusCode == HttpStatusCode.NotFound);
+                throw ex;
+            }
+        }
+    }
+}