Giorgi · Giorgi · Oct 5, 2024 · Oct 5, 2024 · Oct 5, 2024 · Oct 5, 2024
diff --git a/src/EFCore.Visualizer/EFCore.Visualizer.csproj b/src/EFCore.Visualizer/EFCore.Visualizer.csproj
@@ -10,12 +10,16 @@
     <None Remove="Fonts\fa-brands-400.ttf" />
     <None Remove="Fonts\fa-solid-900.ttf" />
     <None Remove="QueryPlanUserControl.xaml" />
+    <None Remove="Resources\Common\template.html" />
     <None Remove="Resources\Postgres\template.html" />
     <None Remove="Resources\Oracle\template.html" />
   </ItemGroup>
 
   <ItemGroup>
     <Content Include="..\IQueryableObjectSource\bin\$(Configuration)\net6\IQueryableObjectSource.dll" Link="netcoreapp\IQueryableObjectSource.dll" />
+    <Content Include="Resources\Common\template.html">
+      <IncludeInVSIX>true</IncludeInVSIX>
+    </Content>
   </ItemGroup>
 
   <ItemGroup>
@@ -58,6 +62,13 @@
     <Content Include="Resources\Oracle\*.*">
       <IncludeInVSIX>true</IncludeInVSIX>
     </Content>
+    <Content Include="Resources\Common\*.*">
+      <IncludeInVSIX>true</IncludeInVSIX>
+    </Content>
+  </ItemGroup>
+
+  <ItemGroup>
+    <Compile Include="..\IQueryableObjectSource\OperationType.cs" Link="OperationType.cs" />
   </ItemGroup>
 
   <ItemGroup>

diff --git a/src/EFCore.Visualizer/QueryPlanUserControl.xaml.cs b/src/EFCore.Visualizer/QueryPlanUserControl.xaml.cs
@@ -7,6 +7,7 @@
 using System.Text;
 using System.Windows;
 using System.Windows.Controls;
