CosmosClientOptions: Adds Private Custom Account Endpoints (#4265)

* Code changes to add regional endpoints for account metadata calls. * Code changes to refactor some codes. * Code changes to add unit tests. * Code changes to make minor code clean-up. * Code changes to fix tests. Refactored API. * Code changes to refactor the enumeration logic inside global endpoint manager. * Code changes to address review comments. * Code changes to fix minor API parameter. * Code changes to update the API naming. * Code changes to update some attribute names. * Code changes to refactor service endpoint creation logic. * Code changes to address review comments. * Code changes to address review comments. * Code changes to update the API contract. * Cosmetic code changes. * Code changes to address review comments.
Azure · Feb 9, 2024 · 0589e8a · 0589e8a
1 parent a9599bd
commit 0589e8a
Show file tree

Hide file tree

Showing 10 changed files with 415 additions and 49 deletions.
diff --git a/Microsoft.Azure.Cosmos/src/ConnectionPolicy.cs b/Microsoft.Azure.Cosmos/src/ConnectionPolicy.cs
@@ -29,6 +29,7 @@ internal sealed class ConnectionPolicy
 
         private Protocol connectionProtocol;
         private ObservableCollection<string> preferredLocations;
+        private ObservableCollection<Uri> accountInitializationCustomEndpoints;
 
         /// <summary>
         /// Initializes a new instance of the <see cref="ConnectionPolicy"/> class to connect to the Azure Cosmos DB service.
@@ -43,6 +44,7 @@ public ConnectionPolicy()
             this.MediaReadMode = MediaReadMode.Buffered;
             this.UserAgentContainer = new UserAgentContainer(clientId: 0);
             this.preferredLocations = new ObservableCollection<string>();
+            this.accountInitializationCustomEndpoints = new ObservableCollection<Uri>();
             this.EnableEndpointDiscovery = true;
             this.MaxConnectionLimit = defaultMaxConcurrentConnectionLimit;
             this.RetryOptions = new RetryOptions();
@@ -90,6 +92,27 @@ public void SetPreferredLocations(IReadOnlyList<string> regions)
             }
         }
 
+        /// <summary>
+        /// Sets the custom private endpoints required to fetch account information from
+        /// private domain names.
+        /// </summary>
+        /// <param name="customEndpoints">An instance of <see cref="IEnumerable{T}"/> containing the custom DNS endpoints
+        /// provided by the customer.</param>
+        public void SetAccountInitializationCustomEndpoints(
+            IEnumerable<Uri> customEndpoints)
+        {
+            if (customEndpoints == null)
+            {
+                throw new ArgumentNullException(nameof(customEndpoints));
+            }
+
+            this.accountInitializationCustomEndpoints.Clear();
+            foreach (Uri endpoint in customEndpoints)
+            {
+                this.accountInitializationCustomEndpoints.Add(endpoint);
+            }
+        }
+
         /// <summary>
         /// Gets or sets the maximum number of concurrent fanout requests sent to the Azure Cosmos DB service.
         /// </summary>
@@ -270,6 +293,24 @@ public Collection<string> PreferredLocations
             }
         }
 
+        /// <summary>
+        /// Gets the custom private endpoints for geo-replicated database accounts in the Azure Cosmos DB service. 
+        /// </summary>
+        /// <remarks>
+        /// <para>
+        /// During the CosmosClient initialization the account information, including the available regions, is obtained from the <see cref="CosmosClient.Endpoint"/>.
+        /// Should the global endpoint become inaccessible, the CosmosClient will attempt to obtain the account information issuing requests to the custom endpoints
+        /// provided in the customAccountEndpoints list.
+        /// </para>
+        /// </remarks>
+        public Collection<Uri> AccountInitializationCustomEndpoints
+        {
+            get
+            {
+                return this.accountInitializationCustomEndpoints;
+            }
+        }
+
         /// <summary>
         /// Gets or sets the flag to enable endpoint discovery for geo-replicated database accounts in the Azure Cosmos DB service.
         /// </summary>

diff --git a/Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs b/Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs
@@ -191,6 +191,42 @@ public string ApplicationName
         /// </example>
         /// <seealso href="https://docs.microsoft.com/azure/cosmos-db/high-availability#high-availability-with-cosmos-db-in-the-event-of-regional-outages">High availability on regional outages</seealso>
         public IReadOnlyList<string> ApplicationPreferredRegions { get; set; }
