Edit Data Service Updates for Edit Data in the MSSQL for VS Code extension (#2530)

lewis-sanchez · web-flow · commit eb45cd11f9a5 · 2025-10-28T12:43:08.000-07:00
* Adds script generation for edit sessions

* Adds tests for ScriptEdits method

* Adds column names to EditSubsetResult

* Returns reverted row after revert operation

* Verifies reverted row contains correct data

* Add flag to indicate if a column is editable

* Populate newly created row in result

* Handle pending row create during delete.

* Minor clean up

* Add row ID validation to edit session

* Fix test assertion for async operation

* Clarifies request result documentation

* Refactor EditSession row ID validation

* Clarifies rename of column property

* Clarify request handlers by using named arguments

* Refactors EditDataService for error handling

* Use fluent assertions

* Use named argument for the rowCount

* Use HandleSessionRequest to get edit scripts

* Clean up comments

* Fix broken tests from HandleSessionRequest refactor

* Add additional logging
diff --git a/src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditColumnInfo.cs b/src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditColumnInfo.cs
@@ -0,0 +1,24 @@
+﻿//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+namespace Microsoft.SqlTools.ServiceLayer.EditData.Contracts
+{
+    /// <summary>
+    /// Metadata about a column in an editable result set
+    /// </summary>
+    public class EditColumnInfo
+    {
+        /// <summary>
+        /// The name of the column
+        /// </summary>
+        public string Name { get; set; }
+
+        /// <summary>
+        /// Whether or not the column can be edited. Columns may not be editable for several reasons,
+        /// such as being computed columns, identity columns, or columns that are part of a primary key.
+        /// </summary>
+        public bool IsEditable { get; set; }
+    }
+}
diff --git a/src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditCreateRowRequest.cs b/src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditCreateRowRequest.cs
@@ -31,6 +31,11 @@ public class EditCreateRowResult
         /// The internal ID of the newly created row
         /// </summary>
         public long NewRowId { get; set; }
+
+        /// <summary>
+        /// The newly created row with all cell data and state information
+        /// </summary>
+        public EditRow Row { get; set; }
     }
 
     public class EditCreateRowRequest
diff --git a/src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditRevertRowRequest.cs b/src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditRevertRowRequest.cs
@@ -21,6 +21,10 @@ public class EditRevertRowParams : RowOperationParams
     /// </summary>
     public class EditRevertRowResult
     {
+        /// <summary>
+        /// The row after the revert was applied, representing the original, unedited state
+        /// </summary>
+        public EditRow Row { get; set; }
     }
 
     public class EditRevertRowRequest
diff --git a/src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditScriptRequest.cs b/src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditScriptRequest.cs
@@ -0,0 +1,32 @@
+﻿//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.EditData.Contracts
+{
+    /// <summary>
+    /// Parameters for a request to get a script with all DML statements for an edit session.
+    /// </summary>
+    public class EditScriptParams : SessionOperationParams
+    {
+
+    }
+
+    /// <summary>
+    /// Returns the scripts with all DML statements for an edit session.
+    /// </summary>
+    public class EditScriptResult
+    {
+        public string[] Scripts { get; set; }
+    }
+
+    public class EditScriptRequest
+    {
+        public static readonly
+            RequestType<EditScriptParams, EditScriptResult> Type =
+            RequestType<EditScriptParams, EditScriptResult>.Create("edit/script");
+    }
+}
diff --git a/src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditSubsetRequest.cs b/src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditSubsetRequest.cs
@@ -43,6 +43,11 @@ public class EditSubsetResult
         /// The requested subset of rows, with information about whether or not the rows are dirty
         /// </summary>
         public EditRow[] Subset { get; set; }
+
+        /// <summary>
+        /// The column information including name and editability for the result set
+        /// </summary>
+        public EditColumnInfo[] ColumnInfo { get; set; }
     }
 
     public class EditSubsetRequest
