Execute selection, from cursor, and current batch (#53)

src/PlanViewer.App/Controls/QuerySessionControl.axaml.cs

@@ -3,6 +3,7 @@
 using System.Diagnostics;
 using System.Linq;
 using System.Text.Json;
+using System.Text.RegularExpressions;
 using System.Threading;
 using System.Threading.Tasks;
 using System.Xml.Linq;
@@ -132,9 +133,25 @@ private void SetupEditorContextMenu()
             QueryEditor.SelectAll();
         };
 
+        var executeFromCursorItem = new MenuItem { Header = "Execute from Cursor" };
+        executeFromCursorItem.Click += async (_, _) =>
+        {
+            var text = GetTextFromCursor();
+            if (!string.IsNullOrWhiteSpace(text))
+                await CaptureAndShowPlan(estimated: false, queryTextOverride: text);
+        };
+
+        var executeCurrentBatchItem = new MenuItem { Header = "Execute Current Batch" };
+        executeCurrentBatchItem.Click += async (_, _) =>
+        {
+            var text = GetCurrentBatch();
+            if (!string.IsNullOrWhiteSpace(text))
+                await CaptureAndShowPlan(estimated: false, queryTextOverride: text);
+        };
+
         QueryEditor.TextArea.ContextMenu = new ContextMenu
         {
-            Items = { cutItem, copyItem, pasteItem, new Separator(), selectAllItem }
+            Items = { cutItem, copyItem, pasteItem, new Separator(), selectAllItem, new Separator(), executeFromCursorItem, executeCurrentBatchItem }
         };
     }
 
@@ -259,6 +276,47 @@ private void OnTextEntered(object? sender, TextInputEventArgs e)
         return (doc.GetText(start, offset - start), start);
     }
 
+    private string? GetSelectedTextOrNull()
+    {
+        var selection = QueryEditor.TextArea.Selection;
+        if (selection.IsEmpty) return null;
+        return selection.GetText();
+    }
+
+    private string GetTextFromCursor()
+    {
+        var doc = QueryEditor.Document;
+        var offset = QueryEditor.CaretOffset;
+        return doc.GetText(offset, doc.TextLength - offset);
+    }
+
+    private string? GetCurrentBatch()
+    {
+        var doc = QueryEditor.Document;
+        var caretOffset = QueryEditor.CaretOffset;
+        var text = doc.Text;
+        var goPattern = new Regex(@"^\s*GO\s*$", RegexOptions.IgnoreCase | RegexOptions.Multiline);
+        var matches = goPattern.Matches(text);
+
+        int batchStart = 0;
+        int batchEnd = text.Length;
+
+        foreach (Match m in matches)
+        {
+            if (m.Index + m.Length <= caretOffset)
+            {
+                batchStart = m.Index + m.Length;
+            }
+            else if (m.Index >= caretOffset)
+            {
+                batchEnd = m.Index;
+                break;
+            }
+        }
+
+        return text[batchStart..batchEnd].Trim();
+    }
+
     private void SetStatus(string text, bool autoClear = true)
     {
         _statusClearCts?.Cancel();
@@ -435,15 +493,17 @@ private async void ExecuteEstimated_Click(object? sender, RoutedEventArgs e)
         await CaptureAndShowPlan(estimated: true);
     }
 
