Implicit Revert Row Cleanup Logic (#297)

This change enhances the way that edit/updateCell and edit/revertCell operations are performed. 

## **THE API BREAKING CHANGES**:
* edit/updateCell now returns an EditCell (a DbCellValue with a dirty flag) and a row dirty flag.
* edit/revertCell now returns an EditCell (a DbCellValue with a dirty flag) and a row dirty flag.

If by setting the value of a cell via edit/updateCell the row no longer has any edits (an "implicit revert"), the entire row's edit will be removed from the cache. Additionally, if by requesting edit/revert all the pending edits for a row are removed, the entire row's edit will be removed from the cache. This will prevent issues where committing will generate an invalid script because it has no pending changes.

* Adding EditCell class
Returning EditCell with EditUpdateCellResult

* Adding code that will remove a row update if the row is clean after a cell update

* Adding code that will return an EditCell and row dirty flag when a cell is reverted.
If the row is reverted by the cell revert, the pending update will be removed

* Comments for edit cell

* Changes as per pull request comments

src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditCell.cs

@@ -0,0 +1,39 @@
+//
+// 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.ServiceLayer.QueryExecution.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.EditData.Contracts
+{
+    /// <summary>
+    /// Cell that wraps info from <see cref="DbCellValue"/> for edit purposes
+    /// </summary>
+    public class EditCell : DbCellValue
+    {
+        /// <summary>
+        /// Default, parameterless constructor to make sure that JSON serializing is happy
+        /// </summary>
+        public EditCell() {}
+
+        /// <summary>
+        /// Constructs a new EditCell based on a DbCellValue
+        /// </summary>
+        /// <param name="dbCellValue">The DbCellValue that will be enhanced</param>
+        /// <param name="isDirty">Whether or not the edit cell is dirty</param>
+        public EditCell(DbCellValue dbCellValue, bool isDirty)
+        {
+            IsDirty = isDirty;
+
+            DisplayValue = dbCellValue.DisplayValue;
+            IsNull = dbCellValue.IsNull;
+            RawObject = dbCellValue.RawObject;
+        }
+
+        /// <summary>
+        /// Whether or not the cell is considered dirty
+        /// </summary>
+        public bool IsDirty { get; set; }
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditCellResult.cs

@@ -0,0 +1,18 @@
+namespace Microsoft.SqlTools.ServiceLayer.EditData.Contracts
+{
+    /// <summary>
+    /// Parameters to return when a cell is updated in isolation
+    /// </summary>
+    public class EditCellResult
+    {
+        /// <summary>
+        /// The cell after the revert was applied
+        /// </summary>
+        public EditCell Cell { get; set; }
+
+        /// <summary>
+        /// Whether or not the row is dirty after the revert has been applied
+        /// </summary>
+        public bool IsRowDirty { get; set; }
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditRevertCellRequest.cs

@@ -7,14 +7,19 @@ using Microsoft.SqlTools.Hosting.Protocol.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.EditData.Contracts
 {
+    /// <summary>
+    /// Parameters for the cell revert request
+    /// </summary>
     public class EditRevertCellParams : RowOperationParams
     {
         public int ColumnId { get; set; }
     }
 
-    public class EditRevertCellResult
+    /// <summary>
+    /// Parameters to return upon successful revert of the cell
+    /// </summary>
+    public class EditRevertCellResult : EditCellResult
     {
-        public string NewValue { get; set; }
     }
 
     public class EditRevertCellRequest

src/Microsoft.SqlTools.ServiceLayer/EditData/Contracts/EditUpdateCellRequest.cs

@@ -26,31 +26,8 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.Contracts
     /// <summary>
     /// Parameters to return upon successful update of the cell
     /// </summary>
-    public class EditUpdateCellResult
+    public class EditUpdateCellResult : EditCellResult
     {
-        /// <summary>
-        /// Whether or not the cell value was modified from the provided string.
-        /// If <c>true</c>, the client should replace the display value of the cell with the value
-        /// in <see cref="NewValue"/>
-        /// </summary>
-        public bool HasCorrections { get; set; }
-
-        /// <summary>
-        /// Whether or not the cell was reverted with the change.
-        /// If <c>true</c>, the client should unmark the cell as having an update and replace the
-        /// display value of the cell with the value in <see cref="NewValue"/>
-        /// </summary>
-        public bool IsRevert { get; set; }
-
-        /// <summary>
-        /// Whether or not the new value of the cell is null
-        /// </summary>
-        public bool IsNull { get; set; }
-
-        /// <summary>
-        /// The new string value of the cell
-        /// </summary>
-        public string NewValue { get; set; }
     }
 
     public class EditUpdateCellRequest

src/Microsoft.SqlTools.ServiceLayer/EditData/EditDataService.cs

@@ -184,14 +184,8 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
         internal Task HandleRevertCellRequest(EditRevertCellParams revertParams,
             RequestContext<EditRevertCellResult> requestContext)
         {
-            return HandleSessionRequest(revertParams, requestContext, session =>
-            {
-                string newValue = session.RevertCell(revertParams.RowId, revertParams.ColumnId);
-                return new EditRevertCellResult
-                {
-                    NewValue = newValue
-                };
-            });
+            return HandleSessionRequest(revertParams, requestContext,
+                session => session.RevertCell(revertParams.RowId, revertParams.ColumnId));
         }
 
         internal Task HandleRevertRowRequest(EditRevertRowParams revertParams,

src/Microsoft.SqlTools.ServiceLayer/EditData/EditSession.cs

@@ -308,7 +308,7 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
         /// <param name="rowId">Internal ID of the row to have its edits reverted</param>
         /// <param name="columnId">Ordinal ID of the column to revert</param>
         /// <returns>String version of the old value for the cell</returns>
-        public string RevertCell(long rowId, int columnId)
+        public EditRevertCellResult RevertCell(long rowId, int columnId)
         {
             ThrowIfNotInitialized();
 
@@ -319,8 +319,12 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
                 throw new ArgumentOutOfRangeException(nameof(rowId), SR.EditDataUpdateNotPending);
             }
 
+            // Update the row
+            EditRevertCellResult revertResult = pendingEdit.RevertCell(columnId);
+            CleanupEditIfRowClean(rowId, revertResult);
+
             // Have the edit base revert the cell
-            return pendingEdit.RevertCell(columnId);
+            return revertResult;
         }
 
         /// <summary>
@@ -410,8 +414,11 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
             // doesn't already exist in the cache
             RowEditBase editRow = EditCache.GetOrAdd(rowId, key => new RowUpdate(rowId, associatedResultSet, objectMetadata));
 
-            // Pass the call to the row update
-            return editRow.SetCell(columnId, newValue);
+            // Update the row
+            EditUpdateCellResult result = editRow.SetCell(columnId, newValue);
+            CleanupEditIfRowClean(rowId, result);
+
+            return result;
         }
 
         #endregion
@@ -521,6 +528,24 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
             }
         }
 
+        /// <summary>
+        /// Removes the edit from the edit cache if the row is no longer dirty
+        /// </summary>
+        /// <param name="rowId">ID of the update to cleanup</param>
+        /// <param name="editCellResult">Result with row dirty flag</param>
+        private void CleanupEditIfRowClean(long rowId, EditCellResult editCellResult)
+        {
+            // If the row is still dirty, don't do anything
+            if (editCellResult.IsRowDirty)
+            {
+                return;
+            }
+
+            // Make an attempt to remove the clean row edit. If this fails, it'll be handled on commit attempt.
+            RowEditBase removedRow;
+            EditCache.TryRemove(rowId, out removedRow);
+        }
+
         #endregion
 
         /// <summary>

src/Microsoft.SqlTools.ServiceLayer/EditData/UpdateManagement/CellUpdate.cs

@@ -6,6 +6,7 @@
 using System;
 using System.Globalization;
 using System.Text.RegularExpressions;
+using Microsoft.SqlTools.ServiceLayer.EditData.Contracts;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.SqlTools.Utility;
 
@@ -101,6 +102,14 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
             }
         }
 
+        /// <summary>
+        /// Generates a new EditCell that represents the contents of the cell update
+        /// </summary>
+        public EditCell AsEditCell
+        {
+            get { return new EditCell(AsDbCellValue, true); }
+        }
+
         /// <summary>
         /// The column that the cell will be placed in
         /// </summary>

src/Microsoft.SqlTools.ServiceLayer/EditData/UpdateManagement/RowCreate.cs

@@ -184,14 +184,16 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
         /// </summary>
         /// <param name="columnId">The ordinal ID of the cell to reset</param>
         /// <returns>The default value for the column, or null if no default is defined</returns>
-        public override string RevertCell(int columnId)
+        public override EditRevertCellResult RevertCell(int columnId)
         {
             // Validate that the column can be reverted
             Validate.IsWithinRange(nameof(columnId), columnId, 0, newCells.Length - 1);
 
             // Remove the cell update from list of set cells
             newCells[columnId] = null;
-            return null;                // @TODO: Return default value when we have support checked in
+            return new EditRevertCellResult {IsRowDirty = true, Cell = null};
+            // @TODO: Return default value when we have support checked in
+            // @TODO: RETURN THE DEFAULT VALUE
         }
 
         /// <summary>
@@ -212,14 +214,11 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
             newCells[columnId] = update;
 
             // Put together a result of the change
-            EditUpdateCellResult eucr = new EditUpdateCellResult
+            return new EditUpdateCellResult
             {
-                HasCorrections = update.ValueAsString != newValue,
-                NewValue = update.ValueAsString != newValue ? update.ValueAsString : null,
-                IsNull = update.Value == DBNull.Value,
-                IsRevert = false            // Editing cells of new rows cannot be reverts
-            }; 
-            return eucr;
+                IsRowDirty = true,                // Row creates will always be dirty
+                Cell = update.AsEditCell
+            };
         }
 
         #endregion

