CosmosClientOptions: Adds Private Custom Account Endpoints

Microsoft.Azure.Cosmos/src/ConnectionPolicy.cs

@@ -29,6 +29,7 @@ internal sealed class ConnectionPolicy
 
         private Protocol connectionProtocol;
         private ObservableCollection<string> preferredLocations;
+        private ObservableCollection<string> regionalEndpoints;
 
         /// <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.regionalEndpoints = new ObservableCollection<string>();
             this.EnableEndpointDiscovery = true;
             this.MaxConnectionLimit = defaultMaxConcurrentConnectionLimit;
             this.RetryOptions = new RetryOptions();
@@ -90,6 +92,27 @@ public void SetPreferredLocations(IReadOnlyList<string> regions)
             }
         }
 
+        /// <summary>
+        /// Sets the regional endpoints required to fetch account information from
+        /// private domain names.
+        /// </summary>
+        /// <param name="regionalEndpoints">An instance of <see cref="ISet{T}"/> containing the regional endpoints
+        /// provided by the customer.</param>
+        public void SetRegionalEndpoints(
+            ISet<string> regionalEndpoints)
+        {
+            if (regionalEndpoints == null)
+            {
+                throw new ArgumentNullException(nameof(regionalEndpoints));
+            }
+
+            this.regionalEndpoints.Clear();
+            foreach (string endpoint in regionalEndpoints)
+            {
+                this.regionalEndpoints.Add(endpoint);
+            }
+        }
+
         /// <summary>
         /// Gets or sets the maximum number of concurrent fanout requests sent to the Azure Cosmos DB service.
         /// </summary>
@@ -270,6 +293,29 @@ public Collection<string> PreferredLocations
             }
         }
 
+        /// <summary>
+        /// Gets and sets the preferred locations (regions) for geo-replicated database accounts in the Azure Cosmos DB service.
+        /// For example, "East US" as the preferred location.
+        /// </summary>
+        /// <remarks>
+        /// <para>
+        /// When <see cref="EnableEndpointDiscovery"/> is true and the value of this property is non-empty,
+        /// the SDK uses the locations in the collection in the order they are specified to perform operations,
+        /// otherwise if the value of this property is not specified,
+        /// the SDK uses the write region as the preferred location for all operations.
+        /// </para>
+        /// <para>
+        /// If <see cref="EnableEndpointDiscovery"/> is set to false, the value of this property is ignored.
+        /// </para>
+        /// </remarks>
+        public Collection<string> RegionalEndpoints
+        {
+            get
+            {
+                return this.regionalEndpoints;
+            }
+        }
+
         /// <summary>
         /// Gets or sets the flag to enable endpoint discovery for geo-replicated database accounts in the Azure Cosmos DB service.
         /// </summary>

Microsoft.Azure.Cosmos/src/CosmosClientOptions.cs

@@ -190,6 +190,37 @@ 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 regional 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 regional endpoints provided in <see cref="RegionalEndpoints"/>.
+        /// </para>
+        /// <para>
+        /// Nevertheless, this parameter remains optional and is recommended for implementation when a customer has configured a private endpoint 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()
+        /// {
+        ///     RegionalEndpoints = new HashSet<string>(){ "custom.p-1.documents.azure.com", "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 ISet<string> RegionalEndpoints { get; set; }
 
         /// <summary>
         /// Get or set the maximum number of concurrent connections allowed for the target
@@ -794,6 +825,11 @@ internal virtual ConnectionPolicy GetConnectionPolicy(int clientId)
                 List<string> mappedRegions = this.ApplicationPreferredRegions.Select(s => mapper.GetCosmosDBRegionName(s)).ToList();
 
                 connectionPolicy.SetPreferredLocations(mappedRegions);
+            }
+
+            if (this.RegionalEndpoints != null)
+            {
+                connectionPolicy.SetRegionalEndpoints(this.RegionalEndpoints);
             }
 
             if (this.MaxRetryAttemptsOnRateLimitedRequests != null)

Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs

@@ -301,6 +301,36 @@ public CosmosClientBuilder WithApplicationPreferredRegions(IReadOnlyList<string>
             return this;
         }
 