-    private async Task CaptureAndShowPlan(bool estimated)
+    private async Task CaptureAndShowPlan(bool estimated, string? queryTextOverride = null)
     {
         if (_connectionString == null || _selectedDatabase == null)
         {
             SetStatus("Connect to a server first", autoClear: false);
             return;
         }
 
-        var queryText = QueryEditor.Text?.Trim();
+        var queryText = queryTextOverride?.Trim()
+                        ?? GetSelectedTextOrNull()?.Trim()
+                        ?? QueryEditor.Text?.Trim();
         if (string.IsNullOrEmpty(queryText))
         {
             SetStatus("Enter a query", autoClear: false);

src/PlanViewer.Core/Services/ActualPlanExecutor.cs

@@ -7,6 +7,7 @@
  */
 
 using System;
+using System.Collections.Generic;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
@@ -55,7 +56,7 @@ public static class ActualPlanExecutor
         sb.AppendLine("SET STATISTICS XML OFF;");
 
         var fullScript = sb.ToString();
-        string? capturedPlanXml = null;
+        var capturedPlanXmls = new List<string>();
 
         /* Override database in connection string */
         var builder = new SqlConnectionStringBuilder(connectionString);
@@ -89,7 +90,7 @@ The plan result set has a single row with a single XML column. */
                 if (value != null && value.TrimStart().StartsWith("<ShowPlanXML", StringComparison.Ordinal))
                 {
                     /* This is a plan XML result set — capture it */
-                    capturedPlanXml = value;
+                    capturedPlanXmls.Add(value);
                 }
                 else
                 {
@@ -105,6 +106,8 @@ The plan result set has a single row with a single XML column. */
         }
         while (await reader.NextResultAsync(cancellationToken));
 
-        return capturedPlanXml;
+        if (capturedPlanXmls.Count == 0) return null;
+        if (capturedPlanXmls.Count == 1) return capturedPlanXmls[0];
+        return EstimatedPlanExecutor.MergeShowPlanXmls(capturedPlanXmls);
     }
 }

src/PlanViewer.Core/Services/EstimatedPlanExecutor.cs

@@ -1,6 +1,9 @@
 using System;
+using System.Collections.Generic;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Xml.Linq;
 using Microsoft.Data.SqlClient;
 
 namespace PlanViewer.Core.Services;
@@ -36,9 +39,9 @@ public static class EstimatedPlanExecutor
             await enableCmd.ExecuteNonQueryAsync(cancellationToken);
         }
 
-        // Execute the query — with SHOWPLAN XML ON, this returns the plan
-        // as a single-row, single-column result set (no actual execution)
-        string? planXml = null;
+        // Execute the query — with SHOWPLAN XML ON, this returns one result set
+        // per statement, each containing a ShowPlanXML document.
+        var planXmls = new List<string>();
         using (var queryCmd = new SqlCommand(queryText, connection))
         {
             queryCmd.CommandTimeout = timeoutSeconds;
@@ -49,12 +52,16 @@ public static class EstimatedPlanExecutor
             });
 
             using var reader = await queryCmd.ExecuteReaderAsync(cancellationToken);
-            if (await reader.ReadAsync(cancellationToken))
+            do
             {
-                var value = reader.GetValue(0)?.ToString();
-                if (value != null && value.TrimStart().StartsWith("<ShowPlanXML", StringComparison.Ordinal))
-                    planXml = value;
+                if (await reader.ReadAsync(cancellationToken))
+                {
+                    var value = reader.GetValue(0)?.ToString();
+                    if (value != null && value.TrimStart().StartsWith("<ShowPlanXML", StringComparison.Ordinal))
+                        planXmls.Add(value);
+                }
             }
+            while (await reader.NextResultAsync(cancellationToken));
         }
 
         // Disable SHOWPLAN XML (best effort — connection is about to close)
@@ -66,6 +73,31 @@ public static class EstimatedPlanExecutor
         }
         catch { /* connection cleanup */ }
 
-        return planXml;
+        if (planXmls.Count == 0) return null;
+        if (planXmls.Count == 1) return planXmls[0];
+        return MergeShowPlanXmls(planXmls);
+    }
+
+    /// <summary>
+    /// Merges multiple ShowPlanXML documents into one by combining all Batch elements.
+    /// </summary>
+    internal static string MergeShowPlanXmls(List<string> planXmls)
+    {
+        XNamespace ns = "http://schemas.microsoft.com/sqlserver/2004/07/showplan";
+        var baseDoc = XDocument.Parse(planXmls[0]);
+        var batchSequence = baseDoc.Root!.Element(ns + "BatchSequence")!;
+
+        for (int i = 1; i < planXmls.Count; i++)
+        {
+            var doc = XDocument.Parse(planXmls[i]);
+            var batches = doc.Root!.Element(ns + "BatchSequence")?.Elements(ns + "Batch");
+            if (batches != null)
+            {
+                foreach (var batch in batches)
+                    batchSequence.Add(batch);
+            }
+        }
+
+        return baseDoc.ToString(SaveOptions.DisableFormatting);
     }
 }