CosmosDb provider support for a current region option

src/EFCore.Cosmos/Infrastructure/CosmosDbContextOptionsBuilder.cs

@@ -30,6 +30,13 @@ public virtual CosmosDbContextOptionsBuilder ExecutionStrategy(
             [NotNull] Func<ExecutionStrategyDependencies, IExecutionStrategy> getExecutionStrategy)
             => WithOption(e => e.WithExecutionStrategyFactory(Check.NotNull(getExecutionStrategy, nameof(getExecutionStrategy))));
 
+        /// <summary>
+        /// Configures the context to use the provided Region.
+        /// </summary>
+        /// <param name="region">CosmosDB region name</param>
+        public virtual CosmosDbContextOptionsBuilder Region(string region)
+            => WithOption(e => e.WithRegion(Check.NotNull(region, nameof(region))));
+
         /// <summary>
         ///     Sets an option by cloning the extension used to store the settings. This ensures the builder
         ///     does not modify options that are already in use elsewhere.

src/EFCore.Cosmos/Infrastructure/Internal/CosmosDbOptionExtension.cs

@@ -18,6 +18,7 @@ public class CosmosDbOptionsExtension : IDbContextOptionsExtension
         private string _databaseName;
         private Func<ExecutionStrategyDependencies, IExecutionStrategy> _executionStrategyFactory;
         private string _logFragment;
+        private string _region;
 
         public CosmosDbOptionsExtension()
         {
@@ -29,6 +30,7 @@ protected CosmosDbOptionsExtension(CosmosDbOptionsExtension copyFrom)
             _authKeyOrResourceToken = copyFrom._authKeyOrResourceToken;
             _databaseName = copyFrom._databaseName;
             _executionStrategyFactory = copyFrom._executionStrategyFactory;
+            _region = copyFrom._region;
         }
 
         public virtual string ServiceEndPoint => _serviceEndPoint;
@@ -64,6 +66,17 @@ public virtual CosmosDbOptionsExtension WithDatabaseName(string database)
             return clone;
         }
 
+        public virtual string Region => _region;
+
+        public virtual CosmosDbOptionsExtension WithRegion(string region)
+        {
+            var clone = Clone();
+
+            clone._region = region;
+
+            return clone;
+        }
+
         /// <summary>
         ///     A factory for creating the default <see cref="IExecutionStrategy" />, or <c>null</c> if none has been
         ///     configured.

src/EFCore.Cosmos/Storage/Internal/CosmosClientWrapper.cs

@@ -33,6 +33,7 @@ public class CosmosClientWrapper : IDisposable
 
         private static readonly string _userAgent = " Microsoft.EntityFrameworkCore.Cosmos/" + ProductInfo.GetVersion();
         public static readonly JsonSerializer Serializer = new JsonSerializer();
+        private string _region;
 
         static CosmosClientWrapper()
         {
@@ -50,18 +51,31 @@ public CosmosClientWrapper(
             _databaseId = options.DatabaseName;
             _endPoint = options.ServiceEndPoint;
             _authKey = options.AuthKeyOrResourceToken;
+            _region = options.Region;
             _executionStrategyFactory = executionStrategyFactory;
             _commandLogger = commandLogger;
         }
 
         private CosmosClient Client =>
             _client
             ?? (_client = new CosmosClient(
-                new CosmosConfiguration(_endPoint, _authKey)
-                {
-                    UserAgentSuffix = _userAgent,
-                    ConnectionMode = ConnectionMode.Direct
-                }));
+                BuildCosmosConfiguration()));
+
+        private CosmosConfiguration BuildCosmosConfiguration()
+        {
+            var configuration = new CosmosConfiguration(_endPoint, _authKey)
+            {
+                UserAgentSuffix = _userAgent,
+                ConnectionMode = ConnectionMode.Direct
+            };
+
+            if (_region != null)
+            {
+                configuration = configuration.UseCurrentRegion(_region);
+            }
+
+            return configuration;
+        }
 
         public bool CreateDatabaseIfNotExists()
             => _executionStrategyFactory.Create().Execute(

test/EFCore.Cosmos.FunctionalTests/CosmosEndToEndTest.cs

@@ -4,6 +4,8 @@
 using System;
 using System.Linq;
 using System.Threading.Tasks;
+using Microsoft.Azure.Cosmos;
+using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
 using Microsoft.EntityFrameworkCore.Cosmos.TestUtilities;
 using Microsoft.EntityFrameworkCore.TestUtilities;
 using Microsoft.EntityFrameworkCore.TestUtilities.Xunit;
@@ -362,6 +364,56 @@ protected override void OnModelCreating(ModelBuilder modelBuilder)
             }
         }
 