+        /// <summary>
+        /// Sets the regional private endpoints 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 regional endpoints
+        /// provided in the regionalEndpoints set.
+        /// </summary>
+        /// <param name="regionalEndpoints">A set of string containing the regional 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 private endpoints 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")
+        /// .WithRegionalEndpoints(new HashSet<string>() { "https://region-1.documents-test.windows-int.net:443/", "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.RegionalEndpoints"/>
+        public CosmosClientBuilder WithRegionalEndpoints(ISet<string> regionalEndpoints)
+        {
+            this.clientOptions.RegionalEndpoints = regionalEndpoints;
+            return this;
+        }
+
         /// <summary>
         /// Limits the operations to the provided endpoint on the CosmosClientBuilder constructor.
         /// </summary>

Microsoft.Azure.Cosmos/src/GatewayAccountReader.cs

@@ -5,10 +5,11 @@
 namespace Microsoft.Azure.Cosmos
 {
     using System;
-    using System.Globalization;
+    using System.Collections.Generic;
     using System.Net.Http;
     using System.Threading;
     using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Core.Trace;
     using Microsoft.Azure.Cosmos.Resource.CosmosExceptions;
     using Microsoft.Azure.Cosmos.Routing;
     using Microsoft.Azure.Cosmos.Tracing;
@@ -22,6 +23,7 @@ internal sealed class GatewayAccountReader
         private readonly AuthorizationTokenProvider cosmosAuthorization;
         private readonly CosmosHttpClient httpClient;
         private readonly Uri serviceEndpoint;
+        private readonly IEnumerator<Uri> serviceEndpointEnumerator;
         private readonly CancellationToken cancellationToken;
 
         // Backlog: Auth abstractions are spilling through. 4 arguments for this CTOR are result of it.
@@ -36,6 +38,24 @@ public GatewayAccountReader(Uri serviceEndpoint,
             this.cosmosAuthorization = cosmosAuthorization ?? throw new ArgumentNullException(nameof(AuthorizationTokenProvider));
             this.connectionPolicy = connectionPolicy;
             this.cancellationToken = cancellationToken;
+
+            List<Uri> serviceEndpoints = new ()
+            {
+                serviceEndpoint
+            };
+
+            if (this.connectionPolicy.RegionalEndpoints != null
+                && this.connectionPolicy.RegionalEndpoints.Count > 0)
+            {
+                foreach (string regionalEndpoint in this.connectionPolicy.RegionalEndpoints)
+                {
+                    // Add all of the regional endpoints to the service endpoints list.
+                    serviceEndpoints.Add(
+                        new Uri(regionalEndpoint));
+                }
+            }
+
+            this.serviceEndpointEnumerator = serviceEndpoints.GetEnumerator();
         }
 
         private async Task<AccountProperties> GetDatabaseAccountAsync(Uri serviceEndpoint)
@@ -86,13 +106,31 @@ await this.cosmosAuthorization.AddAuthorizationHeaderAsync(
 
         public async Task<AccountProperties> InitializeReaderAsync()
         {
-            AccountProperties databaseAccount = await GlobalEndpointManager.GetDatabaseAccountFromAnyLocationsAsync(
-                defaultEndpoint: this.serviceEndpoint,
-                locations: this.connectionPolicy.PreferredLocations,
-                getDatabaseAccountFn: this.GetDatabaseAccountAsync,
-                cancellationToken: this.cancellationToken);
+            int attemptCounter = 1;
+            List<Exception> exceptionList = new ();
+
+            while (this.serviceEndpointEnumerator.MoveNext())
+            {
+                try
+                {
+                    AccountProperties databaseAccount = await GlobalEndpointManager.GetDatabaseAccountFromAnyLocationsAsync(
+                        defaultEndpoint: this.serviceEndpointEnumerator.Current,
+                        locations: this.connectionPolicy.PreferredLocations,
+                        getDatabaseAccountFn: this.GetDatabaseAccountAsync,
+                        cancellationToken: this.cancellationToken);
+
+                    return databaseAccount;
+                }
+                catch (Exception ex)
+                {
+                    DefaultTrace.TraceWarning("Attempt: {0}, Exception occurred while fetching account details: {1}", attemptCounter++, ex.Message);
+                    exceptionList.Add(ex);
+                }
+            }
 
-            return databaseAccount;
+            throw exceptionList.Count > 1
+                ? new AggregateException("Unable to get account information.", exceptionList)
+                : exceptionList[0];
         }
     }
 }

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/Contracts/DotNetSDKAPI.json

