From 2a66ac949b6573af3e7ecbb9c4b47ffa336d8992 Mon Sep 17 00:00:00 2001
From: Fitzchak Yitzchaki <fitzchak@ayende.com>
Date: Tue, 22 Nov 2016 15:56:13 +0200
Subject: [PATCH] #7105 Trace when query is done. This is needed in order to
 get SelectRows statistics. Take #2: Use the DiagnosticSource on
 RelationalDataReader. Take #3: Fix requested changes.

---
 .../Internal/RelationalDiagnostics.cs                 | 11 +++++++++++
 .../Storage/Internal/RelationalCommand.cs             |  6 ++++--
 .../Storage/RelationalDataReader.cs                   | 10 +++++++++-
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/Microsoft.EntityFrameworkCore.Relational/Internal/RelationalDiagnostics.cs b/src/Microsoft.EntityFrameworkCore.Relational/Internal/RelationalDiagnostics.cs
index a3e4f08b73d..0e965b2025b 100644
--- a/src/Microsoft.EntityFrameworkCore.Relational/Internal/RelationalDiagnostics.cs
+++ b/src/Microsoft.EntityFrameworkCore.Relational/Internal/RelationalDiagnostics.cs
@@ -4,6 +4,7 @@
 using System;
 using System.Data.Common;
 using System.Diagnostics;
+using Microsoft.EntityFrameworkCore.Storage;
 
 namespace Microsoft.EntityFrameworkCore.Internal
 {
@@ -15,6 +16,8 @@ internal static class RelationalDiagnostics
         public const string AfterExecuteCommand = NamePrefix + nameof(AfterExecuteCommand);
         public const string CommandExecutionError = NamePrefix + nameof(CommandExecutionError);
 
+        public const string DataReaderDisposing = NamePrefix + nameof(DataReaderDisposing);
+
         public static void WriteCommandBefore(
             this DiagnosticSource diagnosticSource,
             DbCommand command, string executeMethod,
@@ -88,5 +91,13 @@ public static void WriteCommandError(
                     });
             }
         }
+
+        public static void WriteDataReaderDisposing(this DiagnosticSource diagnosticSource, DbDataReader dataReader)
+        {
+            if (diagnosticSource.IsEnabled(DataReaderDisposing))
+            {
+                diagnosticSource.Write(DataReaderDisposing, dataReader);
+            }
+        }
     }
 }
diff --git a/src/Microsoft.EntityFrameworkCore.Relational/Storage/Internal/RelationalCommand.cs b/src/Microsoft.EntityFrameworkCore.Relational/Storage/Internal/RelationalCommand.cs
index 06a49ed1d48..8eb5a08805a 100644
--- a/src/Microsoft.EntityFrameworkCore.Relational/Storage/Internal/RelationalCommand.cs
+++ b/src/Microsoft.EntityFrameworkCore.Relational/Storage/Internal/RelationalCommand.cs
@@ -226,7 +226,8 @@ protected virtual object Execute(
                                 = new RelationalDataReader(
                                     connection,
                                     dbCommand,
-                                    dbCommand.ExecuteReader());
+                                    dbCommand.ExecuteReader(),
+                                    DiagnosticSource);
                         }
                         catch
                         {
@@ -344,7 +345,8 @@ protected virtual async Task<object> ExecuteAsync(
                             result = new RelationalDataReader(
                                 connection,
                                 dbCommand,
-                                await dbCommand.ExecuteReaderAsync(cancellationToken));
+                                await dbCommand.ExecuteReaderAsync(cancellationToken), 
+                                DiagnosticSource);
                         }
                         catch
                         {
diff --git a/src/Microsoft.EntityFrameworkCore.Relational/Storage/RelationalDataReader.cs b/src/Microsoft.EntityFrameworkCore.Relational/Storage/RelationalDataReader.cs
index 6d6e5e57476..cf64fd352d7 100644
--- a/src/Microsoft.EntityFrameworkCore.Relational/Storage/RelationalDataReader.cs
+++ b/src/Microsoft.EntityFrameworkCore.Relational/Storage/RelationalDataReader.cs
@@ -3,7 +3,9 @@
 
 using System;
 using System.Data.Common;
+using System.Diagnostics;
 using JetBrains.Annotations;
+using Microsoft.EntityFrameworkCore.Internal;
 using Microsoft.EntityFrameworkCore.Utilities;
 
 namespace Microsoft.EntityFrameworkCore.Storage
@@ -22,6 +24,7 @@ public class RelationalDataReader : IDisposable
         private readonly IRelationalConnection _connection;
         private readonly DbCommand _command;
         private readonly DbDataReader _reader;
+        private readonly DiagnosticSource _diagnosticSource;
 
         private bool _disposed;
 
@@ -31,17 +34,21 @@ public class RelationalDataReader : IDisposable
         /// <param name="connection"> The connection. </param>
         /// <param name="command"> The command that was executed. </param>
         /// <param name="reader"> The underlying reader for the result set. </param>
+        /// <param name="diagnosticSource"> The diagnostic source. </param>
         public RelationalDataReader(
             [CanBeNull] IRelationalConnection connection,
             [NotNull] DbCommand command,
-            [NotNull] DbDataReader reader)
+            [NotNull] DbDataReader reader,
+            [NotNull] DiagnosticSource diagnosticSource)
         {
             Check.NotNull(command, nameof(command));
             Check.NotNull(reader, nameof(reader));
+            Check.NotNull(diagnosticSource, nameof(diagnosticSource));
 
             _connection = connection;
             _command = command;
             _reader = reader;
+            _diagnosticSource = diagnosticSource;
         }
 
         /// <summary>
@@ -65,6 +72,7 @@ public virtual void Dispose()
         {
             if (!_disposed)
             {
+                _diagnosticSource.WriteDataReaderDisposing(_reader);
                 _reader.Dispose();
                 _command.Dispose();
                 _connection?.Close();