+
+        /// <summary>
+        /// Gets and sets the custom endpoints to use for account initialization for geo-replicated database accounts in the Azure Cosmos DB service. 
+        /// </summary>
+        /// <remarks>
+        /// <para>
+        /// During the CosmosClient initialization the account information, including the available regions, is obtained from the <see cref="CosmosClient.Endpoint"/>.
+        /// Should the global endpoint become inaccessible, the CosmosClient will attempt to obtain the account information issuing requests to the custom endpoints provided in <see cref="AccountInitializationCustomEndpoints"/>.
+        /// </para>
+        /// <para>
+        /// Nevertheless, this parameter remains optional and is recommended for implementation when a customer has configured an endpoint with a custom DNS hostname
+        /// (instead of accountname-region.documents.azure.com) etc. for their Cosmos DB account.
+        /// </para>
+        /// <para>
+        /// See also <seealso href="https://docs.microsoft.com/azure/cosmos-db/sql/troubleshoot-sdk-availability">Diagnose
+        /// and troubleshoot the availability of Cosmos SDKs</seealso> for more details.
+        /// </para>
+        /// </remarks>
+        /// <example>
+        /// <code language="c#">
+        /// <![CDATA[
+        /// CosmosClientOptions clientOptions = new CosmosClientOptions()
+        /// {
+        ///     AccountInitializationCustomEndpoints = new HashSet<Uri>()
+        ///     { 
+        ///         new Uri("custom.p-1.documents.azure.com"),
+        ///         new Uri("custom.p-2.documents.azure.com") 
+        ///     }
+        /// };
+        /// 
+        /// CosmosClient client = new CosmosClient("endpoint", "key", clientOptions);
+        /// ]]>
+        /// </code>
+        /// </example>
+        /// <seealso href="https://docs.microsoft.com/azure/cosmos-db/high-availability#high-availability-with-cosmos-db-in-the-event-of-regional-outages">High availability on regional outages</seealso>
+        public IEnumerable<Uri> AccountInitializationCustomEndpoints { get; set; }
 
         /// <summary>
         /// Get or set the maximum number of concurrent connections allowed for the target
@@ -847,6 +883,11 @@ internal virtual ConnectionPolicy GetConnectionPolicy(int clientId)
                 List<string> mappedRegions = this.ApplicationPreferredRegions.Select(s => mapper.GetCosmosDBRegionName(s)).ToList();
 
                 connectionPolicy.SetPreferredLocations(mappedRegions);
+            }
+
+            if (this.AccountInitializationCustomEndpoints != null)
+            {
+                connectionPolicy.SetAccountInitializationCustomEndpoints(this.AccountInitializationCustomEndpoints);
             }
 
             if (this.MaxRetryAttemptsOnRateLimitedRequests != null)

diff --git a/Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs b/Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs
@@ -308,6 +308,41 @@ public CosmosClientBuilder WithApplicationPreferredRegions(IReadOnlyList<string>
             return this;
         }
 
+        /// <summary>
+        /// Sets the custom endpoints to use for account initialization for geo-replicated database accounts in the Azure Cosmos DB service. 
+        /// During the CosmosClient initialization the account information, including the available regions, is obtained from the <see cref="CosmosClient.Endpoint"/>.
+        /// Should the global endpoint become inaccessible, the CosmosClient will attempt to obtain the account information issuing requests to the custom endpoints
+        /// provided in the customAccountEndpoints list.
+        /// </summary>
+        /// <param name="customAccountEndpoints">An instance of <see cref="IEnumerable{T}"/> of Uri containing the custom private endpoints for the cosmos db account.</param>
+        /// <remarks>
+        ///  This function is optional and is recommended for implementation when a customer has configured one or more endpoints with a custom DNS
+        ///  hostname (instead of accountname-region.documents.azure.com) etc. for their Cosmos DB account.
+        /// </remarks>
+        /// <example>
+        /// The example below creates a new instance of <see cref="CosmosClientBuilder"/> with the regional endpoints.
+        /// <code language="c#">
+        /// <![CDATA[
+        /// CosmosClientBuilder cosmosClientBuilder = new CosmosClientBuilder(
+        ///     accountEndpoint: "https://testcosmos.documents.azure.com:443/",
+        ///     authKeyOrResourceToken: "SuperSecretKey")
+        /// .WithCustomAccountEndpoints(new HashSet<Uri>()
+        ///     { 
+        ///         new Uri("https://region-1.documents-test.windows-int.net:443/"),
+        ///         new Uri("https://region-2.documents-test.windows-int.net:443/") 
+        ///     });
+        /// CosmosClient client = cosmosClientBuilder.Build();
+        /// ]]>
+        /// </code>
+        /// </example>
+        /// <returns>The current <see cref="CosmosClientBuilder"/>.</returns>
+        /// <seealso cref="CosmosClientOptions.AccountInitializationCustomEndpoints"/>
+        public CosmosClientBuilder WithCustomAccountEndpoints(IEnumerable<Uri> customAccountEndpoints)
+        {
+            this.clientOptions.AccountInitializationCustomEndpoints = customAccountEndpoints;
+            return this;
+        }
+
         /// <summary>
         /// Limits the operations to the provided endpoint on the CosmosClientBuilder constructor.
         /// </summary>

diff --git a/Microsoft.Azure.Cosmos/src/GatewayAccountReader.cs b/Microsoft.Azure.Cosmos/src/GatewayAccountReader.cs
@@ -5,7 +5,6 @@
 namespace Microsoft.Azure.Cosmos
 {
     using System;
-    using System.Globalization;
     using System.Net.Http;
     using System.Threading;
     using System.Threading.Tasks;
@@ -89,6 +88,7 @@ public async Task<AccountProperties> InitializeReaderAsync()
             AccountProperties databaseAccount = await GlobalEndpointManager.GetDatabaseAccountFromAnyLocationsAsync(
                 defaultEndpoint: this.serviceEndpoint,
                 locations: this.connectionPolicy.PreferredLocations,
+                accountInitializationCustomEndpoints: this.connectionPolicy.AccountInitializationCustomEndpoints,
                 getDatabaseAccountFn: this.GetDatabaseAccountAsync,
                 cancellationToken: this.cancellationToken);