@@ -2782,6 +2782,18 @@
           ],
           "MethodInfo": "System.Collections.Generic.IReadOnlyList`1[System.String] get_ApplicationPreferredRegions();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
         },
+        "System.Collections.Generic.ISet`1[System.String] get_RegionalEndpoints()[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
+          "Type": "Method",
+          "Attributes": [
+            "CompilerGeneratedAttribute"
+          ],
+          "MethodInfo": "System.Collections.Generic.ISet`1[System.String] get_RegionalEndpoints();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
+        "System.Collections.Generic.ISet`1[System.String] RegionalEndpoints": {
+          "Type": "Property",
+          "Attributes": [],
+          "MethodInfo": "System.Collections.Generic.ISet`1[System.String] RegionalEndpoints;CanRead:True;CanWrite:True;System.Collections.Generic.ISet`1[System.String] get_RegionalEndpoints();IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;Void set_RegionalEndpoints(System.Collections.Generic.ISet`1[System.String]);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
         "System.Collections.ObjectModel.Collection`1[Microsoft.Azure.Cosmos.RequestHandler] CustomHandlers[Newtonsoft.Json.JsonConverterAttribute(typeof(Microsoft.Azure.Cosmos.CosmosClientOptions+ClientOptionJsonConverter))]": {
           "Type": "Property",
           "Attributes": [
@@ -3096,6 +3108,13 @@
           "Attributes": [],
           "MethodInfo": "Void set_PortReuseMode(System.Nullable`1[Microsoft.Azure.Cosmos.PortReuseMode]);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
         },
+        "Void set_RegionalEndpoints(System.Collections.Generic.ISet`1[System.String])[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
+          "Type": "Method",
+          "Attributes": [
+            "CompilerGeneratedAttribute"
+          ],
+          "MethodInfo": "Void set_RegionalEndpoints(System.Collections.Generic.ISet`1[System.String]);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
         "Void set_RequestTimeout(System.TimeSpan)[System.Runtime.CompilerServices.CompilerGeneratedAttribute()]": {
           "Type": "Method",
           "Attributes": [
@@ -4599,6 +4618,11 @@
           "Attributes": [],
           "MethodInfo": "Microsoft.Azure.Cosmos.Fluent.CosmosClientBuilder WithLimitToEndpoint(Boolean);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
         },
+        "Microsoft.Azure.Cosmos.Fluent.CosmosClientBuilder WithRegionalEndpoints(System.Collections.Generic.ISet`1[System.String])": {
+          "Type": "Method",
+          "Attributes": [],
+          "MethodInfo": "Microsoft.Azure.Cosmos.Fluent.CosmosClientBuilder WithRegionalEndpoints(System.Collections.Generic.ISet`1[System.String]);IsAbstract:False;IsStatic:False;IsVirtual:False;IsGenericMethod:False;IsConstructor:False;IsFinal:False;"
+        },
         "Microsoft.Azure.Cosmos.Fluent.CosmosClientBuilder WithRequestTimeout(System.TimeSpan)": {
           "Type": "Method",
           "Attributes": [],

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosClientOptionsUnitTests.cs

@@ -156,7 +156,9 @@ public void VerifyCosmosConfigurationPropertiesGetUpdated()
             Assert.IsFalse(policy.EnablePartitionLevelFailover);
             Assert.IsTrue(clientOptions.EnableAdvancedReplicaSelectionForTcp.Value);
 
-            IReadOnlyList<string> preferredLocations = new List<string>() { Regions.AustraliaCentral, Regions.AustraliaCentral2 };
+            IReadOnlyList<string> preferredLocations = new List<string>() { Regions.AustraliaCentral, Regions.AustraliaCentral2 };
+            ISet<string> regionalEndpoints = new HashSet<string>() { "https://testfed2.documents-test.windows-int.net:443/", "https://testfed4.documents-test.windows-int.net:443/" };
+
             //Verify Direct Mode settings
             cosmosClientBuilder = new CosmosClientBuilder(
                 accountEndpoint: endpoint,
@@ -168,7 +170,8 @@ public void VerifyCosmosConfigurationPropertiesGetUpdated()
                 maxTcpConnectionsPerEndpoint,
                 portReuseMode,
                 enableTcpConnectionEndpointRediscovery)
-                .WithApplicationPreferredRegions(preferredLocations)
+                .WithApplicationPreferredRegions(preferredLocations)
+                .WithRegionalEndpoints(regionalEndpoints)
                 .WithClientTelemetryOptions(new CosmosClientTelemetryOptions()
                 {
                     DisableDistributedTracing = false,
@@ -188,7 +191,8 @@ public void VerifyCosmosConfigurationPropertiesGetUpdated()
             Assert.AreEqual(maxTcpConnectionsPerEndpoint, clientOptions.MaxTcpConnectionsPerEndpoint);
             Assert.AreEqual(portReuseMode, clientOptions.PortReuseMode);
             Assert.IsTrue(clientOptions.EnableTcpConnectionEndpointRediscovery);
-            CollectionAssert.AreEqual(preferredLocations.ToArray(), clientOptions.ApplicationPreferredRegions.ToArray());
+            CollectionAssert.AreEqual(preferredLocations.ToArray(), clientOptions.ApplicationPreferredRegions.ToArray());
+            CollectionAssert.AreEqual(regionalEndpoints.ToArray(), clientOptions.RegionalEndpoints.ToArray());
             Assert.AreEqual(TimeSpan.FromMilliseconds(100), clientOptions.CosmosClientTelemetryOptions.CosmosThresholdOptions.PointOperationLatencyThreshold);
             Assert.AreEqual(TimeSpan.FromMilliseconds(100), clientOptions.CosmosClientTelemetryOptions.CosmosThresholdOptions.NonPointOperationLatencyThreshold);
             Assert.IsFalse(clientOptions.CosmosClientTelemetryOptions.DisableDistributedTracing);
@@ -320,6 +324,13 @@ public void CosmosClientOptions_WhenPartitionLevelFailoverEnabledAndPreferredReg
                         Regions.NorthCentralUS,
                         Regions.WestUS,
                         Regions.EastAsia,
+                        })
+                    .WithRegionalEndpoints(
+                        new HashSet<string>()
+                        {
+                        "https://testfed2.documents-test.windows-int.net:443/",
+                        "https://testfed3.documents-test.windows-int.net:443/",
+                        "https://testfed4.documents-test.windows-int.net:443/",
                         });
 
                 CosmosClientOptions clientOptions = cosmosClientBuilder.Build().ClientOptions;
@@ -337,7 +348,8 @@ public void CosmosClientOptions_WhenPartitionLevelFailoverEnabledAndPreferredReg
                 Assert.IsFalse(clientOptions.AllowBulkExecution);
                 Assert.AreEqual(consistencyLevel, clientOptions.ConsistencyLevel);
                 Assert.IsTrue(clientOptions.EnablePartitionLevelFailover);
-                Assert.IsNotNull(clientOptions.ApplicationPreferredRegions);
+                Assert.IsNotNull(clientOptions.ApplicationPreferredRegions);
+                Assert.IsNotNull(clientOptions.RegionalEndpoints);
             }
             finally
             {

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/CosmosClientTests.cs

@@ -138,7 +138,7 @@ public async Task InvalidKey_ExceptionFullStacktrace(string endpoint, string key
             }
             catch (Exception ex)
             {
-                Assert.IsTrue(ex.StackTrace.Contains("GatewayAccountReader.GetDatabaseAccountAsync"), ex.StackTrace);
+                Assert.IsTrue(ex.StackTrace.Contains("GatewayAccountReader.InitializeReaderAsync"), ex.StackTrace);
             }
         }