+        [ConditionalFact]
+        public void Should_not_throw_if_specified_region_is_right()
+        {
+            var regionName = CosmosRegions.AustraliaCentral;
+
+            using (var testDatabase = CosmosTestStore.CreateInitialized(DatabaseName, o => o.Region(regionName)))
+            {
+                var options = Fixture.CreateOptions(testDatabase);
+
+                var customer = new Customer { Id = 42, Name = "Theon" };
+
+                using (var context = new CustomerContext(options))
+                {
+                    context.Database.EnsureCreated();
+
+                    context.Add(customer);
+
+                    context.SaveChanges();
+                }
+            }
+        }
+
+        [ConditionalFact]
+        public void Should_throw_if_specified_region_is_wrong()
+        {
+            var regionName = "FakeRegion";
+
+            Action a = () =>
+            {
+                using (var testDatabase = CosmosTestStore.CreateInitialized(DatabaseName, o => o.Region(regionName)))
+                {
+                    var options = Fixture.CreateOptions(testDatabase);
+
+                    var customer = new Customer {Id = 42, Name = "Theon"};
+
+                    using (var context = new CustomerContext(options))
+                    {
+                        context.Database.EnsureCreated();
+
+                        context.Add(customer);
+
+                        context.SaveChanges();
+                    }
+                }
+            };
+
+            var ex = Assert.Throws<ArgumentException>(a);
+            Assert.Equal("Current location is not a valid Azure region.", ex.Message);
+        }
+
         [ConditionalFact]
         public async Task Can_add_update_delete_end_to_end_with_conflicting_id()
         {

test/EFCore.Cosmos.FunctionalTests/TestUtilities/CosmosTestStore.cs

@@ -5,6 +5,8 @@
 using System.IO;
 using System.Reflection;
 using System.Threading.Tasks;
+using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure;
+using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
 using Microsoft.EntityFrameworkCore.Cosmos.Storage.Internal;
 using Microsoft.EntityFrameworkCore.Infrastructure;
 using Microsoft.EntityFrameworkCore.TestUtilities;
@@ -18,23 +20,23 @@ public class CosmosTestStore : TestStore
         private readonly TestStoreContext _storeContext;
         private readonly string _dataFilePath;
 
-        public static CosmosTestStore Create(string name) => new CosmosTestStore(name, shared: false);
+        public static CosmosTestStore Create(string name, Action<CosmosDbContextOptionsBuilder> extensionConfiguration = null) => new CosmosTestStore(name, shared: false, extensionConfiguration: extensionConfiguration);
 
-        public static CosmosTestStore CreateInitialized(string name)
-            => (CosmosTestStore)Create(name).Initialize(null, (Func<DbContext>)null, null);
+        public static CosmosTestStore CreateInitialized(string name, Action<CosmosDbContextOptionsBuilder> extensionConfiguration = null)
+            => (CosmosTestStore)Create(name, extensionConfiguration).Initialize(null, (Func<DbContext>)null, null);
 
         public static CosmosTestStore GetOrCreate(string name) => new CosmosTestStore(name);
 
         public static CosmosTestStore GetOrCreate(string name, string dataFilePath)
             => new CosmosTestStore(name, dataFilePath: dataFilePath);
 
-        private CosmosTestStore(string name, bool shared = true, string dataFilePath = null)
+        private CosmosTestStore(string name, bool shared = true, string dataFilePath = null, Action<CosmosDbContextOptionsBuilder> extensionConfiguration = null)
             : base(name, shared)
         {
             ConnectionUri = TestEnvironment.DefaultConnection;
             AuthToken = TestEnvironment.AuthToken;
 
-            _storeContext = new TestStoreContext(this);
+            _storeContext = new TestStoreContext(this, extensionConfiguration);
 
             if (dataFilePath != null)
             {
@@ -148,15 +150,19 @@ public override void Dispose()
         private class TestStoreContext : DbContext
         {
             private readonly CosmosTestStore _testStore;
+            private readonly Action<CosmosDbContextOptionsBuilder> _extensionConfiguration;
 
-            public TestStoreContext(CosmosTestStore testStore)
+            public TestStoreContext(CosmosTestStore testStore,
+                Action<CosmosDbContextOptionsBuilder> extensionConfiguration)
             {
                 _testStore = testStore;
+                _extensionConfiguration = extensionConfiguration;
             }
 
             protected override void OnConfiguring(DbContextOptionsBuilder optionsBuilder)
             {
-                optionsBuilder.UseCosmos(_testStore.ConnectionUri, _testStore.AuthToken, _testStore.Name);
+                var extensionConfiguration = _extensionConfiguration ?? (_ => { });
+                optionsBuilder.UseCosmos(_testStore.ConnectionUri, _testStore.AuthToken, _testStore.Name, extensionConfiguration);
             }
         }
     }

test/EFCore.Cosmos.Tests/Configuration/CosmosDbContextOptionsExtensionsTests.cs

@@ -0,0 +1,44 @@
+using Microsoft.Azure.Cosmos;
+using Microsoft.EntityFrameworkCore.Cosmos.Infrastructure.Internal;
+using Xunit;
+
+namespace Microsoft.EntityFrameworkCore.Cosmos.Configuration
+{
+    public class CosmosDbContextOptionsExtensionsTests
+    {
+        [Fact]
+        public void Can_create_options_with_specified_region()
+        {
+            var regionName = CosmosRegions.EastAsia;
+            var options = new DbContextOptionsBuilder().UseCosmos(
+                "serviceEndPoint",
+                "authKeyOrResourceToken",
+                "databaseName",
+                o => { o.Region(regionName); });
+
+            var extension = options
+                .Options.FindExtension<CosmosDbOptionsExtension>();
+
+            Assert.Equal(regionName, extension.Region);
+        }
+
+        /// <summary>
+        /// The region will be checked by the cosmosdb sdk, because the region list is not constant
+        /// </summary>
+        [Fact]
+        public void Can_create_options_with_wrong_region()
+        {
+            var regionName = "FakeRegion";
+            var options = new DbContextOptionsBuilder().UseCosmos(
+                "serviceEndPoint",
+                "authKeyOrResourceToken",
+                "databaseName",
+                o => { o.Region(regionName); });
+
+            var extension = options
+                .Options.FindExtension<CosmosDbOptionsExtension>();
+
+            Assert.Equal(regionName, extension.Region);
+        }
+    }
+}