src/Microsoft.SqlTools.ServiceLayer/EditData/UpdateManagement/RowDelete.cs

@@ -105,7 +105,7 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
         /// deletion.
         /// </summary>
         /// <param name="columnId">Ordinal of the column to update</param>
-        public override string RevertCell(int columnId)
+        public override EditRevertCellResult RevertCell(int columnId)
         {
             throw new InvalidOperationException(SR.EditDataDeleteSetCell);
         }

src/Microsoft.SqlTools.ServiceLayer/EditData/UpdateManagement/RowEdit.cs

@@ -109,7 +109,7 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
         /// </summary>
         /// <param name="columnId">Ordinal ID of the column to revert</param>
         /// <returns>String value of the original value of the cell</returns>
-        public abstract string RevertCell(int columnId);
+        public abstract EditRevertCellResult RevertCell(int columnId);
 
         /// <summary>
         /// Changes the value a cell in the row.

src/Microsoft.SqlTools.ServiceLayer/EditData/UpdateManagement/RowUpdate.cs

@@ -3,7 +3,6 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
-using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Data;
@@ -167,15 +166,21 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
         /// </summary>
         /// <param name="columnId">Ordinal of the column to revert</param>
         /// <returns>The value that was </returns>
-        public override string RevertCell(int columnId)
+        public override EditRevertCellResult RevertCell(int columnId)
         {
             Validate.IsWithinRange(nameof(columnId), columnId, 0, associatedRow.Count - 1);
 
             // Remove the cell update
+            // NOTE: This is best effort. The only way TryRemove can fail is if it is already
+            //       removed. If this happens, it is OK.
             CellUpdate cellUpdate;
             cellUpdates.TryRemove(columnId, out cellUpdate);
 
-            return associatedRow[columnId].DisplayValue;
+            return new EditRevertCellResult
+            {
+                IsRowDirty = cellUpdates.Count > 0,
+                Cell = new EditCell(associatedRow[columnId], false)
+            };
         }
 
         /// <summary>
