From e5b5a553da59c753d93525ed0a1ea6ea6b16c507 Mon Sep 17 00:00:00 2001
From: Jinan R <jinan.raghavan@agoda.com>
Date: Tue, 27 Feb 2024 11:08:38 +0700
Subject: [PATCH] Fixed issues

---
 Agoda.Frameworks.DB.Tests/DbRepositoryTest.cs | 21 -------------------
 Agoda.Frameworks.DB/DbRepository.cs           |  7 +++++--
 Agoda.Frameworks.DB/IDbRepository.cs          |  2 +-
 3 files changed, 6 insertions(+), 24 deletions(-)

diff --git a/Agoda.Frameworks.DB.Tests/DbRepositoryTest.cs b/Agoda.Frameworks.DB.Tests/DbRepositoryTest.cs
index 2ba7f52..e735f23 100644
--- a/Agoda.Frameworks.DB.Tests/DbRepositoryTest.cs
+++ b/Agoda.Frameworks.DB.Tests/DbRepositoryTest.cs
@@ -667,27 +667,6 @@ public void QueryMultipleAsync_Retry_Failure()
             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>
         {
             public string DbName => "mobile_ro";
diff --git a/Agoda.Frameworks.DB/DbRepository.cs b/Agoda.Frameworks.DB/DbRepository.cs
index 487b5fa..7398139 100644
--- a/Agoda.Frameworks.DB/DbRepository.cs
+++ b/Agoda.Frameworks.DB/DbRepository.cs
@@ -303,7 +303,7 @@ public Task<T> ExecuteReaderAsync<T>(
             string storedProc,
             int timeoutSecs,
             int maxAttemptCount,
-            CancellationToken token,
+            int taskCancellationTimeOutInMilliSecs,
             IDbDataParameter[] parameters,
             Func<SqlDataReader, Task<T>> callback)
         {
@@ -313,6 +313,8 @@ public Task<T> ExecuteReaderAsync<T>(
                 Exception error = null;
                 try
                 {
+                    var taskCancellationSource = new CancellationTokenSource();
+                    taskCancellationSource.CancelAfter(taskCancellationTimeOutInMilliSecs);
                     using (var connection = _generateConnection(connectionStr))
                     {
                         if (connection is SqlConnection sqlConn)
@@ -332,7 +334,8 @@ public Task<T> ExecuteReaderAsync<T>(
                                 CommandTimeout = timeoutSecs
                             };
                             sqlCommand.Parameters.AddRange(parameters);
-                            using (var reader = await sqlCommand.ExecuteReaderAsync(token))
+                          
+                            using (var reader = await sqlCommand.ExecuteReaderAsync(taskCancellationSource.Token))
                             {
                                 return await callback(reader);
                             }
diff --git a/Agoda.Frameworks.DB/IDbRepository.cs b/Agoda.Frameworks.DB/IDbRepository.cs
index 9729254..1d84d8a 100644
--- a/Agoda.Frameworks.DB/IDbRepository.cs
+++ b/Agoda.Frameworks.DB/IDbRepository.cs
@@ -141,7 +141,7 @@ Task<T> ExecuteReaderAsync<T>(
             string storedProc,
             int timeoutSecs,
             int maxAttemptCount,
-            CancellationToken token,
+            int taskCancellationTimeOutInMilliSecs,
             IDbDataParameter[] parameters,
             Func<SqlDataReader, Task<T>> callback);