#941 Fixed the failing tracing scenarios for QueryMultiple operation.

mikependon · Sep 14, 2022 · 65eabbd · 65eabbd
1 parent 8861a1a
commit 65eabbd
Show file tree

Hide file tree

Showing 3 changed files with 1,157 additions and 848 deletions.
diff --git a/RepoDb.Core/RepoDb.Tests/RepoDb.UnitTests/Interfaces/ITraceForDbConnectionTest.cs b/RepoDb.Core/RepoDb.Tests/RepoDb.UnitTests/Interfaces/ITraceForDbConnectionTest.cs
@@ -6,6 +6,7 @@
 using System;
 using System.Collections.Generic;
 using System.Threading;
+using System.Threading.Tasks;
 
 namespace RepoDb.UnitTests.Interfaces
 {
@@ -1620,23 +1621,6 @@ public void TestDbConnectionTraceForAfterMaxViaTableName()
 
         #region MaxAsync
 
-        [TestMethod]
-        public void TestDbConnectionTraceForBeforeMaxAsync()
-        {
-            // Prepare
-            var trace = new Mock<ITrace>();
-            var connection = new TraceDbConnection();
-
-            // Act
-            connection.MaxAsync<TraceEntity>(trace: trace.Object,
-                field: e => e.Id,
-                where: (object)null).Wait();
-
-            // Assert
-            trace.Verify(t =>
-                t.BeforeExecution(It.IsAny<CancellableTraceLog>()), Times.Exactly(1));
-        }
-
         [TestMethod]
         public void TestDbConnectionTraceForAfterMaxAsync()
         {
@@ -1651,7 +1635,8 @@ public void TestDbConnectionTraceForAfterMaxAsync()
 
             // Assert
             trace.Verify(t =>
-                t.AfterExecution(It.IsAny<ResultTraceLog<object>>()), Times.Exactly(1));
+                t.AfterExecutionAsync(It.IsAny<ResultTraceLog<object>>(),
+                    It.IsAny<CancellationToken>()), Times.Exactly(1));
         }
 
         [TestMethod]
@@ -1669,7 +1654,8 @@ public void TestDbConnectionTraceForBeforeMaxAsyncViaTableName()
 
             // Assert
             trace.Verify(t =>
-                t.BeforeExecution(It.IsAny<CancellableTraceLog>()), Times.Exactly(1));
+                t.BeforeExecutionAsync(It.IsAny<CancellableTraceLog>(),
+                    It.IsAny<CancellationToken>()), Times.Exactly(1));
         }
 
         [TestMethod]
@@ -1687,7 +1673,8 @@ public void TestDbConnectionTraceForAfterMaxAsyncViaTableName()
 
             // Assert
             trace.Verify(t =>
-                t.AfterExecution(It.IsAny<ResultTraceLog<object>>()), Times.Exactly(1));
+                t.AfterExecutionAsync(It.IsAny<ResultTraceLog<object>>(),
+                    It.IsAny<CancellationToken>()), Times.Exactly(1));
         }
 
         #endregion
@@ -2029,7 +2016,7 @@ public void TestDbConnectionTraceForAfterMergeAll()
 
             // Assert
             trace.Verify(t =>
-                t.AfterExecution(It.IsAny<ResultTraceLog<object>>()), Times.Exactly(1));
+                t.AfterExecution(It.IsAny<ResultTraceLog<int>>()), Times.Exactly(1));
         }
 
         [TestMethod]
@@ -2063,7 +2050,7 @@ public void TestDbConnectionTraceForAfterMergeAllViaTableName()
 
             // Assert
             trace.Verify(t =>
-                t.AfterExecution(It.IsAny<ResultTraceLog<object>>()), Times.Exactly(1));
+                t.AfterExecution(It.IsAny<ResultTraceLog<int>>()), Times.Exactly(1));
         }
 
         #endregion
@@ -2666,7 +2653,7 @@ public void TestDbConnectionTraceForBeforeQueryMultipleForT3()
         }
 
         [TestMethod]
-        public void TestDbConnectionTraceForAfterQueryMultipleT3()
+        public void TestDbConnectionTraceForAfterQueryMultipleForT3()
         {
             // Prepare
             var trace = new Mock<ITrace>();
@@ -2709,7 +2696,7 @@ public void TestDbConnectionTraceForBeforeQueryMultipleForT4()
         }
 
         [TestMethod]