@@ -203,10 +208,8 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
                 cellUpdates.TryRemove(columnId, out cu);
                 return new EditUpdateCellResult
                 {
-                    HasCorrections = false,
-                    NewValue = associatedRow[columnId].DisplayValue,
-                    IsRevert = true,
-                    IsNull = associatedRow[columnId].IsNull
+                    IsRowDirty = cellUpdates.Count > 0,
+                    Cell = new EditCell(associatedRow[columnId], false)
                 };
             }
 
@@ -214,10 +217,8 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement
             cellUpdates.AddOrUpdate(columnId, update, (i, cu) => update);
             return new EditUpdateCellResult
             {
-                HasCorrections = update.ValueAsString != newValue,
-                NewValue = update.ValueAsString != newValue ? update.ValueAsString : null,
-                IsNull = update.Value == DBNull.Value,
-                IsRevert = false            // If we're in this branch, it is not a revert
+                IsRowDirty = true,
+                Cell = update.AsEditCell
             };
         }
 

test/Microsoft.SqlTools.ServiceLayer.UnitTests/EditData/RowCreateTests.cs

@@ -263,15 +263,17 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             EditUpdateCellResult eucr = rc.SetCell(0, "1");
 
             // Then:
-            // ... The returned value should not have corrections
-            Assert.False(eucr.HasCorrections);
-            Assert.Null(eucr.NewValue);
-            
-            // ... The set value is not null
-            Assert.False(eucr.IsNull);
+            // ... The returned value should be equal to what we provided
+            Assert.NotNull(eucr);
+            Assert.NotNull(eucr.Cell);
+            Assert.Equal("1", eucr.Cell.DisplayValue);
+            Assert.False(eucr.Cell.IsNull);
 
