Added new method with cancellation token

Agoda.Frameworks.DB.Tests/DbRepositoryTest.cs

@@ -1,6 +1,8 @@
 using System;
 using System.Collections.Generic;
 using System.Data;
+using System.Data.SqlClient;
+using System.Threading;
 using System.Threading.Tasks;
 using Dapper;
 using Moq;
@@ -664,6 +666,27 @@
             Assert.IsNotNull(_onQueryCompleteEvents[0].Error);
             Assert.IsNotNull(_onQueryCompleteEvents[1].Error);
         }
+
+        [Test]
+        public async Task ExecuteReaderAsync_CancellationToken_Cancelled()
+        {
+            var cancellationToken = new CancellationTokenSource(TimeSpan.FromDays(1));
+            cancellationToken.Cancel();
+            var maxAttemptCount = 2;
+            Assert.ThrowsAsync<TaskCanceledException>(async () => await _db.ExecuteReaderAsync<string>("mobile_ro", "db.v1.sp_foo", 1,
+                maxAttemptCount,
+                cancellationToken.Token,
+                new IDbDataParameter[]
+                {
+                    new SqlParameter("@param1", "value1"),
+                    new SqlParameter("@param2", "value2")
+
+                }, reader => { return Task.FromResult("");})
+            );
+
+            _dbResources.Verify(x => x.ChooseDb("mobile_ro").UpdateWeight(It.IsAny<string>(), false), Times.Exactly(maxAttemptCount));
+        }
+
 
         protected class FakeStoredProc : IStoredProc<string, string>
         {

Agoda.Frameworks.DB/DbRepository.cs

@@ -3,6 +3,7 @@
 using System.Data;
 using System.Data.SqlClient;
 using System.Diagnostics;
+using System.Threading;
 using System.Threading.Tasks;
 using Agoda.Frameworks.LoadBalancing;
 using Dapper;
@@ -296,6 +297,69 @@ public Task<T> ExecuteReaderAsync<T>(
                 }
             }, ShouldRetry(maxAttemptCount), RaiseOnError);
         }
+
+        public Task<T> ExecuteReaderAsync<T>(
+            string database,
+            string storedProc,
+            int timeoutSecs,
+            int maxAttemptCount,
+            CancellationToken token,
+            IDbDataParameter[] parameters,
+            Func<SqlDataReader, Task<T>> callback)
+        {
+            return _dbResources.ChooseDb(database).ExecuteAsync(async (connectionStr, _) =>
+            {
+                var stopwatch = Stopwatch.StartNew();
+                Exception error = null;
+                try
+                {
+                    using (var connection = _generateConnection(connectionStr))
+                    {
+                        if (connection is SqlConnection sqlConn)
+                        {
+                            await sqlConn.OpenAsync();
+                        }
+                        else
+                        {
+                            connection.Open();
+                        }
+                        SqlCommand sqlCommand = null;
+                        try
+                        {
+                            sqlCommand = new SqlCommand(storedProc, connection as SqlConnection)
+                            {
+                                CommandType = CommandType.StoredProcedure,
+                                CommandTimeout = timeoutSecs
+                            };
+                            sqlCommand.Parameters.AddRange(parameters);
+                            using (var reader = await sqlCommand.ExecuteReaderAsync(token))
+                            {
+                                return await callback(reader);
+                            }
+                        }
+                        finally
+                        {
+                            if (sqlCommand != null)
+                            {
+                                sqlCommand.Parameters.Clear();
+                                sqlCommand.Dispose();
+                            }
+                        }
+                    }
+                }
+                catch (Exception e)
+                {
+                    error = e;
+                    throw;
+                }
+                finally
+                {
+                    stopwatch.Stop();
+                    RaiseOnExecuteReaderComplete(
+                        database, storedProc, stopwatch.ElapsedMilliseconds, error);
+                }
+            }, ShouldRetry(maxAttemptCount), RaiseOnError);
+        }
 
         private IEnumerable<TResult> QueryImpl<TRequest, TResult>(
             IStoredProc<TRequest, TResult> sp,

Agoda.Frameworks.DB/IDbRepository.cs

@@ -2,6 +2,7 @@
 using System.Collections.Generic;
 using System.Data;
 using System.Data.SqlClient;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace Agoda.Frameworks.DB
@@ -26,6 +27,7 @@ Task<T> ExecuteReaderAsync<T>(
             Func<SqlDataReader, Task<T>> callback,
             TimeSpan? timeSpan,
             string cacheKey = "");
+
         Task<object> ExecuteScalarAsync(
             string dbName,
             string sqlCommandString,
@@ -134,6 +136,15 @@ Task<T> ExecuteReaderAsync<T>(
             IDbDataParameter[] parameters,
             Func<SqlDataReader, Task<T>> callback);
 
+        Task<T> ExecuteReaderAsync<T>(
+            string database,
+            string storedProc,
+            int timeoutSecs,
+            int maxAttemptCount,
+            CancellationToken token,
+            IDbDataParameter[] parameters,
+            Func<SqlDataReader, Task<T>> callback);
+
         event EventHandler<DbErrorEventArgs> OnError;
         event EventHandler<QueryCompleteEventArgs> OnQueryComplete;
         event EventHandler<ExecuteReaderCompleteEventArgs> OnExecuteReaderComplete;