diff --git a/docs/SPEC.md b/docs/SPEC.md index 9579626..b38e47c 100644 --- a/docs/SPEC.md +++ b/docs/SPEC.md @@ -108,7 +108,7 @@ ssmsx/ ```jsonc // Request (Rust → C#) -{"id":"abc-123","method":"query.execute","params":{"connectionId":"conn-1","sql":"SELECT 1"}} +{"id":"abc-123","method":"query.execute","params":{"sessionId":"tab-1","connectionId":"conn-1","database":"master","sql":"SELECT 1"}} // Response (C# → Rust) {"id":"abc-123","result":{"columns":[""],"rows":[[1]]}} diff --git a/package-lock.json b/package-lock.json index 1394119..bce1798 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1601,6 +1601,64 @@ "node": ">=14.0.0" } }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/core": { + "version": "1.8.1", + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "@emnapi/wasi-threads": "1.1.0", + "tslib": "^2.4.0" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/runtime": { + "version": "1.8.1", + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/wasi-threads": { + "version": "1.1.0", + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@napi-rs/wasm-runtime": { + "version": "1.1.1", + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "@emnapi/core": "^1.7.1", + "@emnapi/runtime": "^1.7.1", + "@tybys/wasm-util": "^0.10.1" + }, + "funding": { + "type": "github", + "url": "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/sponsors/Brooooooklyn" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@tybys/wasm-util": { + "version": "0.10.1", + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/tslib": { + "version": "2.8.1", + "inBundle": true, + "license": "0BSD", + "optional": true + }, "node_modules/@tailwindcss/oxide-win32-arm64-msvc": { "version": "4.2.1", "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-win32-arm64-msvc/-/oxide-win32-arm64-msvc-4.2.1.tgz", @@ -2374,9 +2432,9 @@ } }, "node_modules/baseline-browser-mapping": { - "version": "2.10.38", - "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.10.38.tgz", - "integrity": "sha512-31/02mVB4yuQU6adKk5SlY6m+mxDwUq5KZkyYgnLrrKl7TEm1+3PyDtDBz2kOv/wxZz41GHsvV1A/u6RmiyBvw==", + "version": "2.11.21", + "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.11.21.tgz", + "integrity": "sha512-uh8vpY/1/YyFkunIDFH/12p7/7VdPKA1hejMVEbdkEaWnUz0Hesvx5EbiU6XxjyHZIOju+ZMbQJkRh+es3/spQ==", "dev": true, "license": "Apache-2.0", "bin": { @@ -2397,9 +2455,9 @@ } }, "node_modules/browserslist": { - "version": "4.28.4", - "resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.4.tgz", - "integrity": "sha512-MTc8i/x9jBQd1iMw2CFGS+rwMa07eYjLR0CCTLDACl9xhxy+nIs3KeML/biicXtk9JrZ6dnnTatmc7ErPXIxqw==", + "version": "4.28.9", + "resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.9.tgz", + "integrity": "sha512-EWazOblFYUvlGZcfGhPUPmYh3nikUxBVb+y9MJun5f3hBi812X+8MSQTujLBtgK3cf51fJWbWfOjyeO954d+Eg==", "dev": true, "funding": [ { @@ -2417,11 +2475,11 @@ ], "license": "MIT", "dependencies": { - "baseline-browser-mapping": "^2.10.38", - "caniuse-lite": "^1.0.30001799", - "electron-to-chromium": "^1.5.376", - "node-releases": "^2.0.48", - "update-browserslist-db": "^1.2.3" + "baseline-browser-mapping": "^2.11.20", + "caniuse-lite": "^1.0.30001810", + "electron-to-chromium": "^1.5.420", + "node-releases": "^2.0.54", + "update-browserslist-db": "^1.3.2" }, "bin": { "browserslist": "cli.js" @@ -2431,9 +2489,9 @@ } }, "node_modules/caniuse-lite": { - "version": "1.0.30001799", - "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001799.tgz", - "integrity": "sha512-hG1bReV+OUU+MOqK4t/ZWI0tZOyz3rqS9XuhOUz1cIcbwBKjOyJEJuw9ER5JuNyqxNk8u/JUVbGibBOL1yrjFw==", + "version": "1.0.30001810", + "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001810.tgz", + "integrity": "sha512-TITQPUkaz+aVk5GL6NhOdwk1aEaNTSDPsGFWrTuhKGtjTF70jL/Oht2W4c6rXUe5fu7Ie19VIahAXHIIiWWNeg==", "dev": true, "funding": [ { @@ -2687,9 +2745,9 @@ } }, "node_modules/electron-to-chromium": { - "version": "1.5.377", - "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.377.tgz", - "integrity": "sha512-cH1jZgJHoezfTnKfKwnScpHywTFVnJUNITDPREFdhNjiuD502+QFpG0Qk7G8jhsV/f+CEAFlIrzP1fT+IMb92g==", + "version": "1.5.422", + "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.422.tgz", + "integrity": "sha512-UvA/32XqrLDdZSn7Jllo1AYNcWji/G0d5M0GTViE7KoGBiMunw3a34Sb2KO4ZZyrSEhqsxFoVhWWJshdyfKqJA==", "dev": true, "license": "ISC" }, @@ -3297,9 +3355,9 @@ "license": "MIT" }, "node_modules/nanoid": { - "version": "3.3.17", - "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.17.tgz", - "integrity": "sha512-xQLf0A3HOMlgHq0n247/LRuAOYmB7dXJ/DvAxGvsSBij45XtBSmQycu+F8ODbHwns/XyFZagyL1+J0Offw1E0g==", + "version": "3.3.18", + "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.18.tgz", + "integrity": "sha512-DTg4MJbGMWkfi6VZFdNt2/caMbQy4Ou+Op/hJQvGEWcnVfoA1QA+xzRKAzw9jD6+GVOOeYr/mIcuDSdug6F6+w==", "funding": [ { "type": "github", @@ -3315,9 +3373,9 @@ } }, "node_modules/node-releases": { - "version": "2.0.48", - "resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.48.tgz", - "integrity": "sha512-1uz8041X6LoI6ZSdZacM9lVY28vuzDlSKitnpbSNK0RfKoIJkX29NBPVEFXhnuSuEOA9Ww0xnPJ+ILWbGAv8DA==", + "version": "2.0.54", + "resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.54.tgz", + "integrity": "sha512-YHs7BmmcsdAI5Ozuf8JZo6PT0mv2GIWC9vMfvUC3dp65M8hn7Ux8CPL+2oBI7juNuj9d0ndhTcznq2ODBps9cQ==", "dev": true, "license": "MIT", "engines": { @@ -3775,9 +3833,9 @@ } }, "node_modules/update-browserslist-db": { - "version": "1.2.3", - "resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.2.3.tgz", - "integrity": "sha512-Js0m9cx+qOgDxo0eMiFGEueWztz+d4+M3rGlmKPT+T4IS/jP4ylw3Nwpu6cpTTP8R1MAC1kF4VbdLt3ARf209w==", + "version": "1.3.2", + "resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.3.2.tgz", + "integrity": "sha512-UQ+MSxlhRm1bzjhU+DcuXfjFO1FzNtqhK5+9Yvlp90ItDLk5vT932A0rFu619nf7RVS+Y/VeaUW1jaRDqZ8VJw==", "dev": true, "funding": [ { diff --git a/package.json b/package.json index 3a7ae5a..cdff52b 100644 --- a/package.json +++ b/package.json @@ -44,7 +44,7 @@ }, "overrides": { "dompurify": "3.4.13", - "nanoid": "3.3.17", + "nanoid": "3.3.18", "postcss": "8.5.23" } } diff --git a/scripts/smoke-sidecar.mjs b/scripts/smoke-sidecar.mjs index 43694cd..0bf9ac0 100644 --- a/scripts/smoke-sidecar.mjs +++ b/scripts/smoke-sidecar.mjs @@ -25,6 +25,7 @@ const child = spawn( ); const pending = new Map(); +const completedResponses = new Map(); const events = []; const errors = []; @@ -47,6 +48,7 @@ stdout.on("line", (line) => { entry.responses.push(message); if (message.error || !entry.streaming || message.result?.done === true) { pending.delete(message.id); + completedResponses.set(message.id, entry.responses); entry.resolve(entry.responses); } }); @@ -67,14 +69,14 @@ child.on("exit", (code) => { let seq = 0; -function send(method, params, { streaming = false, allowError = false } = {}) { +function beginSend(method, params, { streaming = false, allowError = false } = {}) { const id = `smoke-${++seq}`; const request = { id, method, params }; const promise = new Promise((resolve, reject) => { pending.set(id, { resolve, reject, responses: [], streaming }); }); child.stdin.write(`${JSON.stringify(request)}\n`); - return promise.then((responses) => { + const result = promise.then((responses) => { const failed = responses.find((response) => response.error); if (failed && !allowError) { throw new Error(`${method} failed: ${failed.error.code} ${failed.error.message}`); @@ -84,6 +86,40 @@ function send(method, params, { streaming = false, allowError = false } = {}) { } return streaming ? responses.map((response) => response.result) : responses.at(-1).result; }); + return { id, result }; +} + +function send(method, params, options) { + return beginSend(method, params, options).result; +} + +function waitForQueryStart(id, timeoutMs = 5000) { + return new Promise((resolve, reject) => { + const started = setInterval(() => { + const entry = pending.get(id); + const responses = entry?.responses ?? completedResponses.get(id); + const firstResponse = responses?.[0]; + if (firstResponse?.error) { + clearInterval(started); + clearTimeout(timeout); + completedResponses.delete(id); + reject( + new Error( + `query ${id} failed to start: ${firstResponse.error.code} ${firstResponse.error.message}` + ) + ); + } else if (firstResponse?.result?.queryId) { + clearInterval(started); + clearTimeout(timeout); + completedResponses.delete(id); + resolve(firstResponse.result); + } + }, 25); + const timeout = setTimeout(() => { + clearInterval(started); + reject(new Error(`timed out waiting for query ${id} to start`)); + }, timeoutMs); + }); } function assert(condition, message) { @@ -167,6 +203,7 @@ try { assert(savedConnections.length === 1, "duplicate connection.save should not add another recent connection"); const activeConnectionId = duplicateSaved.id; + const primarySessionId = "smoke-primary-session"; const connected = await send("connection.connect", { id: activeConnectionId }); assert(connected.connectionId === activeConnectionId, "connection.connect should activate connection"); @@ -200,6 +237,7 @@ try { await send( "query.execute", { + sessionId: primarySessionId, connectionId: activeConnectionId, database, sql: "SELECT TOP (5) BusinessEntityID, FirstName, LastName FROM Person.Person ORDER BY BusinessEntityID;", @@ -215,6 +253,7 @@ try { await send( "query.execute", { + sessionId: primarySessionId, connectionId: activeConnectionId, database, sql: "PRINT 'hello from ssmsx smoke'; SELECT COUNT(*) AS PersonCount FROM Person.Person;", @@ -229,6 +268,7 @@ try { await send( "query.execute", { + sessionId: primarySessionId, connectionId: activeConnectionId, database, sql: [ @@ -254,6 +294,7 @@ try { await send( "query.execute", { + sessionId: primarySessionId, connectionId: activeConnectionId, database, sql: [ @@ -285,6 +326,7 @@ try { await send( "query.execute", { + sessionId: primarySessionId, connectionId: activeConnectionId, database, sql: [ @@ -311,6 +353,7 @@ try { await send( "query.execute", { + sessionId: primarySessionId, connectionId: activeConnectionId, database, sql: [ @@ -328,6 +371,7 @@ try { const invalidResponses = await send( "query.execute", { + sessionId: primarySessionId, connectionId: activeConnectionId, database, sql: [ @@ -354,40 +398,124 @@ try { "affected-row messages should be delivered before the later query error" ); - const cancelId = `smoke-${++seq}`; - const cancelPromise = new Promise((resolve, reject) => { - pending.set(cancelId, { resolve, reject, responses: [], streaming: true }); - }); - child.stdin.write( - `${JSON.stringify({ - id: cancelId, - method: "query.execute", - params: { + const stateSessionId = "smoke-state-session"; + await send( + "query.execute", + { + sessionId: stateSessionId, + connectionId: activeConnectionId, + database, + sql: "CREATE TABLE #SsmsxSessionState (Value int NOT NULL); INSERT INTO #SsmsxSessionState VALUES (42);", + }, + { streaming: true } + ); + const persistedState = summarizeQueryResults( + await send( + "query.execute", + { + sessionId: stateSessionId, connectionId: activeConnectionId, database, - sql: "WAITFOR DELAY '00:00:10'; SELECT 1 AS ShouldNotNeedToFinish;", + sql: "SELECT Value FROM #SsmsxSessionState;", }, - })}\n` + { streaming: true } + ) ); - const firstBatch = await new Promise((resolve, reject) => { - const started = setInterval(() => { - const entry = pending.get(cancelId); - if (entry?.responses[0]?.result?.queryId) { - clearInterval(started); - resolve(entry.responses[0].result); - } - }, 25); - setTimeout(() => { - clearInterval(started); - reject(new Error("timed out waiting for cancellable queryId")); - }, 5000); + assert(persistedState.rows[0]?.[0] === 42, "a query window should retain its SQL session state"); + + const isolatedState = summarizeQueryResults( + await send( + "query.execute", + { + sessionId: "smoke-isolated-session", + connectionId: activeConnectionId, + database, + sql: "SELECT CASE WHEN OBJECT_ID('tempdb..#SsmsxSessionState') IS NULL THEN 1 ELSE 0 END AS IsIsolated;", + }, + { streaming: true } + ) + ); + assert(isolatedState.rows[0]?.[0] === 1, "SQL session state should not leak between query windows"); + + const sessionClosed = await send("query.sessionClose", { + sessionId: stateSessionId, + connectionId: activeConnectionId, }); - const cancelResult = await send("query.cancel", { queryId: firstBatch.queryId }); - assert(cancelResult.cancelled === true, "query.cancel should cancel the running query"); - const cancelled = summarizeQueryResults(await cancelPromise.then((responses) => responses.map((response) => response.result))); + assert(sessionClosed.closed === true, "query.sessionClose should release an existing query session"); + + const delayedQuery = beginSend( + "query.execute", + { + sessionId: "smoke-concurrent-delayed", + connectionId: activeConnectionId, + database, + sql: "WAITFOR DELAY '00:00:03'; SELECT 1 AS DelayedResult;", + }, + { streaming: true } + ); + await waitForQueryStart(delayedQuery.id); + const fastQuery = beginSend( + "query.execute", + { + sessionId: "smoke-concurrent-fast", + connectionId: activeConnectionId, + database, + sql: "SELECT 1 AS FastResult;", + }, + { streaming: true } + ); + const firstCompleted = await Promise.race([ + delayedQuery.result.then(() => "delayed"), + fastQuery.result.then(() => "fast"), + ]); + assert(firstCompleted === "fast", "a fast query window should not wait behind another window on the same connection"); + const fastResult = summarizeQueryResults(await fastQuery.result); + const delayedResult = summarizeQueryResults(await delayedQuery.result); + assert(fastResult.rows[0]?.[0] === 1, "concurrent fast query should return its result"); + assert(delayedResult.rows[0]?.[0] === 1, "concurrent delayed query should finish normally"); + + const firstCancellationQuery = beginSend( + "query.execute", + { + sessionId: "smoke-cancel-first", + connectionId: activeConnectionId, + database, + sql: "WAITFOR DELAY '00:00:10'; SELECT 1 AS FirstCancellationResult;", + }, + { streaming: true } + ); + const secondCancellationQuery = beginSend( + "query.execute", + { + sessionId: "smoke-cancel-second", + connectionId: activeConnectionId, + database, + sql: "WAITFOR DELAY '00:00:10'; SELECT 1 AS SecondCancellationResult;", + }, + { streaming: true } + ); + const [firstCancellationStart, secondCancellationStart] = await Promise.all([ + waitForQueryStart(firstCancellationQuery.id), + waitForQueryStart(secondCancellationQuery.id), + ]); + const firstCancelResult = await send("query.cancel", { queryId: firstCancellationStart.queryId }); + assert(firstCancelResult.cancelled === true, "query.cancel should cancel the selected query window"); + const firstCancelled = summarizeQueryResults(await firstCancellationQuery.result); + assert( + firstCancelled.messages.some((message) => /cancelled|canceled/i.test(message.text)), + "the selected query window should return a cancellation message" + ); + const secondQueryState = await Promise.race([ + secondCancellationQuery.result.then(() => "finished"), + new Promise((resolve) => setTimeout(() => resolve("running"), 250)), + ]); + assert(secondQueryState === "running", "cancelling one query window should not cancel another"); + const secondCancelResult = await send("query.cancel", { queryId: secondCancellationStart.queryId }); + assert(secondCancelResult.cancelled === true, "the second query window should be independently cancellable"); + const secondCancelled = summarizeQueryResults(await secondCancellationQuery.result); assert( - cancelled.messages.some((message) => /cancelled|canceled/i.test(message.text)), - "cancelled query should return a cancellation message" + secondCancelled.messages.some((message) => /cancelled|canceled/i.test(message.text)), + "the second query window should return its own cancellation message" ); const diagram = await send("explorer.databaseDiagram", { connectionId: activeConnectionId, database }); @@ -400,6 +528,7 @@ try { ); await send("connection.disconnect", { id: activeConnectionId }); + assert(errors.length === 0, `sidecar protocol errors: ${errors.join("; ")}`); console.log( JSON.stringify( @@ -413,6 +542,7 @@ try { selectRows: select.rows.length, multiResultSets: multiResult.resultSets.map((set) => set.rows.length), printMessages: messages.messages.length, + concurrentQueryOrder: firstCompleted, stderr: events.filter(Boolean).slice(-5), }, null, diff --git a/sidecar/src/Ssmsx.Core/Connections/ConnectionManager.cs b/sidecar/src/Ssmsx.Core/Connections/ConnectionManager.cs index f2f996c..1caa215 100644 --- a/sidecar/src/Ssmsx.Core/Connections/ConnectionManager.cs +++ b/sidecar/src/Ssmsx.Core/Connections/ConnectionManager.cs @@ -8,7 +8,7 @@ namespace Ssmsx.Core.Connections; public class ConnectionManager { - private readonly ConcurrentDictionary _connections = new(); + private readonly ConcurrentDictionary _connections = new(); private readonly SqlConnectionFactory _factory; public ConnectionManager(SqlConnectionFactory? factory = null) @@ -42,9 +42,9 @@ public async Task ConnectAsync( // Dispose any existing connection with same ID only after the new // connection has passed cancellation and persistence checks. if (_connections.TryRemove(connectionId, out var existing)) - await existing.DisposeAsync(); + await existing.Connection.DisposeAsync(); - _connections[connectionId] = connection; + _connections[connectionId] = new ActiveConnection(connection, info, credentialStore); return connectionId; } catch @@ -56,17 +56,54 @@ public async Task ConnectAsync( public async Task DisconnectAsync(string connectionId) { - if (_connections.TryRemove(connectionId, out var connection)) - await connection.DisposeAsync(); + if (_connections.TryRemove(connectionId, out var activeConnection)) + await activeConnection.Connection.DisposeAsync(); } public SqlConnection GetConnection(string connectionId) { - if (_connections.TryGetValue(connectionId, out var connection)) - return connection; + if (_connections.TryGetValue(connectionId, out var activeConnection)) + return activeConnection.Connection; throw new InvalidOperationException($"No active connection for '{connectionId}'"); } + /// + /// Opens an independent pooled connection using the profile and credential context + /// that established the currently active saved connection. + /// + public async Task CreateQueryConnectionAsync( + string connectionId, + CancellationToken ct = default) + { + ArgumentException.ThrowIfNullOrWhiteSpace(connectionId); + + if (!_connections.TryGetValue(connectionId, out var activeConnection)) + throw new InvalidOperationException($"No active connection for '{connectionId}'"); + + var connection = await _factory.CreateAsync( + activeConnection.Info, + activeConnection.CredentialStore, + ct: ct); + try + { + ct.ThrowIfCancellationRequested(); + + if (!_connections.TryGetValue(connectionId, out var current) || + !ReferenceEquals(current, activeConnection)) + { + throw new InvalidOperationException( + $"Connection '{connectionId}' changed while the query session was opening"); + } + + return connection; + } + catch + { + await connection.DisposeAsync(); + throw; + } + } + public async Task TestAsync( ConnectionInfo info, ICredentialStore credentialStore, @@ -76,4 +113,9 @@ public async Task TestAsync( await using var connection = await _factory.CreateAsync(info, credentialStore, inlinePassword, ct); // Connection opened and immediately disposed — test passed } + + private sealed record ActiveConnection( + SqlConnection Connection, + ConnectionInfo Info, + ICredentialStore CredentialStore); } diff --git a/sidecar/src/Ssmsx.Core/Query/QueryCancellationManager.cs b/sidecar/src/Ssmsx.Core/Query/QueryCancellationManager.cs index 862752b..a8df4d4 100644 --- a/sidecar/src/Ssmsx.Core/Query/QueryCancellationManager.cs +++ b/sidecar/src/Ssmsx.Core/Query/QueryCancellationManager.cs @@ -23,7 +23,8 @@ public void Register(string queryId, CancellationTokenSource cts) throw new ArgumentException("Query ID cannot be null or empty.", nameof(queryId)); ArgumentNullException.ThrowIfNull(cts); - _activeQueries[queryId] = new QueryState(cts); + if (!_activeQueries.TryAdd(queryId, new QueryState(cts))) + throw new InvalidOperationException($"Query '{queryId}' is already registered"); } /// @@ -40,25 +41,8 @@ public void SetCommand(string queryId, SqlCommand cmd) throw new ArgumentException("Query ID cannot be null or empty.", nameof(queryId)); ArgumentNullException.ThrowIfNull(cmd); - if (_activeQueries.TryGetValue(queryId, out var entry)) - { - entry.Cmd = cmd; - - // If Cancel() was called before the command was registered, we stored - // the intent in Cancelled — apply it now so the query is interrupted - // immediately rather than silently running to completion. - if (entry.Cancelled) - { - try - { - cmd.Cancel(); - } - catch (Exception ex) - { - Console.Error.WriteLine($"Warning: Failed to cancel late-registered SqlCommand for query '{queryId}': {ex.Message}"); - } - } - } + if (_activeQueries.TryGetValue(queryId, out var entry) && entry.SetCommand(cmd)) + CancelCommand(queryId, cmd, "late-registered"); } /// @@ -75,13 +59,12 @@ public bool Cancel(string queryId) if (string.IsNullOrWhiteSpace(queryId)) return false; - // Don't remove the entry — SetCommand may still be racing to register the - // command. Mark Cancelled so a late SetCommand cancels immediately. The - // entry is cleaned up by Remove() when the query actually finishes. + // Don't remove the entry: SetCommand may still be racing to register the + // command. The entry is cleaned up when the query finishes. if (!_activeQueries.TryGetValue(queryId, out var entry)) return false; - entry.Cancelled = true; + var command = entry.MarkCancelled(); try { @@ -92,21 +75,15 @@ public bool Cancel(string queryId) Console.Error.WriteLine($"Warning: CancellationTokenSource for query '{queryId}' was already disposed: {ex.Message}"); } - try - { - entry.Cmd?.Cancel(); - } - catch (Exception ex) - { - Console.Error.WriteLine($"Warning: Failed to cancel SqlCommand for query '{queryId}': {ex.Message}"); - } + if (command is not null) + CancelCommand(queryId, command, "active"); return true; } /// - /// Removes a query from tracking after it has completed (successfully or otherwise). - /// Disposes the CancellationTokenSource. + /// Removes a query from tracking after it has completed and disposes its + /// cancellation token source. /// /// The unique identifier for the query to remove. public void Remove(string queryId) @@ -127,20 +104,48 @@ public void Remove(string queryId) } } - private sealed class QueryState + private static void CancelCommand(string queryId, SqlCommand command, string context) { - public CancellationTokenSource Cts { get; } - public SqlCommand? Cmd { get; set; } + try + { + command.Cancel(); + } + catch (Exception ex) + { + Console.Error.WriteLine( + $"Warning: Failed to cancel {context} SqlCommand for query '{queryId}': {ex.Message}"); + } + } - /// - /// Set to true when Cancel() was called. If the SqlCommand arrives - /// after this, SetCommand will cancel it immediately. - /// - public bool Cancelled { get; set; } + private sealed class QueryState + { + private readonly object _gate = new(); + private SqlCommand? _command; + private bool _cancelled; public QueryState(CancellationTokenSource cts) { Cts = cts; } + + public CancellationTokenSource Cts { get; } + + public bool SetCommand(SqlCommand command) + { + lock (_gate) + { + _command = command; + return _cancelled; + } + } + + public SqlCommand? MarkCancelled() + { + lock (_gate) + { + _cancelled = true; + return _command; + } + } } } diff --git a/sidecar/src/Ssmsx.Core/Query/QueryExecutor.cs b/sidecar/src/Ssmsx.Core/Query/QueryExecutor.cs index 3282d1a..1f3329e 100644 --- a/sidecar/src/Ssmsx.Core/Query/QueryExecutor.cs +++ b/sidecar/src/Ssmsx.Core/Query/QueryExecutor.cs @@ -1,9 +1,9 @@ using System.Data.SqlTypes; +using System.Data; using System.Diagnostics; using System.Text; using System.Text.RegularExpressions; using Microsoft.Data.SqlClient; -using Ssmsx.Core.Connections; using Ssmsx.Protocol.Messages; using Ssmsx.Protocol.Models; @@ -21,7 +21,7 @@ public class QueryExecutor @"^[\t ]*GO(?:[\t ]+(?[1-9][0-9]*))?[\t ]*(?:/\*(?:[^*]|\*(?!/))*\*/[\t ]*)*(?:--[^\r\n]*)?\r?$", RegexOptions.IgnoreCase | RegexOptions.Multiline | RegexOptions.CultureInvariant); - private readonly ConnectionManager _connectionManager; + private readonly QuerySessionManager _sessionManager; private readonly QueryCancellationManager _cancellationManager; /// @@ -35,9 +35,9 @@ public class QueryExecutor /// private const int DefaultCommandTimeoutSeconds = 0; - public QueryExecutor(ConnectionManager connectionManager, QueryCancellationManager cancellationManager) + public QueryExecutor(QuerySessionManager sessionManager, QueryCancellationManager cancellationManager) { - _connectionManager = connectionManager ?? throw new ArgumentNullException(nameof(connectionManager)); + _sessionManager = sessionManager ?? throw new ArgumentNullException(nameof(sessionManager)); _cancellationManager = cancellationManager ?? throw new ArgumentNullException(nameof(cancellationManager)); } @@ -45,6 +45,7 @@ public QueryExecutor(ConnectionManager connectionManager, QueryCancellationManag /// Executes a SQL query, streaming result batches via the onBatch callback. /// Supports multiple result sets and cooperative cancellation. /// + /// The ID of the query window whose SQL session should be used. /// The ID of the active connection to use. /// The database context to execute the query against. /// The SQL text to execute. @@ -52,6 +53,7 @@ public QueryExecutor(ConnectionManager connectionManager, QueryCancellationManag /// Callback invoked for each batch of results. /// Cancellation token for cooperative cancellation. public async Task ExecuteAsync( + string sessionId, string connectionId, string database, string sql, @@ -59,6 +61,8 @@ public async Task ExecuteAsync( Func onBatch, CancellationToken ct) { + if (string.IsNullOrWhiteSpace(sessionId)) + throw new ArgumentException("Session ID cannot be null or empty.", nameof(sessionId)); if (string.IsNullOrWhiteSpace(connectionId)) throw new ArgumentException("Connection ID cannot be null or empty.", nameof(connectionId)); if (string.IsNullOrWhiteSpace(database)) @@ -69,7 +73,9 @@ public async Task ExecuteAsync( throw new ArgumentException("Query ID cannot be null or empty.", nameof(queryId)); ArgumentNullException.ThrowIfNull(onBatch); - var connection = _connectionManager.GetConnection(connectionId); + await using var session = await _sessionManager.AcquireAsync(sessionId, connectionId, ct); + var connection = session.Connection; + ct = session.CancellationToken; var messages = new List(); var stopwatch = new Stopwatch(); long totalRows = 0; @@ -240,9 +246,11 @@ public async Task ExecuteAsync( await onBatch(cancelledBatch); } - catch (SqlException) + catch (SqlException ex) { stopwatch.Stop(); + if (ShouldEvictSession(connection.State, ex.Class)) + session.MarkBroken(); if (messages.Count > 0) { @@ -283,6 +291,9 @@ public async Task ExecuteAsync( }; } + internal static bool ShouldEvictSession(ConnectionState connectionState, byte errorClass) => + connectionState != ConnectionState.Open || errorClass >= 20; + internal static IEnumerable SplitBatches(string sql) { var currentBatch = new StringBuilder(); diff --git a/sidecar/src/Ssmsx.Core/Query/QuerySessionManager.cs b/sidecar/src/Ssmsx.Core/Query/QuerySessionManager.cs new file mode 100644 index 0000000..2d09958 --- /dev/null +++ b/sidecar/src/Ssmsx.Core/Query/QuerySessionManager.cs @@ -0,0 +1,294 @@ +using System.Collections.Concurrent; +using Microsoft.Data.SqlClient; +using Ssmsx.Core.Connections; + +namespace Ssmsx.Core.Query; + +/// +/// Owns one persistent SQL connection per query window and prevents overlapping +/// commands from sharing that session. +/// +public sealed class QuerySessionManager : IAsyncDisposable +{ + private readonly ConcurrentDictionary _sessions = + new(StringComparer.Ordinal); + private readonly Func> _openConnection; + private int _disposed; + + public QuerySessionManager(ConnectionManager connectionManager) + { + ArgumentNullException.ThrowIfNull(connectionManager); + _openConnection = connectionManager.CreateQueryConnectionAsync; + } + + internal QuerySessionManager(Func> openConnection) + { + _openConnection = openConnection ?? throw new ArgumentNullException(nameof(openConnection)); + } + + internal async Task AcquireAsync( + string sessionId, + string connectionId, + CancellationToken ct = default) + { + ArgumentException.ThrowIfNullOrWhiteSpace(sessionId); + ArgumentException.ThrowIfNullOrWhiteSpace(connectionId); + ObjectDisposedException.ThrowIf(Volatile.Read(ref _disposed) != 0, this); + + var session = await GetOrReplaceSessionAsync(sessionId, connectionId); + if (Volatile.Read(ref _disposed) != 0) + { + await RemoveAndDisposeAsync(sessionId, session); + throw new ObjectDisposedException(nameof(QuerySessionManager)); + } + + if (!session.TryBeginExecution()) + throw new InvalidOperationException($"Query session '{sessionId}' is already executing a query"); + + var executionCancellation = CancellationTokenSource.CreateLinkedTokenSource( + ct, + session.LifetimeToken); + try + { + var connection = await session.GetConnectionAsync(executionCancellation.Token); + return new QuerySessionLease( + this, + sessionId, + session, + connection, + executionCancellation); + } + catch + { + executionCancellation.Dispose(); + session.EndExecution(); + await RemoveAndDisposeAsync(sessionId, session); + throw; + } + } + + /// + /// Closes a session only when it still belongs to the expected connection. + /// This prevents a delayed close for an old tab target from closing its replacement. + /// + public Task CloseSessionAsync(string sessionId, string connectionId) + { + ArgumentException.ThrowIfNullOrWhiteSpace(connectionId); + return CloseSessionCoreAsync(sessionId, connectionId); + } + + public async Task CloseConnectionSessionsAsync(string connectionId) + { + ArgumentException.ThrowIfNullOrWhiteSpace(connectionId); + + var removed = new List(); + foreach (var entry in _sessions) + { + if (string.Equals(entry.Value.ConnectionId, connectionId, StringComparison.Ordinal) && + _sessions.TryRemove(entry)) + { + removed.Add(entry.Value); + } + } + + foreach (var session in removed) + await session.DisposeAsync(); + + return removed.Count; + } + + public async ValueTask DisposeAsync() + { + if (Interlocked.Exchange(ref _disposed, 1) != 0) + return; + + foreach (var entry in _sessions) + { + if (_sessions.TryRemove(entry)) + await entry.Value.DisposeAsync(); + } + } + + private async Task GetOrReplaceSessionAsync( + string sessionId, + string connectionId) + { + while (true) + { + var candidate = new QuerySession(connectionId, _openConnection); + if (_sessions.TryAdd(sessionId, candidate)) + return candidate; + + if (!_sessions.TryGetValue(sessionId, out var existing)) + { + await candidate.DisposeAsync(); + continue; + } + + if (string.Equals(existing.ConnectionId, connectionId, StringComparison.Ordinal)) + { + await candidate.DisposeAsync(); + return existing; + } + + if (_sessions.TryUpdate(sessionId, candidate, existing)) + { + await existing.DisposeAsync(); + return candidate; + } + + await candidate.DisposeAsync(); + } + } + + private async Task CloseSessionCoreAsync( + string sessionId, + string expectedConnectionId) + { + ArgumentException.ThrowIfNullOrWhiteSpace(sessionId); + + while (_sessions.TryGetValue(sessionId, out var session)) + { + if (!string.Equals(session.ConnectionId, expectedConnectionId, StringComparison.Ordinal)) + { + return false; + } + + if (_sessions.TryRemove(new KeyValuePair(sessionId, session))) + { + await session.DisposeAsync(); + return true; + } + } + + return false; + } + + private async ValueTask ReleaseAsync( + string sessionId, + QuerySession session, + bool broken) + { + if (broken) + { + // Unpublish the broken session before another execution can lease it. + await RemoveAndDisposeAsync(sessionId, session); + } + + session.EndExecution(); + } + + private async Task RemoveAndDisposeAsync(string sessionId, QuerySession session) + { + if (_sessions.TryRemove(new KeyValuePair(sessionId, session))) + await session.DisposeAsync(); + } + + internal sealed class QuerySessionLease : IAsyncDisposable + { + private readonly QuerySessionManager _manager; + private readonly string _sessionId; + private readonly QuerySession _session; + private readonly CancellationTokenSource _executionCancellation; + private int _disposed; + private int _broken; + + internal QuerySessionLease( + QuerySessionManager manager, + string sessionId, + QuerySession session, + SqlConnection connection, + CancellationTokenSource executionCancellation) + { + _manager = manager; + _sessionId = sessionId; + _session = session; + Connection = connection; + _executionCancellation = executionCancellation; + } + + internal SqlConnection Connection { get; } + + internal CancellationToken CancellationToken => _executionCancellation.Token; + + internal void MarkBroken() => Interlocked.Exchange(ref _broken, 1); + + public async ValueTask DisposeAsync() + { + if (Interlocked.Exchange(ref _disposed, 1) != 0) + return; + + _executionCancellation.Dispose(); + await _manager.ReleaseAsync( + _sessionId, + _session, + Volatile.Read(ref _broken) != 0); + } + } + + internal sealed class QuerySession : IAsyncDisposable + { + private readonly object _connectionGate = new(); + private readonly Func> _openConnection; + private readonly CancellationTokenSource _lifetimeCancellation = new(); + private Task? _connectionTask; + private int _executing; + private int _disposed; + + internal QuerySession( + string connectionId, + Func> openConnection) + { + ConnectionId = connectionId; + _openConnection = openConnection; + } + + internal string ConnectionId { get; } + + internal CancellationToken LifetimeToken => _lifetimeCancellation.Token; + + internal bool TryBeginExecution() => + Interlocked.CompareExchange(ref _executing, 1, 0) == 0; + + internal void EndExecution() => Interlocked.Exchange(ref _executing, 0); + + internal Task GetConnectionAsync(CancellationToken ct) + { + lock (_connectionGate) + { + ObjectDisposedException.ThrowIf( + Volatile.Read(ref _disposed) != 0, + nameof(QuerySession)); + return _connectionTask ??= _openConnection(ConnectionId, ct); + } + } + + public async ValueTask DisposeAsync() + { + if (Interlocked.Exchange(ref _disposed, 1) != 0) + return; + + await _lifetimeCancellation.CancelAsync(); + + Task? connectionTask; + lock (_connectionGate) + connectionTask = _connectionTask; + + if (connectionTask is not null) + { + try + { + var connection = await connectionTask; + await connection.DisposeAsync(); + } + catch (Exception ex) + { + Console.Error.WriteLine( + $"Warning: Failed to close query session for connection '{ConnectionId}': {ex.Message}"); + } + } + + _lifetimeCancellation.Dispose(); + } + } +} diff --git a/sidecar/src/Ssmsx.Protocol/Messages/QueryMessages.cs b/sidecar/src/Ssmsx.Protocol/Messages/QueryMessages.cs index 33c5aa3..e9d7736 100644 --- a/sidecar/src/Ssmsx.Protocol/Messages/QueryMessages.cs +++ b/sidecar/src/Ssmsx.Protocol/Messages/QueryMessages.cs @@ -5,6 +5,9 @@ namespace Ssmsx.Protocol.Messages; public record QueryExecuteParams { + [JsonPropertyName("sessionId")] + public required string SessionId { get; init; } + [JsonPropertyName("connectionId")] public required string ConnectionId { get; init; } @@ -15,6 +18,15 @@ public record QueryExecuteParams public required string Sql { get; init; } } +public record QuerySessionCloseParams +{ + [JsonPropertyName("sessionId")] + public required string SessionId { get; init; } + + [JsonPropertyName("connectionId")] + public required string ConnectionId { get; init; } +} + public record QueryCancelParams { [JsonPropertyName("queryId")] @@ -59,3 +71,9 @@ public record QueryCancelResult [JsonPropertyName("cancelled")] public bool Cancelled { get; init; } } + +public record QuerySessionCloseResult +{ + [JsonPropertyName("closed")] + public bool Closed { get; init; } +} diff --git a/sidecar/src/Ssmsx.Protocol/ProtocolJsonContext.cs b/sidecar/src/Ssmsx.Protocol/ProtocolJsonContext.cs index ed78646..543da70 100644 --- a/sidecar/src/Ssmsx.Protocol/ProtocolJsonContext.cs +++ b/sidecar/src/Ssmsx.Protocol/ProtocolJsonContext.cs @@ -85,6 +85,8 @@ namespace Ssmsx.Protocol; [JsonSerializable(typeof(QueryExecuteResult))] [JsonSerializable(typeof(QueryCancelParams))] [JsonSerializable(typeof(QueryCancelResult))] +[JsonSerializable(typeof(QuerySessionCloseParams))] +[JsonSerializable(typeof(QuerySessionCloseResult))] [JsonSerializable(typeof(List>))] [JsonSerializable(typeof(List))] [JsonSerializable(typeof(List))] diff --git a/sidecar/src/Ssmsx.Sidecar/Program.cs b/sidecar/src/Ssmsx.Sidecar/Program.cs index 1ed64cf..c3fdc5d 100644 --- a/sidecar/src/Ssmsx.Sidecar/Program.cs +++ b/sidecar/src/Ssmsx.Sidecar/Program.cs @@ -20,10 +20,12 @@ var connectionOperations = new ConnectionOperationCoordinator(); var schemaDiscovery = new SchemaDiscoveryService(connectionManager); var queryCancellationManager = new QueryCancellationManager(); -var queryExecutor = new QueryExecutor(connectionManager, queryCancellationManager); +var querySessionManager = new QuerySessionManager(connectionManager); +var queryExecutor = new QueryExecutor(querySessionManager, queryCancellationManager); var intelliSenseService = new IntelliSenseService(connectionManager); var requestCancellation = new ConcurrentDictionary(); var cancelledRequestIds = new ConcurrentDictionary(); +var activeQueryTasks = new ConcurrentDictionary(); var cancelledRequestIdTtl = TimeSpan.FromMinutes(5); await using var stdout = Console.OpenStandardOutput(); @@ -145,9 +147,19 @@ bool TryConsumeCancelledRequestId(string requestId) { var args = Deserialize(p, ProtocolJsonContext.Default.ConnectionDisconnectParams); await connectionManager.DisconnectAsync(args.Id); + await querySessionManager.CloseConnectionSessionsAsync(args.Id); return JsonSerializer.SerializeToElement(true, ProtocolJsonContext.Default.Boolean); }, + ["query.sessionClose"] = async p => + { + var args = Deserialize(p, ProtocolJsonContext.Default.QuerySessionCloseParams); + var closed = await querySessionManager.CloseSessionAsync(args.SessionId, args.ConnectionId); + return JsonSerializer.SerializeToElement( + new QuerySessionCloseResult { Closed = closed }, + ProtocolJsonContext.Default.QuerySessionCloseResult); + }, + ["explorer.databases"] = async p => { var args = Deserialize(p, ProtocolJsonContext.Default.ExplorerDatabasesParams); @@ -252,6 +264,7 @@ static string SidecarVersion() return await connectionOperations.RunAsync(args.Id, async () => { await connectionManager.DisconnectAsync(args.Id); + await querySessionManager.CloseConnectionSessionsAsync(args.Id); try { await credentialStore.DeleteAsync($"ssmsx/{args.Id}"); } catch (Exception ex) { await Console.Error.WriteLineAsync($"Warning: Failed to delete credential for connection '{args.Id}': {ex.Message}"); } var deleted = await connectionStore.DeleteAsync(args.Id); @@ -290,6 +303,7 @@ static string SidecarVersion() return await connectionOperations.RunAsync(args.Id, async () => { var connId = await connectionManager.ConnectAsync(args.Id, connectionStore, credentialStore, ct); + await querySessionManager.CloseConnectionSessionsAsync(args.Id); return JsonSerializer.SerializeToElement( new ConnectionConnectResult { ConnectionId = connId }, ProtocolJsonContext.Default.ConnectionConnectResult); @@ -298,15 +312,11 @@ static string SidecarVersion() }; // --- Request processing --- -// We use a concurrent queue + semaphore so a background stdin reader can enqueue -// requests while the main processing loop handles them one at a time. -// This allows query.cancel to be queued and processed between query batches, -// since async awaits yield control to process the next queued request. +// A dedicated reader keeps cancellation responsive while the main loop dispatches +// query executions without waiting for other query windows to finish. var requestQueue = new BlockingCollection(); -// Background stdin reader — reads lines and enqueues them. -// Cancel requests are handled inline on this thread (SqlCommand.Cancel is thread-safe) -// so they work even while a query is blocking the main processing loop. +// SqlCommand.Cancel is thread-safe, so query cancellation can bypass the queue. var stdinReader = new Thread(() => { string? inputLine; @@ -386,7 +396,6 @@ static string SidecarVersion() }; stdinReader.Start(); -// Main processing loop — processes requests from the queue foreach (var requestLine in requestQueue.GetConsumingEnumerable()) { string requestId = "unknown"; @@ -404,7 +413,7 @@ static string SidecarVersion() // query.execute streams multiple responses, then completes if (request.Method == "query.execute") { - await HandleQueryExecute(request.Id, request.Params); + DispatchQueryExecute(request.Id, request.Params); continue; } @@ -440,13 +449,59 @@ static string SidecarVersion() } } +foreach (var queryId in activeQueryTasks.Keys) + queryCancellationManager.Cancel(queryId); +await Task.WhenAll(activeQueryTasks.Values); +await querySessionManager.DisposeAsync(); + // --- query.execute handler --- -async Task HandleQueryExecute(string requestId, JsonElement? p) +void DispatchQueryExecute(string requestId, JsonElement? p) +{ + var completion = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + if (!activeQueryTasks.TryAdd(requestId, completion.Task)) + { + SendError(requestId, "INVALID_REQUEST", $"A query request with ID '{requestId}' is already active"); + return; + } + + var cts = new CancellationTokenSource(); + try + { + queryCancellationManager.Register(requestId, cts); + } + catch (Exception ex) + { + activeQueryTasks.TryRemove(requestId, out _); + completion.TrySetResult(); + cts.Dispose(); + SendError(requestId, "INVALID_REQUEST", ex.Message); + return; + } + + _ = Task.Run(async () => + { + try + { + await HandleQueryExecute(requestId, p, cts); + } + catch (Exception ex) + { + await Console.Error.WriteLineAsync($"Unexpected error dispatching query '{requestId}': {ex}"); + SendError(requestId, "INTERNAL_ERROR", ex.Message); + } + finally + { + queryCancellationManager.Remove(requestId); + completion.TrySetResult(); + activeQueryTasks.TryRemove(requestId, out _); + } + }); +} + +async Task HandleQueryExecute(string requestId, JsonElement? p, CancellationTokenSource cts) { var args = Deserialize(p, ProtocolJsonContext.Default.QueryExecuteParams); var queryId = requestId; - using var cts = new CancellationTokenSource(); - queryCancellationManager.Register(queryId, cts); // Send immediate "started" response so frontend knows the queryId for cancellation { @@ -462,6 +517,7 @@ async Task HandleQueryExecute(string requestId, JsonElement? p) try { await queryExecutor.ExecuteAsync( + args.SessionId, args.ConnectionId, args.Database, args.Sql, diff --git a/sidecar/tests/Ssmsx.Core.Tests/Connections/ConnectionManagerTests.cs b/sidecar/tests/Ssmsx.Core.Tests/Connections/ConnectionManagerTests.cs index 604fef6..1351e19 100644 --- a/sidecar/tests/Ssmsx.Core.Tests/Connections/ConnectionManagerTests.cs +++ b/sidecar/tests/Ssmsx.Core.Tests/Connections/ConnectionManagerTests.cs @@ -29,7 +29,7 @@ public async Task ConnectAsync_DisposesOpenedConnection_WhenSavedConnectionWasDe var openedConnection = new SqlConnection(); var disposed = false; openedConnection.Disposed += (_, _) => disposed = true; - var factory = new FakeSqlConnectionFactory(async () => + var factory = new FakeSqlConnectionFactory(async (_, _, _) => { Assert.True(await store.DeleteAsync(info.Id)); return openedConnection; @@ -44,14 +44,125 @@ public async Task ConnectAsync_DisposesOpenedConnection_WhenSavedConnectionWasDe Assert.Throws(() => manager.GetConnection(info.Id)); } + [Fact] + public async Task ConnectAsync_PreservesActiveConnection_WhenReplacementFails() + { + var store = new ConnectionStore(_tempDir); + var info = await store.SaveAsync(new ConnectionInfo + { + Id = "replacement-fails", + ServerName = "server.example.com" + }); + var activeConnection = new SqlConnection(); + var managerConnectionAttempts = 0; + var factory = new FakeSqlConnectionFactory((_, _, _) => + { + if (managerConnectionAttempts++ == 0) + return Task.FromResult(activeConnection); + + throw new InvalidOperationException("Replacement failed"); + }); + var manager = new ConnectionManager(factory); + + await manager.ConnectAsync(info.Id, store, new FakeCredentialStore()); + + var error = await Assert.ThrowsAsync(() => + manager.ConnectAsync(info.Id, store, new FakeCredentialStore())); + + Assert.Equal("Replacement failed", error.Message); + Assert.Same(activeConnection, manager.GetConnection(info.Id)); + await manager.DisconnectAsync(info.Id); + } + + [Fact] + public async Task CreateQueryConnectionAsync_ReusesActiveProfileAndCredentialContext() + { + var store = new ConnectionStore(_tempDir); + var info = await store.SaveAsync(new ConnectionInfo + { + Id = "query-session", + ServerName = "server.example.com", + CredentialRef = "ssmsx/query-session" + }); + var credentialStore = new FakeCredentialStore(); + var primaryConnection = new SqlConnection(); + var queryConnection = new SqlConnection(); + var calls = new List<(ConnectionInfo Info, ICredentialStore CredentialStore)>(); + var factory = new FakeSqlConnectionFactory((receivedInfo, receivedCredentialStore, _) => + { + calls.Add((receivedInfo, receivedCredentialStore)); + return Task.FromResult(calls.Count == 1 ? primaryConnection : queryConnection); + }); + var manager = new ConnectionManager(factory); + + await manager.ConnectAsync(info.Id, store, credentialStore); + var created = await manager.CreateQueryConnectionAsync(info.Id); + + Assert.Same(queryConnection, created); + Assert.Equal(2, calls.Count); + Assert.Equal(info.Id, calls[1].Info.Id); + Assert.Same(credentialStore, calls[1].CredentialStore); + + await created.DisposeAsync(); + await manager.DisconnectAsync(info.Id); + } + + [Fact] + public async Task CreateQueryConnectionAsync_RejectsConnectionReplacedDuringOpen() + { + var store = new ConnectionStore(_tempDir); + var info = await store.SaveAsync(new ConnectionInfo + { + Id = "replaced-during-open", + ServerName = "server.example.com" + }); + var credentialStore = new FakeCredentialStore(); + var primaryConnection = new SqlConnection(); + var pendingQueryConnection = new SqlConnection(); + var queryConnectionDisposed = false; + pendingQueryConnection.Disposed += (_, _) => queryConnectionDisposed = true; + var replacementConnection = new SqlConnection(); + var queryOpenStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseQueryOpen = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var callCount = 0; + var factory = new FakeSqlConnectionFactory(async (_, _, _) => + { + var call = Interlocked.Increment(ref callCount); + if (call == 1) + return primaryConnection; + if (call == 2) + { + queryOpenStarted.SetResult(); + await releaseQueryOpen.Task; + return pendingQueryConnection; + } + + return replacementConnection; + }); + var manager = new ConnectionManager(factory); + await manager.ConnectAsync(info.Id, store, credentialStore); + + var pendingQuery = manager.CreateQueryConnectionAsync(info.Id); + await queryOpenStarted.Task; + await manager.ConnectAsync(info.Id, store, credentialStore); + releaseQueryOpen.SetResult(); + + var error = await Assert.ThrowsAsync(() => pendingQuery); + Assert.Contains("changed while the query session was opening", error.Message); + Assert.Same(replacementConnection, manager.GetConnection(info.Id)); + Assert.True(queryConnectionDisposed); + + await manager.DisconnectAsync(info.Id); + } + private sealed class FakeSqlConnectionFactory( - Func> create) : SqlConnectionFactory + Func> create) : SqlConnectionFactory { public override Task CreateAsync( ConnectionInfo info, ICredentialStore credentialStore, string? inlinePassword = null, - CancellationToken ct = default) => create(); + CancellationToken ct = default) => create(info, credentialStore, ct); } private sealed class FakeCredentialStore : ICredentialStore diff --git a/sidecar/tests/Ssmsx.Core.Tests/Query/QueryCancellationManagerTests.cs b/sidecar/tests/Ssmsx.Core.Tests/Query/QueryCancellationManagerTests.cs new file mode 100644 index 0000000..cd0530e --- /dev/null +++ b/sidecar/tests/Ssmsx.Core.Tests/Query/QueryCancellationManagerTests.cs @@ -0,0 +1,79 @@ +using Microsoft.Data.SqlClient; +using Ssmsx.Core.Query; +using Xunit; + +namespace Ssmsx.Core.Tests.Query; + +public sealed class QueryCancellationManagerTests +{ + [Fact] + public void Register_RejectsDuplicateQueryIdWithoutReplacingOriginalState() + { + var manager = new QueryCancellationManager(); + using var original = new CancellationTokenSource(); + using var duplicate = new CancellationTokenSource(); + manager.Register("query-1", original); + + Assert.Throws(() => manager.Register("query-1", duplicate)); + Assert.True(manager.Cancel("query-1")); + Assert.True(original.IsCancellationRequested); + Assert.False(duplicate.IsCancellationRequested); + + manager.Remove("query-1"); + } + + [Fact] + public void RemoveAfterDuplicateRegistrationFailure_CleansUpOriginalStateForReuse() + { + var manager = new QueryCancellationManager(); + using var original = new CancellationTokenSource(); + using var duplicate = new CancellationTokenSource(); + using var replacement = new CancellationTokenSource(); + manager.Register("query-1", original); + + Assert.Throws(() => manager.Register("query-1", duplicate)); + + manager.Remove("query-1"); + manager.Register("query-1", replacement); + + Assert.True(manager.Cancel("query-1")); + Assert.True(replacement.IsCancellationRequested); + manager.Remove("query-1"); + } + + [Fact] + public void CancelBeforeSetCommand_CancelsTokenAndAcceptsLateCommand() + { + var manager = new QueryCancellationManager(); + var cancellation = new CancellationTokenSource(); + using var command = new SqlCommand(); + manager.Register("query-1", cancellation); + + Assert.True(manager.Cancel("query-1")); + manager.SetCommand("query-1", command); + + Assert.True(cancellation.IsCancellationRequested); + manager.Remove("query-1"); + } + + [Fact] + public async Task CancelAndSetCommand_AreSafeWhenTheyRace() + { + var manager = new QueryCancellationManager(); + + for (var index = 0; index < 100; index++) + { + var queryId = $"query-{index}"; + var cancellation = new CancellationTokenSource(); + using var command = new SqlCommand(); + manager.Register(queryId, cancellation); + + await Task.WhenAll( + Task.Run(() => manager.Cancel(queryId)), + Task.Run(() => manager.SetCommand(queryId, command))); + + Assert.True(cancellation.IsCancellationRequested); + manager.Remove(queryId); + } + } +} diff --git a/sidecar/tests/Ssmsx.Core.Tests/Query/QueryExecutorTests.cs b/sidecar/tests/Ssmsx.Core.Tests/Query/QueryExecutorTests.cs index 32a16c9..70f6800 100644 --- a/sidecar/tests/Ssmsx.Core.Tests/Query/QueryExecutorTests.cs +++ b/sidecar/tests/Ssmsx.Core.Tests/Query/QueryExecutorTests.cs @@ -1,3 +1,4 @@ +using System.Data; using Ssmsx.Core.Query; using Xunit; @@ -5,6 +6,19 @@ namespace Ssmsx.Core.Tests.Query; public class QueryExecutorTests { + [Theory] + [InlineData(ConnectionState.Open, 16, false)] + [InlineData(ConnectionState.Open, 20, true)] + [InlineData(ConnectionState.Closed, 16, true)] + [InlineData(ConnectionState.Broken, 16, true)] + public void ShouldEvictSession_OnlyEvictsFatalOrUnusableConnections( + ConnectionState state, + byte errorClass, + bool expected) + { + Assert.Equal(expected, QueryExecutor.ShouldEvictSession(state, errorClass)); + } + [Theory] [InlineData(-1, null)] [InlineData(0, "(0 rows affected)")] diff --git a/sidecar/tests/Ssmsx.Core.Tests/Query/QuerySessionManagerTests.cs b/sidecar/tests/Ssmsx.Core.Tests/Query/QuerySessionManagerTests.cs new file mode 100644 index 0000000..41bc5f7 --- /dev/null +++ b/sidecar/tests/Ssmsx.Core.Tests/Query/QuerySessionManagerTests.cs @@ -0,0 +1,158 @@ +using Microsoft.Data.SqlClient; +using Ssmsx.Core.Query; +using Xunit; + +namespace Ssmsx.Core.Tests.Query; + +public sealed class QuerySessionManagerTests +{ + [Fact] + public async Task AcquireAsync_ReusesConnectionWithinSession() + { + var openedConnections = new List(); + await using var manager = CreateManager(openedConnections); + + SqlConnection firstConnection; + await using (var first = await manager.AcquireAsync("tab-1", "connection-1")) + firstConnection = first.Connection; + + await using var second = await manager.AcquireAsync("tab-1", "connection-1"); + + Assert.Same(firstConnection, second.Connection); + Assert.Single(openedConnections); + } + + [Fact] + public async Task AcquireAsync_UsesIndependentConnectionsForDifferentSessions() + { + var openedConnections = new List(); + await using var manager = CreateManager(openedConnections); + + await using var first = await manager.AcquireAsync("tab-1", "connection-1"); + await using var second = await manager.AcquireAsync("tab-2", "connection-1"); + + Assert.NotSame(first.Connection, second.Connection); + Assert.Equal(2, openedConnections.Count); + } + + [Fact] + public async Task AcquireAsync_RejectsOverlappingExecutionWithinSession() + { + var openedConnections = new List(); + await using var manager = CreateManager(openedConnections); + await using var first = await manager.AcquireAsync("tab-1", "connection-1"); + + var error = await Assert.ThrowsAsync(() => + manager.AcquireAsync("tab-1", "connection-1")); + + Assert.Contains("already executing", error.Message); + Assert.Single(openedConnections); + } + + [Fact] + public async Task AcquireAsync_ReplacesRetargetedSessionAndIgnoresStaleClose() + { + var openedConnections = new List(); + await using var manager = CreateManager(openedConnections); + + SqlConnection oldConnection; + await using (var first = await manager.AcquireAsync("tab-1", "connection-old")) + oldConnection = first.Connection; + + var oldDisposed = false; + oldConnection.Disposed += (_, _) => oldDisposed = true; + + await using (var replacement = await manager.AcquireAsync("tab-1", "connection-new")) + { + Assert.NotSame(oldConnection, replacement.Connection); + } + + Assert.True(oldDisposed); + Assert.False(await manager.CloseSessionAsync("tab-1", "connection-old")); + + await using var current = await manager.AcquireAsync("tab-1", "connection-new"); + Assert.Same(openedConnections[1], current.Connection); + } + + [Fact] + public async Task CloseSessionAsync_CancelsExecutionAndDisposesConnection() + { + var openedConnections = new List(); + await using var manager = CreateManager(openedConnections); + var lease = await manager.AcquireAsync("tab-1", "connection-1"); + var connectionDisposed = false; + lease.Connection.Disposed += (_, _) => connectionDisposed = true; + + Assert.True(await manager.CloseSessionAsync("tab-1", "connection-1")); + + Assert.True(lease.CancellationToken.IsCancellationRequested); + Assert.True(connectionDisposed); + await lease.DisposeAsync(); + } + + [Fact] + public async Task CloseConnectionSessionsAsync_ClosesOnlyMatchingSessions() + { + var openedConnections = new List(); + await using var manager = CreateManager(openedConnections); + + await using (var first = await manager.AcquireAsync("tab-1", "connection-1")) { } + await using (var second = await manager.AcquireAsync("tab-2", "connection-1")) { } + await using (var other = await manager.AcquireAsync("tab-3", "connection-2")) { } + + var matchingDisposed = new[] { false, false }; + openedConnections[0].Disposed += (_, _) => matchingDisposed[0] = true; + openedConnections[1].Disposed += (_, _) => matchingDisposed[1] = true; + + Assert.Equal(2, await manager.CloseConnectionSessionsAsync("connection-1")); + Assert.All(matchingDisposed, Assert.True); + + await using var retained = await manager.AcquireAsync("tab-3", "connection-2"); + Assert.Same(openedConnections[2], retained.Connection); + } + + [Fact] + public async Task BrokenSession_IsUnpublishedBeforeAConcurrentNextExecution() + { + var openedConnections = new List(); + await using var manager = CreateManager(openedConnections); + + var first = await manager.AcquireAsync("tab-1", "connection-1"); + var firstConnection = first.Connection; + first.MarkBroken(); + var release = first.DisposeAsync().AsTask(); + + await using var second = await manager.AcquireAsync("tab-1", "connection-1"); + await release; + + Assert.NotSame(firstConnection, second.Connection); + Assert.Equal(2, openedConnections.Count); + } + + [Fact] + public async Task FailedConnectionOpen_IsEvictedBeforeRetry() + { + var attempts = 0; + await using var manager = new QuerySessionManager((_, _) => + { + attempts++; + return attempts == 1 + ? Task.FromException(new InvalidOperationException("open failed")) + : Task.FromResult(new SqlConnection()); + }); + + await Assert.ThrowsAsync(() => + manager.AcquireAsync("tab-1", "connection-1")); + + await using var retry = await manager.AcquireAsync("tab-1", "connection-1"); + Assert.Equal(2, attempts); + } + + private static QuerySessionManager CreateManager(List openedConnections) => + new((_, _) => + { + var connection = new SqlConnection(); + openedConnections.Add(connection); + return Task.FromResult(connection); + }); +} diff --git a/sidecar/tests/Ssmsx.Protocol.Tests/JsonRpcTests.cs b/sidecar/tests/Ssmsx.Protocol.Tests/JsonRpcTests.cs index 47b1aa3..5c36f13 100644 --- a/sidecar/tests/Ssmsx.Protocol.Tests/JsonRpcTests.cs +++ b/sidecar/tests/Ssmsx.Protocol.Tests/JsonRpcTests.cs @@ -72,4 +72,60 @@ public void ReassignColorProfileMessages_RoundTripExpectedContract() Assert.Equal("{\"fromProfileId\":\"custom\",\"toProfileId\":\"red\"}", parametersJson); Assert.Equal("{\"updatedCount\":3}", resultJson); } + + [Fact] + public void QueryExecuteParams_RoundTripExpectedContract() + { + var parameters = new QueryExecuteParams + { + SessionId = "query-tab-1", + ConnectionId = "connection-1", + Database = "AdventureWorks2022", + Sql = "SELECT 1" + }; + + var json = JsonSerializer.Serialize( + parameters, + ProtocolJsonContext.Default.QueryExecuteParams); + var roundTripped = JsonSerializer.Deserialize( + json, + ProtocolJsonContext.Default.QueryExecuteParams); + + Assert.Equal( + "{\"sessionId\":\"query-tab-1\",\"connectionId\":\"connection-1\",\"database\":\"AdventureWorks2022\",\"sql\":\"SELECT 1\"}", + json); + Assert.NotNull(roundTripped); + Assert.Equal(parameters, roundTripped); + } + + [Fact] + public void QueryExecuteParams_WithoutSessionId_IsRejected() + { + var json = """{"connectionId":"connection-1","database":"master","sql":"SELECT 1"}"""; + + Assert.Throws(() => JsonSerializer.Deserialize( + json, + ProtocolJsonContext.Default.QueryExecuteParams)); + } + + [Fact] + public void QuerySessionCloseMessages_RoundTripExpectedContract() + { + var parameters = new QuerySessionCloseParams + { + SessionId = "query-tab-1", + ConnectionId = "connection-1" + }; + var result = new QuerySessionCloseResult { Closed = true }; + + var parametersJson = JsonSerializer.Serialize( + parameters, + ProtocolJsonContext.Default.QuerySessionCloseParams); + var resultJson = JsonSerializer.Serialize( + result, + ProtocolJsonContext.Default.QuerySessionCloseResult); + + Assert.Equal("{\"sessionId\":\"query-tab-1\",\"connectionId\":\"connection-1\"}", parametersJson); + Assert.Equal("{\"closed\":true}", resultJson); + } } diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index 053123f..c321444 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -1,31 +1,58 @@ use crate::sidecar::SidecarManager; use tauri::Emitter; +const QUERY_EXECUTE_METHOD: &str = "query.execute"; +const QUERY_SESSION_CLOSE_METHOD: &str = "query.sessionClose"; + +fn query_execute_params( + session_id: String, + connection_id: String, + database: String, + sql: String, +) -> serde_json::Value { + serde_json::json!({ + "sessionId": session_id, + "connectionId": connection_id, + "database": database, + "sql": sql + }) +} + +fn query_session_close_params(session_id: String, connection_id: String) -> serde_json::Value { + serde_json::json!({ + "sessionId": session_id, + "connectionId": connection_id + }) +} + #[tauri::command] pub async fn query_execute( sidecar: tauri::State<'_, SidecarManager>, app_handle: tauri::AppHandle, request_id: String, + session_id: String, connection_id: String, database: String, sql: String, ) -> Result { - let params = serde_json::json!({ - "connectionId": connection_id, - "database": database, - "sql": sql - }); + let params = query_execute_params( + session_id.clone(), + connection_id.clone(), + database.clone(), + sql, + ); // Use the caller-provided request_id so the frontend can store it in // state BEFORE invoke returns — otherwise streaming events can arrive // before the frontend knows which tab they belong to. let (request_id, mut rx) = sidecar - .send_streaming_request_with_id("query.execute", Some(params), request_id) + .send_streaming_request_with_id(QUERY_EXECUTE_METHOD, Some(params), request_id) .await?; log::debug!( - "Query execution started: request_id='{}', connection_id='{}', database='{}'", + "Query execution started: request_id='{}', session_id='{}', connection_id='{}', database='{}'", request_id, + session_id, connection_id, database ); @@ -83,6 +110,29 @@ pub async fn query_execute( .map_err(|e| format!("Failed to serialize query response: {}", e)) } +#[tauri::command] +pub async fn query_session_close( + sidecar: tauri::State<'_, SidecarManager>, + session_id: String, + connection_id: String, +) -> Result { + log::debug!( + "Closing query session: session_id='{}', connection_id='{}'", + session_id, + connection_id + ); + let params = query_session_close_params(session_id.clone(), connection_id); + let result = sidecar + .send_request(QUERY_SESSION_CLOSE_METHOD, Some(params)) + .await?; + serde_json::to_string(&result).map_err(|e| { + format!( + "Failed to serialize close response for query session '{}': {}", + session_id, e + ) + }) +} + #[tauri::command] pub async fn intellisense_get_metadata( sidecar: tauri::State<'_, SidecarManager>, @@ -121,3 +171,42 @@ pub async fn query_cancel( ) }) } + +#[cfg(test)] +mod tests { + use super::{ + query_execute_params, query_session_close_params, QUERY_EXECUTE_METHOD, + QUERY_SESSION_CLOSE_METHOD, + }; + + #[test] + fn query_execute_bridge_includes_the_window_session() { + assert_eq!(QUERY_EXECUTE_METHOD, "query.execute"); + assert_eq!( + query_execute_params( + "tab-7".into(), + "connection-2".into(), + "warehouse".into(), + "select 1".into(), + ), + serde_json::json!({ + "sessionId": "tab-7", + "connectionId": "connection-2", + "database": "warehouse", + "sql": "select 1" + }) + ); + } + + #[test] + fn query_session_close_bridge_uses_the_window_session() { + assert_eq!(QUERY_SESSION_CLOSE_METHOD, "query.sessionClose"); + assert_eq!( + query_session_close_params("tab-7".into(), "connection-2".into()), + serde_json::json!({ + "sessionId": "tab-7", + "connectionId": "connection-2" + }) + ); + } +} diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 37de096..cdfb99a 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -95,6 +95,7 @@ pub fn run() { commands::explorer::explorer_database_diagram, commands::query::query_execute, commands::query::query_cancel, + commands::query::query_session_close, commands::query::intellisense_get_metadata, ]) .run(tauri::generate_context!()) diff --git a/src/features/query/api/queryApi.ts b/src/features/query/api/queryApi.ts index a7c7a58..968ec42 100644 --- a/src/features/query/api/queryApi.ts +++ b/src/features/query/api/queryApi.ts @@ -8,12 +8,14 @@ interface QueryExecuteResponse { export async function queryExecute( requestId: string, + sessionId: string, connectionId: string, database: string, sql: string ): Promise { const result = await invoke("query_execute", { requestId, + sessionId, connectionId, database, sql, @@ -25,6 +27,13 @@ export async function queryCancel(queryId: string): Promise { await invoke("query_cancel", { queryId }); } +export async function querySessionClose( + sessionId: string, + connectionId: string +): Promise { + await invoke("query_session_close", { sessionId, connectionId }); +} + export async function onQueryResults( handler: (payload: QueryBatchPayload) => void ): Promise { diff --git a/src/features/query/store/queryStore.ts b/src/features/query/store/queryStore.ts index 413589e..1f5f9b9 100644 --- a/src/features/query/store/queryStore.ts +++ b/src/features/query/store/queryStore.ts @@ -11,6 +11,7 @@ import type { import { queryExecute, queryCancel, + querySessionClose, intellisenseGetMetadata, type IntelliSenseMetadata, } from "../api/queryApi"; @@ -137,6 +138,12 @@ function canUseLocalStorage(): boolean { return typeof window !== "undefined" && typeof window.localStorage !== "undefined"; } +function closeQuerySession(sessionId: string, connectionId: string): void { + void querySessionClose(sessionId, connectionId).catch((error) => { + console.error(`Failed to close query session for tab '${sessionId}':`, error); + }); +} + function readSavedSession(): SavedQuerySession | null { if (!canUseLocalStorage()) return null; @@ -221,7 +228,12 @@ export const useQueryStore = create((set, get) => ({ })); }, - removeTab: (id) => + removeTab: (id) => { + const tab = get().tabs.find((candidate) => candidate.id === id); + if (tab && tab.kind !== "diagram" && tab.connectionId) { + closeQuerySession(id, tab.connectionId); + } + set((state) => { const tabs = state.tabs.filter((t) => t.id !== id); const { [id]: _sql, ...restSql } = state.tabSql; @@ -241,7 +253,8 @@ export const useQueryStore = create((set, get) => ({ : null : state.activeTabId, }; - }), + }); + }, consumePendingCursorOffset: (tabId) => set((state) => { @@ -251,22 +264,43 @@ export const useQueryStore = create((set, get) => ({ setActiveTab: (id) => set({ activeTabId: id }), - updateTab: (tabId, patch) => + updateTab: (tabId, patch) => { + const tab = get().tabs.find((candidate) => candidate.id === tabId); + if ( + tab && + tab.kind !== "diagram" && + tab.connectionId && + patch.connectionId !== undefined && + patch.connectionId !== tab.connectionId + ) { + closeQuerySession(tabId, tab.connectionId); + } + set((state) => { + const currentTab = state.tabs.find((candidate) => candidate.id === tabId); + const connectionChanged = + patch.connectionId !== undefined && + patch.connectionId !== currentTab?.connectionId; const execution = state.executionInfo[tabId]; + const { [tabId]: _execution, ...executionWithoutRetargetedTab } = state.executionInfo; + const { [tabId]: _results, ...resultsWithoutRetargetedTab } = state.results; return { tabs: state.tabs.map((tab) => tab.id === tabId ? { ...tab, ...patch } : tab ), executionInfo: - patch.database !== undefined && execution?.state === "executing" + connectionChanged + ? executionWithoutRetargetedTab + : patch.database !== undefined && execution?.state === "executing" ? { ...state.executionInfo, [tabId]: { ...execution, databaseSyncEnabled: false }, } : state.executionInfo, + results: connectionChanged ? resultsWithoutRetargetedTab : state.results, }; - }), + }); + }, updateSql: (tabId, sql) => set((state) => ({ @@ -288,7 +322,14 @@ export const useQueryStore = create((set, get) => ({ ), })), - closeOtherTabs: (tabId) => + closeOtherTabs: (tabId) => { + const removedTabs = get().tabs.filter((tab) => tab.id !== tabId); + for (const tab of removedTabs) { + if (tab.kind !== "diagram" && tab.connectionId) { + closeQuerySession(tab.id, tab.connectionId); + } + } + set((state) => { const kept = state.tabs.filter((t) => t.id === tabId); const removedIds = state.tabs @@ -314,9 +355,16 @@ export const useQueryStore = create((set, get) => ({ results, pendingCursorOffsets, }; - }), + }); + }, + + closeAllTabs: () => { + for (const tab of get().tabs) { + if (tab.kind !== "diagram" && tab.connectionId) { + closeQuerySession(tab.id, tab.connectionId); + } + } - closeAllTabs: () => set({ tabs: [], activeTabId: null, @@ -324,7 +372,8 @@ export const useQueryStore = create((set, get) => ({ executionInfo: {}, results: {}, pendingCursorOffsets: {}, - }), + }); + }, saveSession: () => { if (!canUseLocalStorage()) return; @@ -519,7 +568,7 @@ export const useQueryStore = create((set, get) => ({ })); try { - await queryExecute(requestId, tab.connectionId, tab.database, queryText); + await queryExecute(requestId, tabId, tab.connectionId, tab.database, queryText); // requestId is already in executionInfo — no follow-up set needed } catch (e) { console.error(`Query execution failed for tab '${tabId}':`, e); diff --git a/tests/frontend/newQueryTemplateBrowser.test.ts b/tests/frontend/newQueryTemplateBrowser.test.ts index c560e40..1148048 100644 --- a/tests/frontend/newQueryTemplateBrowser.test.ts +++ b/tests/frontend/newQueryTemplateBrowser.test.ts @@ -1,5 +1,5 @@ import assert from "node:assert/strict"; -import { existsSync } from "node:fs"; +import { existsSync, mkdtempSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import test from "node:test"; @@ -13,6 +13,7 @@ const PROCESSED_TEMPLATE = " BEGIN TRANSACTION;\n SELECT 1;\nCOMMIT; \n"; const CURSOR_SENTINEL = "/* here */"; const SCREENSHOT_PATH = join(tmpdir(), "ssmsx-new-query-template-browser.png"); const MESSAGES_SCREENSHOT_PATH = join(tmpdir(), "ssmsx-query-messages-select-all.png"); +const VITE_CACHE_DIR = mkdtempSync(join(tmpdir(), "ssmsx-new-query-template-vite-")); interface FixtureWindow extends Window { ssmsxNewQueryTemplateFixture?: { @@ -30,6 +31,7 @@ test("query acceptance preserves template whitespace and scopes message selectio root: process.cwd(), logLevel: "error", plugins: [react(), tailwindcss()], + cacheDir: VITE_CACHE_DIR, server: { host: "127.0.0.1", port: 0 }, }); await vite.listen(); @@ -143,5 +145,6 @@ test("query acceptance preserves template whitespace and scopes message selectio } finally { await browser?.close(); await vite.close(); + rmSync(VITE_CACHE_DIR, { recursive: true, force: true }); } }); diff --git a/tests/frontend/objectExplorerStyles.test.ts b/tests/frontend/objectExplorerStyles.test.ts index 9d5d453..1838c67 100644 --- a/tests/frontend/objectExplorerStyles.test.ts +++ b/tests/frontend/objectExplorerStyles.test.ts @@ -1,10 +1,15 @@ import assert from "node:assert/strict"; +import { mkdtempSync, rmSync } from "node:fs"; +import { join } from "node:path"; +import { tmpdir } from "node:os"; import test from "node:test"; import react from "@vitejs/plugin-react"; import tailwindcss from "@tailwindcss/vite"; import { chromium, type Browser } from "@playwright/test"; import { createServer } from "vite"; +const VITE_CACHE_DIR = mkdtempSync(join(tmpdir(), "ssmsx-object-explorer-vite-")); + function relativeLuminance(rgb: number[]): number { const channels = rgb.map((value) => { const channel = value / 255; @@ -32,6 +37,7 @@ test("Object Explorer applies profile styles per subtree without tinting its she root: process.cwd(), logLevel: "error", plugins: [react(), tailwindcss()], + cacheDir: VITE_CACHE_DIR, server: { host: "127.0.0.1", port: 0 }, }); await vite.listen(); @@ -140,5 +146,6 @@ test("Object Explorer applies profile styles per subtree without tinting its she } finally { await browser?.close(); await vite.close(); + rmSync(VITE_CACHE_DIR, { recursive: true, force: true }); } }); diff --git a/tests/frontend/queryUi.test.tsx b/tests/frontend/queryUi.test.tsx index fd4837e..72a1d09 100644 --- a/tests/frontend/queryUi.test.tsx +++ b/tests/frontend/queryUi.test.tsx @@ -1,6 +1,7 @@ // @vitest-environment jsdom import { cleanup, fireEvent, render, screen } from "@testing-library/react"; +import { invoke } from "@tauri-apps/api/core"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { QueryResultsTable } from "../../src/features/query/components/QueryResultsTable"; import { QueryTabBar } from "../../src/features/query/components/QueryTabBar"; @@ -18,6 +19,10 @@ vi.mock("@monaco-editor/react", () => ({ default: () =>
, })); +vi.mock("@tauri-apps/api/core", () => ({ + invoke: vi.fn().mockResolvedValue("{}"), +})); + vi.mock("@tanstack/react-virtual", () => ({ useVirtualizer: ({ count, @@ -63,6 +68,7 @@ const tabs: QueryTab[] = [ ]; function resetStores() { + vi.mocked(invoke).mockReset().mockResolvedValue("{}"); window.localStorage.clear(); useQueryStore.setState({ tabs: [], @@ -241,6 +247,119 @@ describe("query tab session and profile colours", () => { expect(useQueryStore.getState().pendingCursorOffsets).toEqual({}); }); + it("passes the query window id to the sidecar as its session id", async () => { + useQueryStore.setState({ + tabs: [tabs[0]], + activeTabId: "unpinned", + tabSql: { unpinned: "select 1" }, + }); + useConnectionStore.setState({ activeConnectionIds: ["prod"] }); + + await useQueryStore.getState().executeQuery("unpinned"); + + expect(invoke).toHaveBeenCalledWith("query_execute", { + requestId: expect.any(String), + sessionId: "unpinned", + connectionId: "prod", + database: "master", + sql: "select 1", + }); + }); + + it("does not queue a second execution from the same query window", async () => { + useQueryStore.setState({ + tabs: [tabs[0]], + activeTabId: "unpinned", + tabSql: { unpinned: "select 1" }, + }); + useConnectionStore.setState({ activeConnectionIds: ["prod"] }); + + await useQueryStore.getState().executeQuery("unpinned"); + await useQueryStore.getState().executeQuery("unpinned"); + + expect(invoke).toHaveBeenCalledTimes(1); + }); + + it("closes query sessions when tabs close or change connections", () => { + useQueryStore.setState({ tabs, activeTabId: "unpinned" }); + + useQueryStore.getState().removeTab("unpinned"); + expect(invoke).toHaveBeenLastCalledWith("query_session_close", { + sessionId: "unpinned", + connectionId: "prod", + }); + + useQueryStore.setState({ tabs, activeTabId: "unpinned" }); + useQueryStore.getState().closeOtherTabs("unpinned"); + expect(invoke).toHaveBeenLastCalledWith("query_session_close", { + sessionId: "pinned", + connectionId: "prod", + }); + + useQueryStore.setState({ tabs, activeTabId: "unpinned" }); + useQueryStore.getState().closeAllTabs(); + expect(invoke).toHaveBeenNthCalledWith(3, "query_session_close", { + sessionId: "unpinned", + connectionId: "prod", + }); + expect(invoke).toHaveBeenNthCalledWith(4, "query_session_close", { + sessionId: "pinned", + connectionId: "prod", + }); + + useQueryStore.setState({ tabs: [tabs[0]], activeTabId: "unpinned" }); + useQueryStore.getState().updateTab("unpinned", { database: "app" }); + expect(invoke).toHaveBeenCalledTimes(4); + + useQueryStore.getState().updateTab("unpinned", { + connectionId: "reporting", + }); + expect(invoke).toHaveBeenLastCalledWith("query_session_close", { + sessionId: "unpinned", + connectionId: "prod", + }); + }); + + it("retargeting an executing tab drops old execution state and ignores late results", () => { + useQueryStore.setState({ + tabs: [tabs[0]], + activeTabId: "unpinned", + executionInfo: { + unpinned: { + state: "executing", + queryId: "old-query", + requestId: "old-request", + startTime: Date.now(), + }, + }, + results: { + unpinned: { + columns: ["old"], + rows: [["old result"]], + resultSets: [], + messages: [], + executionTimeMs: null, + totalRows: 1, + }, + }, + }); + + useQueryStore.getState().updateTab("unpinned", { connectionId: "reporting" }); + expect(useQueryStore.getState().executionInfo.unpinned).toBeUndefined(); + expect(useQueryStore.getState().results.unpinned).toBeUndefined(); + + useQueryStore.getState().handleResultsBatch({ + queryId: "old-query", + requestId: "old-request", + batch: 1, + done: true, + columns: ["late"], + rows: [["late result"]], + }); + + expect(useQueryStore.getState().results.unpinned).toBeUndefined(); + }); + it("renders wrapping tab bands with both profile colours on active and inactive tabs", () => { useConnectionStore.setState({ connections: [