fix: Adding retry support

Adding options to create an AsyncDbSession with retry support
for high latency connection usage

src/Codibre.MSSqlSession/Impl/AsyncDbSession.cs

@@ -18,10 +18,9 @@ internal sealed class AsyncDbSession : IAsyncDbSession
     private static readonly Task<IDisposable?> _nullTask = Task.FromResult((IDisposable?)null);
     private static readonly object _updateToken = new();
     private static DateTime _omitLogDeadline = DateTime.Now;
-    private readonly bool _customPool;
     private readonly AsyncLocal<AsyncDbStorage?> _asyncStorage = new();
     private readonly ILogger<IAsyncDbSession> _logger;
-    private readonly string _connectionString;
+    private readonly AsyncSqlOptions _options;
     public DbTransaction? Transaction
     {
         get => _asyncStorage.Value?.Transaction;
@@ -64,11 +63,20 @@ public AsyncDbSession(
         string connectionString,
         bool customPool,
         ILogger<AsyncDbSession> logger
+    ) : this(logger, new AsyncSqlOptions
+    {
+        ConnectionString = connectionString,
+        CustomPool = customPool
+    })
+    { }
+
+    public AsyncDbSession(
+        ILogger<AsyncDbSession> logger,
+        AsyncSqlOptions options
     )
     {
         _logger = logger;
-        _customPool = customPool;
-        _connectionString = connectionString;
+        _options = options;
     }
 
     public async ValueTask Clear()
@@ -142,10 +150,9 @@ private static void UpdateOmitLogDeadline(DateTime now)
     {
         var disposer = new CrossedDisposer(this);
         var connection = new AsyncSqlConnection(
-            _connectionString,
+            _options,
             disposer,
             _asyncStorage,
-            _customPool,
             _logger
         );
         _asyncStorage.Value = new AsyncDbStorage(connection);

src/Codibre.MSSqlSession/Impl/AsyncSqlConnection.cs

@@ -19,27 +19,26 @@ namespace Codibre.MSSqlSession.Impl;
 // Threads
 internal class AsyncSqlConnection : DbConnection
 {
-    private string _connectionString;
+    private readonly AsyncSqlOptions _options;
     private string _dataSource;
     private bool _opened = false;
     private SqlConnection? _connection;
     private object? _pooledConnection;
     private readonly CrossedDisposer _disposer;
-    private readonly bool _customPool;
     private readonly AsyncLocal<AsyncDbStorage?> _asyncDbStorage;
     private readonly ILogger _logger;
     public override string ConnectionString
     {
-        get => _connectionString;
+        get => _options.ConnectionString;
         set
         {
-            if (_connectionString != value)
+            if (_options.ConnectionString != value)
             {
                 _connection = null;
                 ReleaseConnection();
             }
-            _connectionString = value;
-            _dataSource = GetDataSource(_connectionString);
+            _options.ConnectionString = value;
+            _dataSource = GetDataSource(_options.ConnectionString);
         }
     }
 
@@ -58,18 +57,16 @@ private static string GetDataSource(string connectionString)
     public override ConnectionState State => _opened && _connection is not null ? _connection.State : ConnectionState.Closed;
 
     internal AsyncSqlConnection(
-        string connectionString,
+        AsyncSqlOptions options,
         CrossedDisposer disposer,
         AsyncLocal<AsyncDbStorage?> currentTransaction,
-        bool customPool,
         ILogger logger
         )
     {
-        _connectionString = connectionString;
+        _options = options;
         _disposer = disposer;
         _asyncDbStorage = currentTransaction;
-        _customPool = customPool;
-        _dataSource = GetDataSource(connectionString);
+        _dataSource = GetDataSource(_options.ConnectionString);
         _logger = logger;
     }
 
@@ -79,7 +76,7 @@ private SqlConnection Connection
         {
             if (_connection is null)
             {
-                if (_customPool) (_connection, _pooledConnection) = SqlConnectionFactory.GetConnection(ConnectionString, _logger);
+                if (_options.CustomPool) (_connection, _pooledConnection) = SqlConnectionFactory.GetConnection(ConnectionString, _logger);
                 else _connection = new SqlConnection(ConnectionString);
             }
             return _connection;
@@ -145,7 +142,9 @@ protected override DbTransaction BeginDbTransaction(IsolationLevel isolationLeve
 
     protected override DbCommand CreateDbCommand()
     {
-        var command = new AsyncDbCommand(Connection.CreateCommand(), this);
+        var sqlCommand = Connection.CreateCommand();
+        if (_options.RetryLogicBaseProvider is not null) sqlCommand.RetryLogicProvider = _options.RetryLogicBaseProvider;
+        var command = new AsyncDbCommand(sqlCommand, this);
         var value = _asyncDbStorage.Value;
         if (value != null)
         {

src/Codibre.MSSqlSession/Impl/AsyncSqlOptions.cs

@@ -0,0 +1,10 @@
+using Microsoft.Data.SqlClient;
+
+namespace Codibre.MSSqlSession.Impl;
+
+public class AsyncSqlOptions
+{
+    public string ConnectionString { get; set; }
+    public bool CustomPool { get; set; }
+    public SqlRetryLogicBaseProvider? RetryLogicBaseProvider { get; set; }
+}

src/Codibre.MSSqlSession/Impl/BatchQuery.Query.cs

@@ -0,0 +1,70 @@
+using Codibre.MSSqlSession.Extensions;
+using Codibre.MSSqlSession.Impl.Utils;
+
+namespace Codibre.MSSqlSession.Impl;
+
+internal partial class BatchQuery : IBatchQuery
+{
+    private readonly IAsyncDbSession _session;
+    private readonly IScriptBuilder _builder;
+    private readonly HookStorage _hookStorage;
+    internal BatchQuery(IAsyncDbSession session)
+    {
+        _session = session;
+        _builder = new ScriptBuilder(session.Connection);
+        _hookStorage = new(
+            new(),
+            _builder
+        );
+    }
+
+    public string Sql => _builder.Sql;
+    public int QueryCount => _builder.QueryCount;
+    public int ParamCount => _builder.ParamCount;
+
+    public void AddNoResultScript(FormattableString builtScript) => _builder.Add(builtScript);
+
+    public IResultHook<T> QueryFirstHook<T>(FormattableString builtScript)
+    {
+        var token = new object();
+        return _hookStorage.AddRequiredHook<T>(builtScript, token,
+            async (reader, setResult) => setResult(token, await reader.ReadFirstAsync<T>())
+        );
+    }
+    public IResultHook<T?> QueryFirstOrDefaultHook<T>(FormattableString builtScript)
+    {
+        var token = new object();
+        return _hookStorage.AddNullableHook<T?>(builtScript, token,
+            async (reader, setResult) => setResult(token, await reader.ReadFirstOrDefaultAsync<T>())
+        );
+    }
+
+    public IResultHook<IEnumerable<T>> QueryHook<T>(FormattableString builtScript)
+    {
+        var token = new object();
+        return _hookStorage.AddRequiredHook<IEnumerable<T>>(builtScript, token,
+            async (reader, setResult) => setResult(token, await reader.ReadAsync<T>())
+        );
+    }
+
+    public Task RunQueries(TimeSpan? customTimeout = null)
+        => _hookStorage.RunQueries(customTimeout);
+
+    public Task Execute(TimeSpan? customTimeout = null)
+        => _hookStorage.Execute(customTimeout);
+
+    public void Clear() => _hookStorage.Clear();
+
+    private IAsyncEnumerable<IList<KeyValuePair<TInput, TOutput>>> InternalPrepareEnumerable<TInput, TOutput>(
+        IEnumerable<TInput> enumerable,
+        Func<TInput, IBatchQuery, ValueTask<TOutput>> PreRunQuery
+    ) => _hookStorage
+            .PrepareEnumerable(enumerable, (current) => PreRunQuery(current, this))
+            .OnStartAsync(async () => _session.ConnectionAcquired ? null : await _session.StartSession());
+
+    public IAsyncEnumerable<KeyValuePair<TInput, TOutput>> PrepareEnumerable<TInput, TOutput>(
+        IEnumerable<TInput> enumerable,
+        Func<TInput, IBatchQuery, ValueTask<TOutput>> PreRunQuery
+    ) => InternalPrepareEnumerable(enumerable, PreRunQuery)
+            .SelectMany(x => x.ToAsyncEnumerable());
+}

src/Codibre.MSSqlSession/Impl/BatchQuery.cs → src/Codibre.MSSqlSession/Impl/BatchQuery.Transaction.cs

@@ -1,76 +1,12 @@
-using Codibre.MSSqlSession.Extensions;
-using Codibre.MSSqlSession.Impl.Utils;
+using Codibre.MSSqlSession.Impl.Utils;
 
 namespace Codibre.MSSqlSession.Impl;
 
-internal class BatchQuery : IBatchQuery
+internal partial class BatchQuery
 {
     private static readonly FormattableString _beginTran = $"BEGIN TRAN;";
     private static readonly FormattableString _commitTran = $"COMMIT TRAN";
-    private readonly IAsyncDbSession _session;
-    private readonly IScriptBuilder _builder;
     private readonly TransactionControl _transactionControl = new();
-    private readonly HookStorage _hookStorage;
-    internal BatchQuery(IAsyncDbSession session)
-    {
-        _session = session;
-        _builder = new ScriptBuilder(session.Connection);
-        _hookStorage = new(
-            new(),
-            _builder
-        );
-    }
-
-    public string Sql => _builder.Sql;
-    public int QueryCount => _builder.QueryCount;
-    public int ParamCount => _builder.ParamCount;
-
-    public void AddNoResultScript(FormattableString builtScript) => _builder.Add(builtScript);
-
-    public IResultHook<T> QueryFirstHook<T>(FormattableString builtScript)
-    {
-        var token = new object();
-        return _hookStorage.AddRequiredHook<T>(builtScript, token,
-            async (reader, setResult) => setResult(token, await reader.ReadFirstAsync<T>())
-        );
-    }
-    public IResultHook<T?> QueryFirstOrDefaultHook<T>(FormattableString builtScript)
-    {
-        var token = new object();
-        return _hookStorage.AddNullableHook<T?>(builtScript, token,
-            async (reader, setResult) => setResult(token, await reader.ReadFirstOrDefaultAsync<T>())
-        );
-    }
-
-    public IResultHook<IEnumerable<T>> QueryHook<T>(FormattableString builtScript)
-    {
-        var token = new object();
-        return _hookStorage.AddRequiredHook<IEnumerable<T>>(builtScript, token,
-            async (reader, setResult) => setResult(token, await reader.ReadAsync<T>())
-        );
-    }
-
-    public Task RunQueries(TimeSpan? customTimeout = null)
-        => _hookStorage.RunQueries(customTimeout);
-
-    public Task Execute(TimeSpan? customTimeout = null)
-        => _hookStorage.Execute(customTimeout);
-
-    public void Clear() => _hookStorage.Clear();
-
-    private IAsyncEnumerable<IList<KeyValuePair<TInput, TOutput>>> InternalPrepareEnumerable<TInput, TOutput>(
-        IEnumerable<TInput> enumerable,
-        Func<TInput, IBatchQuery, ValueTask<TOutput>> PreRunQuery
-    ) => _hookStorage
-            .PrepareEnumerable(enumerable, (current) => PreRunQuery(current, this))
-            .OnStartAsync(async () => _session.ConnectionAcquired ? null : await _session.StartSession());
-
-    public IAsyncEnumerable<KeyValuePair<TInput, TOutput>> PrepareEnumerable<TInput, TOutput>(
-        IEnumerable<TInput> enumerable,
-        Func<TInput, IBatchQuery, ValueTask<TOutput>> PreRunQuery
-    ) => InternalPrepareEnumerable(enumerable, PreRunQuery)
-            .SelectMany(x => x.ToAsyncEnumerable());
-
     private async ValueTask InternalFlushTransaction()
     {
         if (!_transactionControl.Open)