-            // ... The result is not an implicit revert
-            Assert.False(eucr.IsRevert);
+            // ... The returned value should be dirty
+            Assert.NotNull(eucr.Cell.IsDirty);
+
+            // ... The row should still be dirty
+            Assert.True(eucr.IsRowDirty);
 
             // ... There should be a cell update in the cell list
             Assert.NotNull(rc.newCells[0]);
@@ -303,15 +305,17 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             EditUpdateCellResult eucr = rc.SetCell(0, "1000");
 
             // Then:
-            // ... The returned value should have corrections
-            Assert.True(eucr.HasCorrections);
-            Assert.NotEmpty(eucr.NewValue);
+            // ... The returned value should be equal to what we provided
+            Assert.NotNull(eucr);
+            Assert.NotNull(eucr.Cell);
+            Assert.NotEqual("1000", eucr.Cell.DisplayValue);
+            Assert.False(eucr.Cell.IsNull);
 
-            // ... The set value is not null
-            Assert.False(eucr.IsNull);
+            // ... The returned value should be dirty
+            Assert.NotNull(eucr.Cell.IsDirty);
 
-            // ... The result is not an implicit revert
-            Assert.False(eucr.IsRevert);
+            // ... The row should still be dirty
+            Assert.True(eucr.IsRowDirty);
 
             // ... There should be a cell update in the cell list
             Assert.NotNull(rc.newCells[0]);
@@ -327,15 +331,17 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             EditUpdateCellResult eucr = rc.SetCell(0, "NULL");
 
             // Then:
-            // ... The returned value should not have corrections
-            Assert.False(eucr.HasCorrections);
-            Assert.Null(eucr.NewValue);
+            // ... The returned value should be equal to what we provided
+            Assert.NotNull(eucr);
+            Assert.NotNull(eucr.Cell);
+            Assert.NotEmpty(eucr.Cell.DisplayValue);
+            Assert.True(eucr.Cell.IsNull);
 
-            // ... The set value is null
-            Assert.True(eucr.IsNull);
+            // ... The returned value should be dirty
+            Assert.NotNull(eucr.Cell.IsDirty);
 
-            // ... The result is not an implicit revert
-            Assert.False(eucr.IsRevert);
+            // ... The row should still be dirty
+            Assert.True(eucr.IsRowDirty);
 
             // ... There should be a cell update in the cell list
             Assert.NotNull(rc.newCells[0]);
@@ -362,11 +368,18 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             RowCreate rc = await GetStandardRowCreate();
 
             // If: I attempt to revert a cell that has not been set
-            string result = rc.RevertCell(0);
+            EditRevertCellResult result = rc.RevertCell(0);
 
-            // Then: We should get null back
+            // Then:
+            // ... We should get a result back
+            Assert.NotNull(result);
+
+            // ... We should get a null cell back
             // @TODO: Check for a default value when we support it
-            Assert.Null(result);
+            Assert.Null(result.Cell);
+
+            // ... The row should be dirty
+            Assert.True(result.IsRowDirty);
 
             // ... The cell should no longer be set
             Assert.Null(rc.newCells[0]);
@@ -380,11 +393,18 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             rc.SetCell(0, "1");
 
             // If: I attempt to revert a cell that was set
-            string result = rc.RevertCell(0);
+            EditRevertCellResult result = rc.RevertCell(0);
 
             // Then:
-            // ... We should get null back
-            Assert.Null(result);
+            // ... We should get a result back
+            Assert.NotNull(result);
+
+            // ... We should get a null cell back
+            // @TODO: Check for a default value when we support it
+            Assert.Null(result.Cell);
+
+            // ... The row should be dirty
+            Assert.True(result.IsRowDirty);
 
             // ... The cell should no longer be set
             Assert.Null(rc.newCells[0]);