diff --git a/src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs b/src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs
@@ -77,33 +77,66 @@ internal EditDataService(QueryExecutionService qes, ConnectionService cs, IEditM
         public void InitializeService(ServiceHost serviceHost)
         {
             // Register handlers for requests
-            serviceHost.SetRequestHandler(EditCreateRowRequest.Type, HandleCreateRowRequest, true);
-            serviceHost.SetRequestHandler(EditDeleteRowRequest.Type, HandleDeleteRowRequest, true);
-            serviceHost.SetRequestHandler(EditDisposeRequest.Type, HandleDisposeRequest, true);
-            serviceHost.SetRequestHandler(EditInitializeRequest.Type, HandleInitializeRequest, true);
-            serviceHost.SetRequestHandler(EditRevertCellRequest.Type, HandleRevertCellRequest, true);
-            serviceHost.SetRequestHandler(EditRevertRowRequest.Type, HandleRevertRowRequest, true);
-            serviceHost.SetRequestHandler(EditSubsetRequest.Type, HandleSubsetRequest, true);
-            serviceHost.SetRequestHandler(EditUpdateCellRequest.Type, HandleUpdateCellRequest, true);
-            serviceHost.SetRequestHandler(EditCommitRequest.Type, HandleCommitRequest, true);
+            serviceHost.SetRequestHandler(EditCreateRowRequest.Type, HandleCreateRowRequest, isParallelProcessingSupported: true);
+            serviceHost.SetRequestHandler(EditDeleteRowRequest.Type, HandleDeleteRowRequest, isParallelProcessingSupported: true);
+            serviceHost.SetRequestHandler(EditDisposeRequest.Type, HandleDisposeRequest, isParallelProcessingSupported: true);
+            serviceHost.SetRequestHandler(EditInitializeRequest.Type, HandleInitializeRequest, isParallelProcessingSupported: true);
+            serviceHost.SetRequestHandler(EditRevertCellRequest.Type, HandleRevertCellRequest, isParallelProcessingSupported: true);
+            serviceHost.SetRequestHandler(EditRevertRowRequest.Type, HandleRevertRowRequest, isParallelProcessingSupported: true);
+            serviceHost.SetRequestHandler(EditSubsetRequest.Type, HandleSubsetRequest, isParallelProcessingSupported: true);
+            serviceHost.SetRequestHandler(EditUpdateCellRequest.Type, HandleUpdateCellRequest, isParallelProcessingSupported: true);
+            serviceHost.SetRequestHandler(EditCommitRequest.Type, HandleCommitRequest, isParallelProcessingSupported: true);
+            serviceHost.SetRequestHandler(EditScriptRequest.Type, HandleEditScriptRequest, isParallelProcessingSupported: true);
         }
 
         #region Request Handlers
 
         internal async Task HandleSessionRequest<TResult>(SessionOperationParams sessionParams,
             RequestContext<TResult> requestContext, Func<EditSession, TResult> sessionOperation)
         {
-            EditSession editSession = GetActiveSessionOrThrow(sessionParams.OwnerUri);
+            try
+            {
+                EditSession editSession = GetActiveSessionOrThrow(sessionParams.OwnerUri);
+
+                // Get the result from execution of the editSession operation
+                TResult result = sessionOperation(editSession);
+                await requestContext.SendResult(result);
+            }
+            catch (Exception e)
+            {
+                await requestContext.SendError(e.Message);
+            }
+        }
 
-            // Get the result from execution of the editSession operation
-            TResult result = sessionOperation(editSession);
-            await requestContext.SendResult(result);
+        internal async Task HandleSessionRequestAsync<TResult>(SessionOperationParams sessionParams,
+            RequestContext<TResult> requestContext, Func<EditSession, Task<TResult>> sessionOperationAsync)
+        {
+            try
+            {
+                EditSession editSession = GetActiveSessionOrThrow(sessionParams.OwnerUri);
+
+                // Get the result from execution of the async editSession operation
+                TResult result = await sessionOperationAsync(editSession);
+                await requestContext.SendResult(result);
+            }
+            catch (Exception e)
+            {
+                await requestContext.SendError(e.Message);
+            }
         }
 
         internal Task HandleCreateRowRequest(EditCreateRowParams createParams,
             RequestContext<EditCreateRowResult> requestContext)
         {
-            return HandleSessionRequest(createParams, requestContext, s => s.CreateRow());
+            return HandleSessionRequestAsync(createParams, requestContext, async session =>
+            {
+                EditCreateRowResult result = session.CreateRow();
+
+                EditRow[] rows = await session.GetRows(result.NewRowId, rowCount: 1);
+                result.Row = rows.Length > 0 ? rows[0] : null;
+
+                return result;
+            });
         }
 
         internal Task HandleDeleteRowRequest(EditDeleteRowParams deleteParams,
@@ -156,7 +189,7 @@ internal async Task HandleInitializeRequest(EditInitializeParams initParams,
                 throw new InvalidOperationException(SR.EditDataSessionAlreadyExists);
             }
 
-            context.ResultSetHandler = (ResultSetEventParams resultSetEventParams) => { session.UpdateColumnInformationWithMetadata(resultSetEventParams.ResultSetSummary.ColumnInfo); };
+            context.ResultSetHandler = resultSetEventParams => { session.UpdateColumnInformationWithMetadata(resultSetEventParams.ResultSetSummary.ColumnInfo); };
 
             // Initialize the session
             session.Initialize(initParams, connector, queryRunner, executionSuccessHandler, executionFailureHandler);
@@ -175,26 +208,31 @@ internal Task HandleRevertCellRequest(EditRevertCellParams revertParams,
         internal Task HandleRevertRowRequest(EditRevertRowParams revertParams,
             RequestContext<EditRevertRowResult> requestContext)
         {
-            return HandleSessionRequest(revertParams, requestContext, session =>
+            return HandleSessionRequestAsync(revertParams, requestContext, async session =>
             {
-                session.RevertRow(revertParams.RowId);
-                return new EditRevertRowResult();
+                EditRow revertedRow = await session.RevertRow(revertParams.RowId);
+
+                return new EditRevertRowResult
+                {
+                    Row = revertedRow
+                };
             });
         }
 
-        internal async Task HandleSubsetRequest(EditSubsetParams subsetParams,
+        internal Task HandleSubsetRequest(EditSubsetParams subsetParams,
             RequestContext<EditSubsetResult> requestContext)
         {
-            EditSession session = GetActiveSessionOrThrow(subsetParams.OwnerUri);
-
-            EditRow[] rows = await session.GetRows(subsetParams.RowStartIndex, subsetParams.RowCount);
-            EditSubsetResult result = new EditSubsetResult
+            return HandleSessionRequestAsync(subsetParams, requestContext, async session =>
             {
-                RowCount = rows.Length,
-                Subset = rows
-            };
-
-            await requestContext.SendResult(result);
+                EditRow[] rows = await session.GetRows(subsetParams.RowStartIndex, subsetParams.RowCount);
+                
+                return new EditSubsetResult
+                {
+                    RowCount = rows.Length,
+                    Subset = rows,
+                    ColumnInfo = session.GetColumnInfo()
+                };
+            });
         }
 
         internal Task HandleUpdateCellRequest(EditUpdateCellParams updateParams,
@@ -229,6 +267,15 @@ internal async Task HandleCommitRequest(EditCommitParams commitParams,
             }
         }
 
+        internal Task HandleEditScriptRequest(EditScriptParams scriptParams, RequestContext<EditScriptResult> requestContext)
+        {
+            return HandleSessionRequest(scriptParams, requestContext, session =>
+            {
+                string[] scripts = session.ScriptEdits();
+                return new EditScriptResult { Scripts = scripts };
+            });
+        }
+
         #endregion
 
         #region Private Helpers
diff --git a/src/Microsoft.SqlTools.ServiceLayer/EditData/EditSession.cs b/src/Microsoft.SqlTools.ServiceLayer/EditData/EditSession.cs
diff --git a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/EditData/ServiceIntegrationTests.cs b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/EditData/ServiceIntegrationTests.cs
diff --git a/test/Microsoft.SqlTools.ServiceLayer.UnitTests/EditData/SessionTests.cs b/test/Microsoft.SqlTools.ServiceLayer.UnitTests/EditData/SessionTests.cs

Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,10 @@ public class EditRevertRowParams : RowOperationParams`
`21`	`21`	`/// </summary>`
`22`	`22`	`public class EditRevertRowResult`
`23`	`23`	`{`
	`24`	`+ /// <summary>`
	`25`	`+ /// The row after the revert was applied, representing the original, unedited state`
	`26`	`+ /// </summary>`
	`27`	`+ public EditRow Row { get; set; }`
`24`	`28`	`}`
`25`	`29`
`26`	`30`	`public class EditRevertRowRequest`