From 3b8b0cdc8ad0faf1f58f49585c1bd2a6d79bc32f Mon Sep 17 00:00:00 2001
From: mgravell <marc.gravell@gmail.com>
Date: Fri, 30 Aug 2019 15:23:15 +0100
Subject: [PATCH] add APIs for working with the connection pools and client-id

---
 .../DbConnectionExtensions.cs                 | 129 ++++++++++++++++++
 Dapper.Tests/ProviderTests.cs                 |  59 +++++++-
 2 files changed, 187 insertions(+), 1 deletion(-)
 create mode 100644 Dapper.ProviderTools/DbConnectionExtensions.cs

diff --git a/Dapper.ProviderTools/DbConnectionExtensions.cs b/Dapper.ProviderTools/DbConnectionExtensions.cs
new file mode 100644
index 000000000..05d21d5a7
--- /dev/null
+++ b/Dapper.ProviderTools/DbConnectionExtensions.cs
@@ -0,0 +1,129 @@
+using System;
+using System.Collections.Concurrent;
+using System.Data.Common;
+using System.Linq.Expressions;
+using System.Reflection;
+
+namespace Dapper.ProviderTools
+{
+    /// <summary>
+    /// Helper utilties for working with database connections
+    /// </summary>
+    public static class DbConnectionExtensions
+    {
+        /// <summary>
+        /// Attempt to get the client connection id for a given connection
+        /// </summary>
+        public static bool TryGetClientConnectionId(this DbConnection connection, out Guid clientConnectionId)
+        {
+            clientConnectionId = default;
+            return connection != null && ByTypeHelpers.Get(connection.GetType()).TryGetClientConnectionId(
+                connection, out clientConnectionId);
+        }
+
+        /// <summary>
+        /// Clear all pools associated with the provided connection type
+        /// </summary>
+        public static bool TryClearAllPools(this DbConnection connection)
+            => connection != null && ByTypeHelpers.Get(connection.GetType()).TryClearAllPools();
+
+        /// <summary>
+        /// Clear the pools associated with the provided connection
+        /// </summary>
+        public static bool TryClearPool(this DbConnection connection)
+            => connection != null && ByTypeHelpers.Get(connection.GetType()).TryClearPool(connection);
+
+        private sealed class ByTypeHelpers
+        {
+            private static readonly ConcurrentDictionary<Type, ByTypeHelpers> s_byType
+                = new ConcurrentDictionary<Type, ByTypeHelpers>();
+            private readonly Func<DbConnection, Guid>? _getClientConnectionId;
+
+            private readonly Action<DbConnection>? _clearPool;
+            private readonly Action? _clearAllPools;
+
+            public bool TryGetClientConnectionId(DbConnection connection, out Guid clientConnectionId)
+            {
+                if (_getClientConnectionId == null)
+                {
+                    clientConnectionId = default;
+                    return false;
+                }
+                clientConnectionId = _getClientConnectionId(connection);
+                return true;
+            }
+
+            public bool TryClearPool(DbConnection connection)
+            {
+                if (_clearPool == null) return false;
+                _clearPool(connection);
+                return true;
+            }
+
+            public bool TryClearAllPools()
+            {
+                if (_clearAllPools == null) return false;
+                _clearAllPools();
+                return true;
+            }
+
+            public static ByTypeHelpers Get(Type type)
+            {
+                if (!s_byType.TryGetValue(type, out var value))
+                {
+                    s_byType[type] = value = new ByTypeHelpers(type);
+                }
+                return value;
+            }
+
+            private ByTypeHelpers(Type type)
+            {
+                _getClientConnectionId = TryGetInstanceProperty<Guid>("ClientConnectionId", type);
+
+                try
+                {
+                    var clearAllPools = type.GetMethod("ClearAllPools", BindingFlags.Public | BindingFlags.Static,
+                        null, Type.EmptyTypes, null);
+                    if (clearAllPools != null)
+                    {
+                        _clearAllPools = (Action)Delegate.CreateDelegate(typeof(Action), clearAllPools);
+                    }
+                }
+                catch { }
+
+                try
+                {
+                    var clearPool = type.GetMethod("ClearPool", BindingFlags.Public | BindingFlags.Static,
+                        null, new[] { type }, null);
+                    if (clearPool != null)
+                    {
+                        var p = Expression.Parameter(typeof(DbConnection), "connection");
+                        var body = Expression.Call(clearPool, Expression.Convert(p, type));
+                        var lambda = Expression.Lambda<Action<DbConnection>>(body, p);
+                        _clearPool = lambda.Compile();
+                    }
+                }
+                catch { }
+            }
+
+            private static Func<DbConnection, T>? TryGetInstanceProperty<T>(string name, Type type)
+            {
+                try
+                {
+                    var prop = type.GetProperty(name, BindingFlags.Public | BindingFlags.Instance);
+                    if (prop == null || !prop.CanRead) return null;
+                    if (prop.PropertyType != typeof(T)) return null;
+
+                    var p = Expression.Parameter(typeof(DbConnection), "connection");
+                    var body = Expression.Property(Expression.Convert(p, type), prop);
+                    var lambda = Expression.Lambda<Func<DbConnection, T>>(body, p);
+                    return lambda.Compile();
+                }
+                catch
+                {
+                    return null;
+                }
+            }
+        }
+    }
+}
diff --git a/Dapper.Tests/ProviderTests.cs b/Dapper.Tests/ProviderTests.cs
index bb3c0886d..b51ed2d2e 100644
--- a/Dapper.Tests/ProviderTests.cs
+++ b/Dapper.Tests/ProviderTests.cs
@@ -1,4 +1,5 @@
-using System.Data.Common;
+using System;
+using System.Data.Common;
 using Dapper.ProviderTools;
 using Xunit;
 
