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;