读写分离自动切换到写库链接在操作写库后

Author: xuejmnet

src/ShardingCore/Core/ShardingConfigurations/ShardingConfigOptions.cs

@@ -9,6 +9,10 @@ namespace ShardingCore.Core.ShardingConfigurations
 {
     public class ShardingConfigOptions
     {
+        /// <summary>
+        /// 写操作数据库后自动使用写库链接防止读库链接未同步无法查询到数据
+        /// </summary>
+        public bool AutoUseWriteConnectionStringAfterWriteDb { get; set; } = false;
         /// <summary>
         /// 当查询遇到没有路由被命中时是否抛出错误
         /// </summary>

src/ShardingCore/Sharding/ShardingDbContextExecutors/ShardingDbContextExecutor.cs

@@ -13,6 +13,7 @@
 using ShardingCore.Core.VirtualRoutes.TableRoutes.RouteTails.Abstractions;
 using ShardingCore.Core.DbContextCreator;
 using ShardingCore.Core.RuntimeContexts;
+using ShardingCore.Core.ShardingConfigurations;
 using ShardingCore.Core.VirtualRoutes.Abstractions;
 using ShardingCore.Extensions;
 using ShardingCore.Sharding.Abstractions;
@@ -38,6 +39,7 @@ public class ShardingDbContextExecutor : IShardingDbContextExecutor
         //private readonly ConcurrentDictionary<string, ConcurrentDictionary<string, DbContext>> _dbContextCaches = new ConcurrentDictionary<string, ConcurrentDictionary<string, DbContext>>();
         private readonly ConcurrentDictionary<string, IDataSourceDbContext> _dbContextCaches = new ConcurrentDictionary<string, IDataSourceDbContext>();
         private readonly IShardingRuntimeContext _shardingRuntimeContext;
+        private readonly ShardingConfigOptions _shardingConfigOptions;
         private readonly IVirtualDataSource _virtualDataSource;
         private readonly IDataSourceRouteManager _dataSourceRouteManager;
         private readonly ITableRouteManager _tableRouteManager;
@@ -66,6 +68,7 @@ public ShardingDbContextExecutor(DbContext shardingDbContext)
             //初始化
             _shardingRuntimeContext = shardingDbContext.GetShardingRuntimeContext();
             _shardingRuntimeContext.GetOrCreateShardingRuntimeModel(shardingDbContext);
+            _shardingConfigOptions = _shardingRuntimeContext.GetShardingConfigOptions();
             _virtualDataSource = _shardingRuntimeContext.GetVirtualDataSource();
             _dataSourceRouteManager = _shardingRuntimeContext.GetDataSourceRouteManager();
             _tableRouteManager = _shardingRuntimeContext.GetTableRouteManager();
@@ -154,6 +157,7 @@ private string GetTableTail<TEntity>(string dataSourceName,TEntity entity) where
                 i += await dbContextCache.Value.SaveChangesAsync(acceptAllChangesOnSuccess, cancellationToken);
             }
 
+            AutoUseWriteConnectionString();
             return i;
         }
 
@@ -165,6 +169,7 @@ public int SaveChanges(bool acceptAllChangesOnSuccess)
                 i += dbContextCache.Value.SaveChanges(acceptAllChangesOnSuccess);
             }
 
+            AutoUseWriteConnectionString();
             return i;
         }
 
@@ -182,6 +187,8 @@ public void Rollback()
             {
                 dbContextCache.Value.Rollback();
             }
+
+            AutoUseWriteConnectionString();
         }
 
         public void Commit()
@@ -190,6 +197,8 @@ public void Commit()
             {
                 dbContextCache.Value.Commit(_dbContextCaches.Count);
             }
+
+            AutoUseWriteConnectionString();
         }
 
         public IDictionary<string, IDataSourceDbContext> GetCurrentDbContexts()
@@ -215,6 +224,8 @@ public void Dispose()
             {
                 await dbContextCache.Value.RollbackAsync(cancellationToken);
             }
+
+            AutoUseWriteConnectionString();
         }
 
         public async Task CommitAsync(CancellationToken cancellationToken = new CancellationToken())
@@ -223,6 +234,8 @@ public void Dispose()
             {
                 await dbContextCache.Value.CommitAsync(_dbContextCaches.Count, cancellationToken);
             }
+
+            AutoUseWriteConnectionString();
         }
         public async ValueTask DisposeAsync()
         {
@@ -232,5 +245,16 @@ public async ValueTask DisposeAsync()
             }
         }
 #endif
+
+        /// <summary>
+        /// 自动切换成写库连接
+        /// </summary>
+        private void AutoUseWriteConnectionString()
+        {
+            if (_shardingConfigOptions.AutoUseWriteConnectionStringAfterWriteDb)
+            {
+                ((IShardingDbContext)_shardingDbContext).ReadWriteSeparationWriteOnly();
+            }
+        }
     }
 }