Add support for LINQ to SQL typed data context. Closes #4

Giorgi · Mar 20, 2016 · 779cbda · 779cbda
1 parent 34a1b54
commit 779cbda
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 22 deletions.
diff --git a/src/QueryPlanVisualizer/Helpers/DatabaseHelper.cs b/src/QueryPlanVisualizer/Helpers/DatabaseHelper.cs
@@ -1,7 +1,8 @@
 using System;
-using System.Data;
 using System.Data.Common;
+using System.Data.Linq;
 using System.Linq;
+using System.Reflection;
 using System.Threading.Tasks;
 using LINQPad;
 
@@ -11,14 +12,34 @@ internal abstract class DatabaseHelper
     {
         private DbConnection _dbConnection;
 
-        public static DatabaseHelper Create(DataContextBase dataContextBase)
+        public static DatabaseHelper Create<T>(DataContextBase dataContextBase, IQueryable<T> queryable)
         {
-            if (dataContextBase == null)
+            if (dataContextBase != null)
             {
-                return new EntityFrameworkDatabaseHelper();
+                return new LinqToSqlDatabaseHelper(dataContextBase);
             }
 
-            return new LinqToSqlDatabaseHelper(dataContextBase.Connection);
+            var table = queryable as ITable;
+
+            if (table != null)
+            {
+                return new LinqToSqlDatabaseHelper(table.Context);
+            }
+
+            var dataQueryType = typeof(DataContext).Assembly.GetType("System.Data.Linq.DataQuery`1");
+
+            if (queryable.GetType().GetGenericTypeDefinition() == dataQueryType)
+            {
+                var contextField = queryable.GetType().GetField("context", BindingFlags.Instance | BindingFlags.NonPublic | BindingFlags.GetField);
+                var context = contextField?.GetValue(queryable) as DataContext;
+
+                if (context != null)
+                {
+                    return new LinqToSqlDatabaseHelper(context);
+                }
+            }
+
+            return new EntityFrameworkDatabaseHelper();
         }
 
         public DbConnection Connection
@@ -86,8 +107,7 @@ public virtual async Task CreateIndexAsync(string script)
                 Connection.Close();
             }
         }
-
-        protected abstract DbCommand CreateCommand(IQueryable queryable);
 
+        protected abstract DbCommand CreateCommand(IQueryable queryable);
     }
 }
diff --git a/src/QueryPlanVisualizer/Helpers/LinqToSqlDatabaseHelper.cs b/src/QueryPlanVisualizer/Helpers/LinqToSqlDatabaseHelper.cs
@@ -1,20 +1,22 @@
-using System;
 using System.Data.Common;
+using System.Data.Linq;
 using System.Linq;
-using LINQPad;
 
 namespace ExecutionPlanVisualizer.Helpers
 {
     internal class LinqToSqlDatabaseHelper : DatabaseHelper
     {
-        public LinqToSqlDatabaseHelper(DbConnection connection)
+        private readonly DataContext dataContext;
+
+        public LinqToSqlDatabaseHelper(DataContext dataContext)
         {
-            Connection = connection;
+            this.dataContext = dataContext;
+            Connection = dataContext.Connection;
         }
 
         protected override DbCommand CreateCommand(IQueryable queryable)
         {
-            return Util.CurrentDataContext.GetCommand(queryable);
+            return dataContext.GetCommand(queryable);
         }
     }
 }
diff --git a/src/QueryPlanVisualizer/QueryPlanUserControl.cs b/src/QueryPlanVisualizer/QueryPlanUserControl.cs
@@ -4,7 +4,6 @@
 using System.IO;
 using System.Windows.Forms;
 using ExecutionPlanVisualizer.Helpers;
-using LINQPad;
 
 namespace ExecutionPlanVisualizer
 {

diff --git a/src/QueryPlanVisualizer/QueryPlanVisualizer.cs b/src/QueryPlanVisualizer/QueryPlanVisualizer.cs
@@ -25,16 +25,15 @@ public static IQueryable<T> DumpPlan<T>(this IQueryable<T> queryable, bool dumpD
 
         private static void DumpPlanInternal<T>(IQueryable<T> queryable, bool dumpData, bool addNewPanel)
         {
-            var databaseHelper = DatabaseHelper.Create(Util.CurrentDataContext);
-
-            var sqlConnection = Util.CurrentDataContext?.Connection as SqlConnection;
-
-            if (databaseHelper is LinqToSqlDatabaseHelper && sqlConnection == null)
+            if (Util.CurrentDataContext != null && !(Util.CurrentDataContext.Connection is SqlConnection))
             {
                 var control = new Label { Text = "Query Plan Visualizer supports only Sql Server" };
                 PanelManager.DisplayControl(control, ExecutionPlanPanelTitle);
+                return;
             }
 
+            var databaseHelper = DatabaseHelper.Create(Util.CurrentDataContext, queryable);
+
             if (dumpData)
             {
                 queryable.Dump();
@@ -60,9 +59,9 @@ private static void DumpPlanInternal<T>(IQueryable<T> queryable, bool dumpData,
                 if (control == null || addNewPanel)
                 {
                     control = new QueryPlanUserControl()
-                              {
-                                  DatabaseHelper = databaseHelper
-                              };
+                    {
+                        DatabaseHelper = databaseHelper
+                    };
 
                     control.IndexCreated += (sender, args) =>
                     {
@@ -79,7 +78,7 @@ private static void DumpPlanInternal<T>(IQueryable<T> queryable, bool dumpData,
             }
             catch (Exception exception)
             {
-                var control = new Label {Text = exception.ToString()};
+                var control = new Label { Text = exception.ToString() };
                 PanelManager.DisplayControl(control, ExecutionPlanPanelTitle);
             }
         }