@@ -24,6 +25,62 @@ public void BulkCopy_MicrosoftDataSqlClient()
             }
         }
 
+        [Fact]
+        public void ClientId_SystemDataSqlClient()
+            => TestClientId<SystemSqlClientProvider>();
+
+        [Fact]
+        public void ClientId_MicrosoftDataSqlClient()
+            => TestClientId<MicrosoftSqlClientProvider>();
+
+
+        [Fact]
+        public void ClearPool_SystemDataSqlClient()
+            => ClearPool<SystemSqlClientProvider>();
+
+        [Fact]
+        public void ClearPool_MicrosoftDataSqlClient()
+            => ClearPool<MicrosoftSqlClientProvider>();
+
+        [Fact]
+        public void ClearAllPools_SystemDataSqlClient()
+            => ClearAllPools<SystemSqlClientProvider>();
+
+        [Fact]
+        public void ClearAllPools_MicrosoftDataSqlClient()
+            => ClearAllPools<MicrosoftSqlClientProvider>();
+
+        private static void TestClientId<T>()
+             where T : SqlServerDatabaseProvider, new()
+        {
+            var provider = new T();
+            using (var conn = provider.GetOpenConnection())
+            {
+                Assert.True(conn.TryGetClientConnectionId(out var id));
+                Assert.NotEqual(Guid.Empty, id);
+            }
+        }
+
+        private static void ClearPool<T>()
+     where T : SqlServerDatabaseProvider, new()
+        {
+            var provider = new T();
+            using (var conn = provider.GetOpenConnection())
+            {
+                Assert.True(conn.TryClearPool());
+            }
+        }
+
+        private static void ClearAllPools<T>()
+     where T : SqlServerDatabaseProvider, new()
+        {
+            var provider = new T();
+            using (var conn = provider.GetOpenConnection())
+            {
+                Assert.True(conn.TryClearAllPools());
+            }
+        }
+
         private static void Test<T>(DbConnection connection)
         {
             using (var bcp = BulkCopy.TryCreate(connection))