From 655da45b343041d2dd8193f87d646a1f07e0bd00 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Halil=20=C4=B0brahim=20Kalkan?= <hi_kalkan@yahoo.com>
Date: Thu, 27 Jul 2017 17:35:33 +0300
Subject: [PATCH] Added UnitOfWorkInterceptorRegistrar.

---
 .../Volo/Abp/Uow/UnitOfWorkInterceptor.cs     | 16 +-----
 .../Abp/Uow/UnitOfWorkInterceptorRegistrar.cs | 49 +++++++++++++++++++
 2 files changed, 50 insertions(+), 15 deletions(-)
 create mode 100644 src/Volo.Abp/Volo/Abp/Uow/UnitOfWorkInterceptorRegistrar.cs

diff --git a/src/Volo.Abp/Volo/Abp/Uow/UnitOfWorkInterceptor.cs b/src/Volo.Abp/Volo/Abp/Uow/UnitOfWorkInterceptor.cs
index 5270243ad4..b486d84e09 100644
--- a/src/Volo.Abp/Volo/Abp/Uow/UnitOfWorkInterceptor.cs
+++ b/src/Volo.Abp/Volo/Abp/Uow/UnitOfWorkInterceptor.cs
@@ -1,23 +1,9 @@
-using System.Reflection;
-using System.Threading.Tasks;
-using Microsoft.Extensions.DependencyInjection;
-using Volo.Abp.Application.Services;
+using System.Threading.Tasks;
 using Volo.Abp.DynamicProxy;
 using Volo.DependencyInjection;
 
 namespace Volo.Abp.Uow
 {
-    public static class UnitOfWorkInterceptorRegistrar
-    {
-        public static void RegisterIfNeeded(IOnServiceRegistredContext context)
-        {
-            if (typeof(IApplicationService).GetTypeInfo().IsAssignableFrom(context.ImplementationType))
-            {
-                context.Interceptors.Add<UnitOfWorkInterceptor>();
-            }
-        }
-    }
-
     public class UnitOfWorkInterceptor : AbpInterceptor, ITransientDependency
     {
         private readonly IUnitOfWorkManager _unitOfWorkManager;
diff --git a/src/Volo.Abp/Volo/Abp/Uow/UnitOfWorkInterceptorRegistrar.cs b/src/Volo.Abp/Volo/Abp/Uow/UnitOfWorkInterceptorRegistrar.cs
new file mode 100644
index 0000000000..edac0be69c
--- /dev/null
+++ b/src/Volo.Abp/Volo/Abp/Uow/UnitOfWorkInterceptorRegistrar.cs
@@ -0,0 +1,49 @@
+using System.Linq;
+using System.Reflection;
+using Microsoft.Extensions.DependencyInjection;
+using Volo.Abp.Application.Services;
+using Volo.Abp.Domain.Repositories;
+
+namespace Volo.Abp.Uow
+{
+    public static class UnitOfWorkInterceptorRegistrar
+    {
+        public static void RegisterIfNeeded(IOnServiceRegistredContext context)
+        {
+            if (IsUnitOfWorkType(context.ImplementationType.GetTypeInfo()))
+            {
+                context.Interceptors.Add<UnitOfWorkInterceptor>();
+            }
+        }
+
+        public static bool IsUnitOfWorkType(TypeInfo implementationType)
+        {
+            //Explicitly defined UnitOfWorkAttribute
+            if (HasUnitOfWorkAttribute(implementationType) || AnyMethodHasUnitOfWorkAttribute(implementationType))
+            {
+                return true;
+            }
+
+            //Conventional classes
+            if (typeof(IApplicationService).GetTypeInfo().IsAssignableFrom(implementationType) ||
+                typeof(IRepository).GetTypeInfo().IsAssignableFrom(implementationType))
+            {
+                return true;
+            }
+
+            return false;
+        }
+
+        private static bool AnyMethodHasUnitOfWorkAttribute(TypeInfo implementationType)
+        {
+            return implementationType
+                .GetMethods(BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic)
+                .Any(HasUnitOfWorkAttribute);
+        }
+
+        private static bool HasUnitOfWorkAttribute(MemberInfo methodInfo)
+        {
+            return methodInfo.IsDefined(typeof(UnitOfWorkAttribute), true);
+        }
+    }
+}
\ No newline at end of file