+using IQueryableObjectSource;
 
 namespace EFCore.Visualizer
 {
@@ -42,6 +43,7 @@ private void QueryPlanUserControlUnloaded(object sender, RoutedEventArgs e)
         protected override async void OnInitialized(EventArgs e)
 #pragma warning restore VSTHRD100 // Avoid async void methods
         {
+            var query = string.Empty;
             try
             {
                 base.OnInitialized(e);
@@ -50,6 +52,8 @@ protected override async void OnInitialized(EventArgs e)
                 var environment = await CoreWebView2Environment.CreateAsync(userDataFolder: Path.Combine(AssemblyLocation, "WVData"));
                 await webView.EnsureCoreWebView2Async(environment);
 
+                query = await GetQueryAsync();
+
 #if !DEBUG
                 webView.CoreWebView2.Settings.AreBrowserAcceleratorKeysEnabled = false;
                 webView.CoreWebView2.Settings.AreDefaultContextMenusEnabled = false; 
@@ -84,10 +88,34 @@ protected override async void OnInitialized(EventArgs e)
             }
             catch (Exception ex)
             {
+                if (!string.IsNullOrEmpty(query))
+                    webView.CoreWebView2.NavigateToString(query);
                 MessageBox.Show(ex.Message, "Error", MessageBoxButton.OK, MessageBoxImage.Error);
             }
         }
 
+        private async Task<string> GetQueryAsync()
+        {
+            var query = string.Empty;
+            var response = await visualizerTarget.ObjectSource.RequestDataAsync(ConvertEnumToReadOnlySequence(OperationType.GetQuery),CancellationToken.None);
+            if (response.HasValue)
+            {
+                using var stream = response.Value.AsStream();
+                using var binaryReader = new BinaryReader(stream, Encoding.Default);
+                var isError = binaryReader.ReadBoolean();
+                if (!isError)
+                {
+                    query = binaryReader.ReadString();
+                }
+            }
+            return query;
+        }
+        private ReadOnlySequence<byte> ConvertEnumToReadOnlySequence(OperationType operation)
+        {
+            var operationByte = new byte[] { (byte)operation };
+            return new ReadOnlySequence<byte>(operationByte);
+        }
+
         private void ButtonReviewClick(object sender, RoutedEventArgs e)
         {
             StartProcess("https://marketplace.visualstudio.com/items?itemName=GiorgiDalakishvili.EFCoreVisualizer&ssr=false#review-details");

diff --git a/src/EFCore.Visualizer/Resources/Common/template.html b/src/EFCore.Visualizer/Resources/Common/template.html
@@ -0,0 +1,36 @@
+<!DOCTYPE html>
+<html>
+<head>
+    <title>Query Plan Visualizer</title>
+    <style>
+        body {
+            font-family: Arial, sans-serif;
+            margin: 20px;
+        }
+
+        h2 {
+            color: #333;
+        }
+
+        .query-box {
+            background-color: #f4f4f4;
+            padding: 10px;
+            border-radius: 5px;
+            border: 1px solid #ccc;
+            overflow-x: auto;
+            max-width: 100%;
+        }
+
+        pre {
+            white-space: pre-wrap;
+            word-wrap: break-word;
+        }
+    </style>
+</head>
+<body>
+    <h2>SQL Query</h2>
+    <div class='query-box'>
+        <pre>{query}</pre>
+    </div>
+</body>
+</html>
diff --git a/src/IQueryableObjectSource/EFCoreQueryableObjectSource.cs b/src/IQueryableObjectSource/EFCoreQueryableObjectSource.cs
@@ -6,6 +6,10 @@
 using System.Linq;
 using System.Text;
 using System.Text.Encodings.Web;
+using System.Threading.Tasks;
+using System.Net;
+using Microsoft.EntityFrameworkCore.Storage;
+using System.Diagnostics;
 
 namespace IQueryableObjectSource
 {
@@ -22,55 +26,123 @@ public override void TransferData(object target, Stream incomingData, Stream out
 
             try
             {
-                using var command = queryable.CreateDbCommand();
-                var provider = GetDatabaseProvider(command);
-
-                if (provider == null)
+                var operationType = GetOperationType(incomingData);
+                switch (operationType)
                 {
-                    return;
+                    case OperationType.GetQuery:
+                        HandleGetQuery(queryable, outgoingData);
+                        break;
+
+                    case OperationType.NotSupported:
+                        throw new InvalidOperationException("Unknown operation type."); 
+
+                    default:
+                        HandleGetQueryPlan(queryable, incomingData, outgoingData);
+                        break;
                 }
+            }
+            catch (Exception ex)
+            {
+                WriteError(outgoingData, ex.Message);
+            }
+        }
+        private void HandleGetQuery(IQueryable queryable, Stream outgoingData)
+        {
+            using var queryWriter = new BinaryWriter(outgoingData, Encoding.Default, true);
+            queryWriter.Write(false); // Indicates no error
+            queryWriter.Write(GenerateHtml(queryable.ToQueryString()));
+        }
+        private void HandleGetQueryPlan(IQueryable queryable, Stream incomingData, Stream outgoingData)
+        {
+            using var command = queryable.CreateDbCommand();
+            var provider = GetDatabaseProvider(command);
 
-                var query = queryable.ToQueryString();
-                var rawPlan = provider.ExtractPlan();
+            if (provider == null)
+            {
+                return;
+            }
 
-                var buffer = new byte[3];
-                var isBackgroundDarkColor = false;
+            var query = queryable.ToQueryString();
+            var rawPlan = provider.ExtractPlan();
 
-                var r = 255;
-                var g = 255;
-                var b = 255;
+            var (r, g, b) = ReadBackgroundColor(incomingData);
+            var isBackgroundDarkColor = r * 0.2126 + g * 0.7152 + b * 0.0722 < 255 / 2.0; 
 
-                if (incomingData.Read(buffer, 0, buffer.Length) == buffer.Length)
-                {
-                    r = buffer[0];
-                    g = buffer[1];
-                    b = buffer[2];
-                }
+            var planFile = GeneratePlanFile(provider, query, rawPlan, r, g, b, isBackgroundDarkColor);
+
+            using var writer = new BinaryWriter(outgoingData, Encoding.Default, true);
+            writer.Write(false); // Indicates no error
+            writer.Write(planFile);
+        }
+        private (int r, int g, int b) ReadBackgroundColor(Stream incomingData)
+        {
+            var buffer = new byte[3];
+            var r = 255;
+            var g = 255;
+            var b = 255;
+
+            if (incomingData.Read(buffer, 0, buffer.Length) == buffer.Length)
+            {
+                r = buffer[0];
+                g = buffer[1];
+                b = buffer[2];
+            }
+
+            return (r, g, b);
+        }
+
+        private string GeneratePlanFile(DatabaseProvider provider, string query, string rawPlan, int r, int g, int b, bool isBackgroundDarkColor)
+        {
+            var planDirectory = provider.GetPlanDirectory(ResourcesLocation);
+            var planFile = Path.Combine(planDirectory, Path.ChangeExtension(Path.GetRandomFileName(), "html"));
 
-                isBackgroundDarkColor = r * 0.2126 + g * 0.7152 + b * 0.0722 < 255 / 2.0;
+            var planPageHtml = File.ReadAllText(Path.Combine(planDirectory, "template.html"))
+                .Replace("{backColor}", $"rgb({r} {g} {b})")
+                .Replace("{textColor}", isBackgroundDarkColor ? "white" : "black")
+                .Replace("{plan}", JavaScriptEncoder.UnsafeRelaxedJsonEscaping.Encode(rawPlan).Replace("'", "\\'"))
+                .Replace("{query}", JavaScriptEncoder.UnsafeRelaxedJsonEscaping.Encode(query).Replace("'", "\\'"));
 
-                var planFile = Path.Combine(provider.GetPlanDirectory(ResourcesLocation), Path.ChangeExtension(Path.GetRandomFileName(), "html"));
+            File.WriteAllText(planFile, planPageHtml);
 
-                var planPageHtml = File.ReadAllText(Path.Combine(provider.GetPlanDirectory(ResourcesLocation), "template.html"))
-                    .Replace("{backColor}", $"rgb({r} {g} {b})")
-                    .Replace("{textColor}", isBackgroundDarkColor ? "white" : "black")
-                    .Replace("{plan}", JavaScriptEncoder.UnsafeRelaxedJsonEscaping.Encode(rawPlan).Replace("'", "\\'"))
-                    .Replace("{query}", JavaScriptEncoder.UnsafeRelaxedJsonEscaping.Encode(query).Replace("'", "\\'"));
+            return planFile;
+        }
 
-                File.WriteAllText(planFile, planPageHtml);
+        private void WriteError(Stream outgoingData, string errorMessage)
+        {
+            using var writer = new BinaryWriter(outgoingData, Encoding.Default, true);
+            writer.Write(true); // Indicates an error occurred
+            writer.Write(errorMessage);
+        }
 
-                using var writer = new BinaryWriter(outgoingData, Encoding.Default, true);
-                writer.Write(false);
-                writer.Write(planFile);
+        public static OperationType GetOperationType(Stream stream)
+        {
+            try
+            {
+                var operationBuffer = new byte[1];
+                stream.Read(operationBuffer, 0, 1);
+                return (OperationType)operationBuffer[0];
             }
-            catch (Exception ex)
+            catch (Exception)
+            {
+                return OperationType.NotSupported;
+            }
+        }
+
+        private static string GenerateHtml(string query)
+        {
+            string escapedQuery = WebUtility.HtmlEncode(query);
+            string templatePath = Path.Combine(ResourcesLocation,"Common", "template.html");
+            if (!File.Exists(templatePath))
             {
-                using var writer = new BinaryWriter(outgoingData, Encoding.Default, true);
-                writer.Write(true);
-                writer.Write(ex.Message);
+                throw new FileNotFoundException("Common Query template file not found", templatePath);
             }
+
+            string templateContent = File.ReadAllText(templatePath);
+            string finalHtml = templateContent.Replace("{query}", escapedQuery);
+            return finalHtml;
         }
 
+
         private static DatabaseProvider GetDatabaseProvider(DbCommand command)
         {
             return command.GetType().FullName switch

diff --git a/src/IQueryableObjectSource/OperationType.cs b/src/IQueryableObjectSource/OperationType.cs
@@ -0,0 +1,14 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+
+namespace IQueryableObjectSource
+{
+    public enum OperationType : byte
+    {
+        GetQuery = 1,
+        NotSupported = 9
+    }
+}