test/Microsoft.SqlTools.ServiceLayer.UnitTests/EditData/RowEditBaseTests.cs

@@ -320,7 +320,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
                 throw new NotImplementedException();
             }
 
-            public override string RevertCell(int columnId)
+            public override EditRevertCellResult RevertCell(int columnId)
             {
                 throw new NotImplementedException();
             }

test/Microsoft.SqlTools.ServiceLayer.UnitTests/EditData/RowUpdateTests.cs

@@ -49,15 +49,19 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             EditUpdateCellResult eucr = ru.SetCell(0, "col1");
 
             // Then:
-            // ... The returned value should not have corrections
-            Assert.False(eucr.HasCorrections);
-            Assert.Null(eucr.NewValue);
+            // ... A edit cell was returned
+            Assert.NotNull(eucr);
+            Assert.NotNull(eucr.Cell);
 
-            // ... The set value is not null
-            Assert.False(eucr.IsNull);
+            // ... The new value we provided should be returned
+            Assert.Equal("col1", eucr.Cell.DisplayValue);
+            Assert.False(eucr.Cell.IsNull);
 
-            // ... The result is not an implicit revert
-            Assert.False(eucr.IsRevert);
+            // ... The row is still dirty
+            Assert.True(eucr.IsRowDirty);
+
+            // ... The cell should be dirty
+            Assert.True(eucr.Cell.IsDirty);
 
             // ... There should be a cell update in the cell list
             Assert.Contains(0, ru.cellUpdates.Keys);
@@ -93,15 +97,20 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             EditUpdateCellResult eucr = ru.SetCell(0, "1000");
 
             // Then:
-            // ... The returned value should have corrections
-            Assert.True(eucr.HasCorrections);
-            Assert.NotEmpty(eucr.NewValue);
+            // ... A edit cell was returned
+            Assert.NotNull(eucr);
+            Assert.NotNull(eucr.Cell);
 
-            // ... The set value is not null
-            Assert.False(eucr.IsNull);
+            // ... The value we used won't be returned
+            Assert.NotEmpty(eucr.Cell.DisplayValue);
+            Assert.NotEqual("1000", eucr.Cell.DisplayValue);
+            Assert.False(eucr.Cell.IsNull);
 
-            // ... The result is not an implicit revert
-            Assert.False(eucr.IsRevert);
+            // ... The cell should be dirty
+            Assert.True(eucr.Cell.IsDirty);
+
+            // ... The row is still dirty
+            Assert.True(eucr.IsRowDirty);
 
             // ... There should be a cell update in the cell list
             Assert.Contains(0, ru.cellUpdates.Keys);
@@ -122,15 +131,22 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             Common.AddCells(ru, true);
 
             // ... Then I update a cell back to it's old value
-            var output = ru.SetCell(1, (string) rs.GetRow(0)[1].RawObject);
+            var eucr = ru.SetCell(1, (string) rs.GetRow(0)[1].RawObject);
 
             // Then:
-            // ... The output should indicate a revert
-            Assert.NotNull(output);
-            Assert.True(output.IsRevert);
-            Assert.False(output.HasCorrections);
-            Assert.False(output.IsNull);
-            Assert.Equal(rs.GetRow(0)[1].DisplayValue, output.NewValue);
+            // ... A edit cell was returned
+            Assert.NotNull(eucr);
+            Assert.NotNull(eucr.Cell);
+
+            // ... The new value we provided should be returned
+            Assert.Equal(rs.GetRow(0)[1].DisplayValue, eucr.Cell.DisplayValue);
+            Assert.False(eucr.Cell.IsNull);
+
+            // ... The cell should be clean
+            Assert.False(eucr.Cell.IsDirty);
+
+            // ... The row is still dirty
+            Assert.True(eucr.IsRowDirty);
 
             // ... It should be formatted as an update script
             Regex r = new Regex(@"UPDATE .+ SET (.*) WHERE");
@@ -143,6 +159,39 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             Assert.All(updateSplit, s => Assert.Equal(2, s.Split('=').Length));
         }
 
