From 5bdb4202e3340b4fb6a613965141dec77c52c8e7 Mon Sep 17 00:00:00 2001
From: Erik Darling <2136037+erikdarlingdata@users.noreply.github.com>
Date: Mon, 6 Apr 2026 10:22:29 -0400
Subject: [PATCH] Add connection toolbar to plan viewer for schema lookups
 (#166)

Plan viewer now has Connect button + server status + database dropdown
in the toolbar, enabling schema lookups on plans opened from file.
Plans loaded from connected sessions inherit the connection and show
the server name and database pre-populated.

- Connect button shows connection dialog, populates database dropdown
- Database switching updates the connection string for schema lookups
- Connection services passed to PlanViewerControl from both MainWindow
  and QuerySessionControl
- Schema lookup items disabled when no connection, enabled after connect

Closes #166.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../Controls/PlanViewerControl.axaml          | 14 +++
 .../Controls/PlanViewerControl.axaml.cs       | 97 +++++++++++++++++++
 .../Controls/QuerySessionControl.axaml.cs     |  9 ++
 src/PlanViewer.App/MainWindow.axaml.cs        |  2 +
 4 files changed, 122 insertions(+)

diff --git a/src/PlanViewer.App/Controls/PlanViewerControl.axaml b/src/PlanViewer.App/Controls/PlanViewerControl.axaml
index ef2ed43..1eec781 100644
--- a/src/PlanViewer.App/Controls/PlanViewerControl.axaml
+++ b/src/PlanViewer.App/Controls/PlanViewerControl.axaml
@@ -14,6 +14,20 @@
                 BorderBrush="{DynamicResource BorderBrush}" BorderThickness="0,0,0,1">
             <DockPanel>
                 <StackPanel Orientation="Horizontal" DockPanel.Dock="Left">
+                    <Button x:Name="PlanConnectButton" Content="Connect" Click="PlanConnect_Click"
+                            Height="28" Padding="10,0" FontSize="12"
+                            Theme="{StaticResource AppButton}"
+                            ToolTip.Tip="Connect to a SQL Server for schema lookups"/>
+                    <TextBlock x:Name="PlanServerLabel" Text=""
+                               VerticalAlignment="Center" FontSize="12"
+                               Foreground="{DynamicResource ForegroundBrush}" Margin="6,0,0,0"/>
+                    <TextBlock Text="|" VerticalAlignment="Center"
+                               Foreground="{DynamicResource ForegroundBrush}" Margin="6,0"/>
+                    <ComboBox x:Name="PlanDatabaseBox" Width="180" Height="28" FontSize="12"
+                              IsEnabled="False" PlaceholderText="Database"
+                              SelectionChanged="PlanDatabase_SelectionChanged"/>
+                    <TextBlock Text="|" VerticalAlignment="Center"
+                               Foreground="{DynamicResource ForegroundBrush}" Margin="6,0"/>
                     <Button Content="+" Click="ZoomIn_Click" Width="28" Height="28" Padding="0" FontSize="16"
                             FontWeight="Bold" ToolTip.Tip="Zoom In"
                             Theme="{StaticResource AppButton}"/>
diff --git a/src/PlanViewer.App/Controls/PlanViewerControl.axaml.cs b/src/PlanViewer.App/Controls/PlanViewerControl.axaml.cs
index 3eef9f4..1051b23 100644
--- a/src/PlanViewer.App/Controls/PlanViewerControl.axaml.cs
+++ b/src/PlanViewer.App/Controls/PlanViewerControl.axaml.cs
@@ -16,7 +16,11 @@
 using Avalonia.Controls.Templates;
 using Avalonia.Platform.Storage;
 using AvaloniaEdit.TextMate;
+using Microsoft.Data.SqlClient;
+using PlanViewer.App.Dialogs;
+using PlanViewer.Core.Interfaces;
 using PlanViewer.App.Helpers;
+using PlanViewer.App.Services;
 using PlanViewer.App.Mcp;
 using PlanViewer.Core.Models;
 using PlanViewer.Core.Output;
@@ -173,6 +177,33 @@ public ServerMetadata? Metadata
     /// </summary>
     public string? ConnectionString { get; set; }
 
+    // Connection state for plans that connect via the toolbar
+    private ServerConnection? _planConnection;
+    private ICredentialService? _planCredentialService;
+    private ConnectionStore? _planConnectionStore;
+    private string? _planSelectedDatabase;
+
+    /// <summary>
+    /// Provide credential service and connection store so the plan viewer can show a connection dialog.
+    /// </summary>
+    public void SetConnectionServices(ICredentialService credentialService, ConnectionStore connectionStore)
+    {
+        _planCredentialService = credentialService;
+        _planConnectionStore = connectionStore;
+    }
+
+    /// <summary>
+    /// Update the connection UI to reflect an active connection (used when connection is inherited).
+    /// </summary>
+    public void SetConnectionStatus(string serverName, string? database)
+    {
+        PlanServerLabel.Text = serverName;
+        PlanServerLabel.Foreground = Brushes.LimeGreen;
+        PlanConnectButton.Content = "Reconnect";
+        if (database != null)
+            _planSelectedDatabase = database;
+    }
+
     // Events for MainWindow to wire up advice/repro actions
     public event EventHandler? HumanAdviceRequested;
     public event EventHandler? RobotAdviceRequested;
@@ -3348,6 +3379,72 @@ private IBrush FindBrushResource(string key)
 
     #endregion
 
+    #region Plan Viewer Connection
+
+    private async void PlanConnect_Click(object? sender, RoutedEventArgs e)
+    {
+        if (_planCredentialService == null || _planConnectionStore == null) return;
+
+        var dialog = new ConnectionDialog(_planCredentialService, _planConnectionStore);
+        var topLevel = TopLevel.GetTopLevel(this);
+        if (topLevel is not Window parentWindow) return;
+
+        var result = await dialog.ShowDialog<bool?>(parentWindow);
+        if (result != true || dialog.ResultConnection == null) return;
+
+        _planConnection = dialog.ResultConnection;
+        _planSelectedDatabase = dialog.ResultDatabase;
+        ConnectionString = _planConnection.GetConnectionString(_planCredentialService, _planSelectedDatabase);
+
+        PlanServerLabel.Text = _planConnection.ServerName;
+        PlanServerLabel.Foreground = Brushes.LimeGreen;
+        PlanConnectButton.Content = "Reconnect";
+
+        // Populate database dropdown
+        try
+        {
+            var connStr = _planConnection.GetConnectionString(_planCredentialService, "master");
+            await using var conn = new SqlConnection(connStr);
+            await conn.OpenAsync();
+
+            var databases = new List<string>();
+            using var cmd = new SqlCommand(
+                "SET TRANSACTION ISOLATION LEVEL READ UNCOMMITTED; SELECT name FROM sys.databases WHERE state_desc = 'ONLINE' ORDER BY name", conn);
+            using var reader = await cmd.ExecuteReaderAsync();
+            while (await reader.ReadAsync())
+                databases.Add(reader.GetString(0));
+
+            PlanDatabaseBox.ItemsSource = databases;
+            PlanDatabaseBox.IsEnabled = true;
+
+            if (_planSelectedDatabase != null)
+            {
+                for (int i = 0; i < PlanDatabaseBox.Items.Count; i++)
+                {
+                    if (PlanDatabaseBox.Items[i]?.ToString() == _planSelectedDatabase)
+                    {
+                        PlanDatabaseBox.SelectedIndex = i;
+                        break;
+                    }
+                }
+            }
+        }
+        catch
+        {
+            PlanDatabaseBox.IsEnabled = false;
+        }
+    }
+
+    private void PlanDatabase_SelectionChanged(object? sender, SelectionChangedEventArgs e)
+    {
+        if (_planConnection == null || _planCredentialService == null || PlanDatabaseBox.SelectedItem == null) return;
+
+        _planSelectedDatabase = PlanDatabaseBox.SelectedItem.ToString();
+        ConnectionString = _planConnection.GetConnectionString(_planCredentialService, _planSelectedDatabase);
+    }
+
+    #endregion
+
     #region Schema Lookup
 
     private static bool IsTempObject(string objectName)
diff --git a/src/PlanViewer.App/Controls/QuerySessionControl.axaml.cs b/src/PlanViewer.App/Controls/QuerySessionControl.axaml.cs
index f396c26..7da0deb 100644
--- a/src/PlanViewer.App/Controls/QuerySessionControl.axaml.cs
+++ b/src/PlanViewer.App/Controls/QuerySessionControl.axaml.cs
@@ -1075,6 +1075,9 @@ private async Task CaptureAndShowPlan(bool estimated, string? queryTextOverride
             var viewer = new PlanViewerControl();
             viewer.Metadata = _serverMetadata;
             viewer.ConnectionString = _connectionString;
+            viewer.SetConnectionServices(_credentialService, _connectionStore);
+            if (_serverConnection != null)
+                viewer.SetConnectionStatus(_serverConnection.ServerName, _selectedDatabase);
             viewer.OpenInEditorRequested += OnOpenInEditorRequested;
             viewer.LoadPlan(planXml, tabLabel, queryText);
             loadingTab.Content = viewer;
@@ -1158,6 +1161,9 @@ private void AddPlanTab(string planXml, string queryText, bool estimated, string
         var viewer = new PlanViewerControl();
         viewer.Metadata = _serverMetadata;
         viewer.ConnectionString = _connectionString;
+        viewer.SetConnectionServices(_credentialService, _connectionStore);
+        if (_serverConnection != null)
+            viewer.SetConnectionStatus(_serverConnection.ServerName, _selectedDatabase);
         viewer.OpenInEditorRequested += OnOpenInEditorRequested;
         viewer.LoadPlan(planXml, label, queryText);
 
@@ -1848,6 +1854,9 @@ private async void GetActualPlan_Click(object? sender, RoutedEventArgs e)
             var actualViewer = new PlanViewerControl();
             actualViewer.Metadata = _serverMetadata;
             actualViewer.ConnectionString = _connectionString;
+            actualViewer.SetConnectionServices(_credentialService, _connectionStore);
+            if (_serverConnection != null)
+                actualViewer.SetConnectionStatus(_serverConnection.ServerName, _selectedDatabase);
             actualViewer.OpenInEditorRequested += OnOpenInEditorRequested;
             actualViewer.LoadPlan(actualPlanXml, tabLabel, queryText);
             loadingTab.Content = actualViewer;
diff --git a/src/PlanViewer.App/MainWindow.axaml.cs b/src/PlanViewer.App/MainWindow.axaml.cs
index 92de531..457c9d5 100644
--- a/src/PlanViewer.App/MainWindow.axaml.cs
+++ b/src/PlanViewer.App/MainWindow.axaml.cs
@@ -361,6 +361,7 @@ private void LoadPlanFile(string filePath)
                 return;
 
             var viewer = new PlanViewerControl();
+            viewer.SetConnectionServices(_credentialService, _connectionStore);
             viewer.LoadPlan(xml, fileName);
             viewer.SourceFilePath = filePath;
 
@@ -399,6 +400,7 @@ private async Task PasteXmlAsync()
             return;
 
         var viewer = new PlanViewerControl();
+        viewer.SetConnectionServices(_credentialService, _connectionStore);
         viewer.LoadPlan(xml, "Pasted Plan");
 
         var content = CreatePlanTabContent(viewer);