-        public void TestDbConnectionTraceForAfterQueryMultipleT4()
+        public void TestDbConnectionTraceForAfterQueryMultipleForT4()
         {
             // Prepare
             var trace = new Mock<ITrace>();
@@ -2755,7 +2742,7 @@ public void TestDbConnectionTraceForBeforeQueryMultipleForT5()
         }
 
         [TestMethod]
-        public void TestDbConnectionTraceForAfterQueryMultipleT5()
+        public void TestDbConnectionTraceForAfterQueryMultipleForT5()
         {
             // Prepare
             var trace = new Mock<ITrace>();
@@ -2804,7 +2791,7 @@ public void TestDbConnectionTraceForBeforeQueryMultipleForT6()
         }
 
         [TestMethod]
-        public void TestDbConnectionTraceForAfterQueryMultipleT6()
+        public void TestDbConnectionTraceForAfterQueryMultipleForT6()
         {
             // Prepare
             var trace = new Mock<ITrace>();
@@ -2856,7 +2843,7 @@ public void TestDbConnectionTraceForBeforeQueryMultipleForT7()
         }
 
         [TestMethod]
-        public void TestDbConnectionTraceForAfterQueryMultipleT7()
+        public void TestDbConnectionTraceForAfterQueryMultipleForT7()
         {
             // Prepare
             var trace = new Mock<ITrace>();

diff --git a/RepoDb.Core/RepoDb/Extensions/DbConnectionExtension.cs b/RepoDb.Core/RepoDb/Extensions/DbConnectionExtension.cs
@@ -1251,6 +1251,7 @@ public static IDataReader ExecuteReader(this IDbConnection connection,
         /// <param name="entityType"></param>
         /// <param name="dbFields"></param>
         /// <param name="skipCommandArrayParametersCheck"></param>
+        /// <param name="beforeExecutionCallback"></param>
         /// <returns></returns>
         internal static IDataReader ExecuteReaderInternal(this IDbConnection connection,
             string commandText,
@@ -1262,7 +1263,8 @@ internal static IDataReader ExecuteReaderInternal(this IDbConnection connection,
             ITrace trace,
             Type entityType,
             IEnumerable<DbField> dbFields,
-            bool skipCommandArrayParametersCheck)
+            bool skipCommandArrayParametersCheck,
+            Action<DbCommand> beforeExecutionCallback = null)
         {
             // Variables
             var setting = DbSettingMapper.Get(connection);
@@ -1280,6 +1282,9 @@ internal static IDataReader ExecuteReaderInternal(this IDbConnection connection,
             // Ensure the DbCommand disposal
             try
             {
+                // A hacky solution for other operations (i.e.: QueryMultiple)
+                beforeExecutionCallback?.Invoke(command);
+
                 // Before Execution
                 var traceResult = Tracer
                     .InvokeBeforeExecution(traceKey, trace, command);
@@ -1371,6 +1376,7 @@ public static Task<IDataReader> ExecuteReaderAsync(this IDbConnection connection
         /// <param name="entityType"></param>
         /// <param name="dbFields"></param>
         /// <param name="skipCommandArrayParametersCheck"></param>
+        /// <param name="beforeExecutionCallbackAsync"></param>
         /// <returns></returns>
         internal static async Task<IDataReader> ExecuteReaderAsyncInternal(this IDbConnection connection,
             string commandText,
@@ -1383,7 +1389,8 @@ internal static async Task<IDataReader> ExecuteReaderAsyncInternal(this IDbConne
             CancellationToken cancellationToken,
             Type entityType,
             IEnumerable<DbField> dbFields,
-            bool skipCommandArrayParametersCheck)
+            bool skipCommandArrayParametersCheck,
+            Func<DbCommand, CancellationToken, Task> beforeExecutionCallbackAsync = null)
         {
             // Variables
             var setting = connection.GetDbSetting();
@@ -1402,6 +1409,12 @@ internal static async Task<IDataReader> ExecuteReaderAsyncInternal(this IDbConne
             // Ensure the DbCommand disposal
             try
             {
+                // A hacky solution for other operations (i.e.: QueryMultipleAsync)
+                if (beforeExecutionCallbackAsync != null)
+                {
+                    await beforeExecutionCallbackAsync(command, cancellationToken);
+                }
+
                 // Before Execution
                 var traceResult = await Tracer
                     .InvokeBeforeExecutionAsync(traceKey, trace, command, cancellationToken);