+        public async Task SetCellImplicitRowRevertTests()
+        {
+            // Setup: Create a fake column to update
+            DbColumn[] columns = Common.GetColumns(true);
+            ResultSet rs = await Common.GetResultSet(columns, true);
+            EditTableMetadata etm = Common.GetStandardMetadata(columns);
+
+            // If:
+            // ... I add updates to one cell in the row
+            RowUpdate ru = new RowUpdate(0, rs, etm);
+            ru.SetCell(1, "qqq");
+
+            // ... Then I update the cell to its original value
+            var eucr = ru.SetCell(1, (string) rs.GetRow(0)[1].RawObject);
+
+            // Then:
+            // ... An edit cell should have been returned
+            Assert.NotNull(eucr);
+            Assert.NotNull(eucr.Cell);
+
+            // ... The old value should be returned
+            Assert.Equal(rs.GetRow(0)[1].DisplayValue, eucr.Cell.DisplayValue);
+            Assert.False(eucr.Cell.IsNull);
+
+            // ... The cell should be clean
+            Assert.False(eucr.Cell.IsDirty);
+
+            // ... The row should be clean
+            Assert.False(eucr.IsRowDirty);
+
+            // TODO: Make sure that the script and command things will return null
+        }
+
         [Theory]
         [InlineData(true)]
         [InlineData(false)]
@@ -368,10 +417,19 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             RowUpdate ru = new RowUpdate(0, rs, etm);
 
             // If: I attempt to revert a cell that has not been set
-            string result = ru.RevertCell(0);
+            EditRevertCellResult result = ru.RevertCell(0);
 
-            // Then: We should get the original value back
-            Assert.NotEmpty(result);
+            // Then:
+            // ... We should get a result back
+            Assert.NotNull(result);
+
+            // ... We should get the original value back
+            // @TODO: Check for a default value when we support it
+            Assert.NotNull(result.Cell);
+            Assert.Equal(rs.GetRow(0)[0].DisplayValue, result.Cell.DisplayValue);
+
+            // ... The row should be clean
+            Assert.False(result.IsRowDirty);
 
             // ... The cell should no longer be set
             Assert.DoesNotContain(0, ru.cellUpdates.Keys);
@@ -386,14 +444,53 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             var rs = await Common.GetResultSet(columns, false);
             var etm = Common.GetStandardMetadata(columns);
             RowUpdate ru = new RowUpdate(0, rs, etm);
-            ru.SetCell(0, "1");
+            ru.SetCell(0, "qqq");
+            ru.SetCell(1, "qqq");
 
             // If: I attempt to revert a cell that was set
-            string result = ru.RevertCell(0);
+            EditRevertCellResult result = ru.RevertCell(0);
 
             // Then:
+            // ... We should get a result back
+            Assert.NotNull(result);
+
             // ... We should get the original value back
-            Assert.NotEmpty(result);
+            // @TODO: Check for a default value when we support it
+            Assert.NotNull(result.Cell);
+            Assert.Equal(rs.GetRow(0)[0].DisplayValue, result.Cell.DisplayValue);
+
+            // ... The row should be dirty still
+            Assert.True(result.IsRowDirty);
+
+            // ... The cell should no longer be set
+            Assert.DoesNotContain(0, ru.cellUpdates.Keys);
+        }
+
+        [Fact]
+        public async Task RevertCellRevertsRow()
+        {
+            // Setup:
+            // ... Create a row update
+            var columns = Common.GetColumns(false);
+            var rs = await Common.GetResultSet(columns, false);
+            var etm = Common.GetStandardMetadata(columns);
+            RowUpdate ru = new RowUpdate(0, rs, etm);
+            ru.SetCell(0, "qqq");
+
+            // If: I attempt to revert a cell that was set
+            EditRevertCellResult result = ru.RevertCell(0);
+
+            // Then:
+            // ... We should get a result back
+            Assert.NotNull(result);
+
+            // ... We should get the original value back
+            // @TODO: Check for a default value when we support it
+            Assert.NotNull(result.Cell);
+            Assert.Equal(rs.GetRow(0)[0].DisplayValue, result.Cell.DisplayValue);
+
+            // ... The row should now be reverted
+            Assert.False(result.IsRowDirty);
 
             // ... The cell should no longer be set
             Assert.DoesNotContain(0, ru.cellUpdates.Keys);

test/Microsoft.SqlTools.ServiceLayer.UnitTests/EditData/ServiceIntegrationTests.cs

@@ -10,6 +10,7 @@ using Microsoft.SqlTools.ServiceLayer.EditData;
 using Microsoft.SqlTools.ServiceLayer.EditData.Contracts;
 using Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Test.Common;
 using Microsoft.SqlTools.ServiceLayer.UnitTests.Utility;
 using Moq;
@@ -190,10 +191,8 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             var edit = new Mock<RowEditBase>();
             edit.Setup(e => e.SetCell(It.IsAny<int>(), It.IsAny<string>())).Returns(new EditUpdateCellResult
             {
-                NewValue = string.Empty,
-                HasCorrections = true,
-                IsRevert = false,
-                IsNull = false
+                IsRowDirty = true,
+                Cell = new EditCell(new DbCellValue(), true)
             });
             session.EditCache[0] = edit.Object;
 
@@ -202,9 +201,8 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
                 .AddResultValidation(eucr =>
                 {
                     Assert.NotNull(eucr);
-                    Assert.NotNull(eucr.NewValue);
-                    Assert.False(eucr.IsRevert);
-                    Assert.False(eucr.IsNull);
+                    Assert.NotNull(eucr.Cell);
+                    Assert.True(eucr.IsRowDirty);
                 })
                 .Complete();
             await eds.HandleUpdateCellRequest(new EditUpdateCellParams { OwnerUri = Constants.OwnerUri, RowId = 0}, efv.Object);

test/Microsoft.SqlTools.ServiceLayer.UnitTests/EditData/SessionTests.cs

@@ -627,6 +627,30 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             Assert.Throws<InvalidOperationException>(() => s.RevertCell(0, 0));
         }
 
+        [Fact]
+        public async Task RevertCellRowRevert()
+        {
+            // Setup:
+            // ... Create a session with a proper query and metadata
+            EditSession s = await GetBasicSession();
+
+            // ... Add a edit that we will say the edit was reverted if we update a cell
+            var mockEdit = new Mock<RowEditBase>();
+            mockEdit.Setup(e => e.RevertCell(It.IsAny<int>()))
+                .Returns(new EditRevertCellResult {IsRowDirty = false});
+            s.EditCache[0] = mockEdit.Object;
+
+            // If: I update a cell that will return an implicit revert
+            s.RevertCell(0, 0);
+
+            // Then:
+            // ... Set cell should have been called on the mock update once
+            mockEdit.Verify(e => e.RevertCell(0), Times.Once);
+
+            // ... The mock update should no longer be in the edit cache
+            Assert.Empty(s.EditCache);
+        }
+
         #endregion
 
         #region Update Cell Tests
@@ -653,7 +677,8 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
 
             // ... Add a mock edit to the edit cache to cause the .TryAdd to fail
             var mockEdit = new Mock<RowEditBase>();
-            mockEdit.Setup(e => e.SetCell(It.IsAny<int>(), It.IsAny<string>()));
+            mockEdit.Setup(e => e.SetCell(It.IsAny<int>(), It.IsAny<string>()))
+                .Returns(new EditUpdateCellResult {IsRowDirty = true});
             s.EditCache[0] = mockEdit.Object;
 
             // If: I update a cell on a row that already has a pending edit
@@ -681,6 +706,32 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             Assert.IsType<RowUpdate>(s.EditCache[0]);
         }
 
+        [Fact]
+        public async Task UpdateCellRowRevert()
+        {
+            // Setup:
+            // ... Create a session with a proper query and metadata
+            EditSession s = await GetBasicSession();
+
+            // ... Add a edit that we will say the edit was reverted if we update a cell
+            var mockEdit = new Mock<RowEditBase>();
+            mockEdit.Setup(e => e.SetCell(It.IsAny<int>(), It.IsAny<string>()))
+                .Returns(new EditUpdateCellResult {IsRowDirty = false});
+            s.EditCache[0] = mockEdit.Object;
+
+            // If: I update a cell that will return an implicit revert
+            s.UpdateCell(0, 0, null);
+
+            // Then:
+            // ... Set cell should have been called on the mock update once
+            mockEdit.Verify(e => e.SetCell(0, null), Times.Once);
+
+            // ... The mock update should no longer be in the edit cache
+            Assert.Empty(s.EditCache);
+        }
+
+
+
         #endregion
 
         #region SubSet Tests