Support TOP n When Initializing Edit Sessions (#290)

* Start of edit initialize filtering

* Implementation to utilize filtering

* Unit tests for the refactoring

* Unit tests for query construction

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

@@ -0,0 +1,18 @@
+//
+// 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>
+    /// Parameters for filtering a the rows in a table to make querying easier
+    /// </summary>
+    public class EditInitializeFiltering
+    {
+        /// <summary>
+        /// Limit the records queried from the database to this many. If null, all rows are returned
+        /// </summary>
+        public int? LimitResults { get; set; }
+    }
+}

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

@@ -12,6 +12,11 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData.Contracts
     /// </summary>
     public class EditInitializeParams : SessionOperationParams
     {
+        /// <summary>
+        /// Filtering parameters
+        /// </summary>
+        public EditInitializeFiltering Filters { get; set; }
+
         /// <summary>
         /// The object to use for generating an edit script
         /// </summary>

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

@@ -162,14 +162,14 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
                 Validate.IsNotNullOrWhitespaceString(nameof(initParams.ObjectType), initParams.ObjectType);
 
                 // Create a session and add it to the session list
-                EditSession session = new EditSession(metadataFactory, initParams.ObjectName, initParams.ObjectType);
+                EditSession session = new EditSession(metadataFactory);
                 if (!ActiveSessions.TryAdd(initParams.OwnerUri, session))
                 {
                     throw new InvalidOperationException(SR.EditDataSessionAlreadyExists);
                 }
 
                 // Initialize the session
-                session.Initialize(connector, queryRunner, executionSuccessHandler, executionFailureHandler);
+                session.Initialize(initParams, connector, queryRunner, executionSuccessHandler, executionFailureHandler);
             }
             catch (Exception e)
             {

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

@@ -9,6 +9,7 @@ using System.Collections.Generic;
 using System.Data.Common;
 using System.IO;
 using System.Linq;
+using System.Text;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.ServiceLayer.EditData.Contracts;
 using Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement;
@@ -29,25 +30,17 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
 
         private readonly IEditMetadataFactory metadataFactory;
         private EditTableMetadata objectMetadata;
-        private readonly string objectName;
-        private readonly string objectType;
 
         /// <summary>
         /// Constructs a new edit session bound to the result set and metadat object provided
         /// </summary>
         /// <param name="metaFactory">Factory for creating metadata</param>
-        /// <param name="objName">The name of the object to edit</param>
+        public EditSession(IEditMetadataFactory metaFactory)
-        /// <param name="objType">The type of the object to edit</param>
-        public EditSession(IEditMetadataFactory metaFactory, string objName, string objType)
         {
             Validate.IsNotNull(nameof(metaFactory), metaFactory);
-            Validate.IsNotNullOrWhitespaceString(nameof(objName), objName);
-            Validate.IsNotNullOrWhitespaceString(nameof(objType), objType);
 
             // Setup the internal state
             metadataFactory = metaFactory;
-            objectName = objName;
-            objectType = objType;
         }
 
         #region Properties
@@ -86,7 +79,22 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
 
         #region Public Methods
 
-        public void Initialize(Connector connector, QueryRunner queryRunner, Func<Task> successHandler, Func<Exception, Task> errorHandler)
+        /// <summary>
+        /// Initializes the edit session, asynchronously, by retrieving metadata about the table to
+        /// edit and querying the table for the rows of the table.
+        /// </summary>
+        /// <param name="initParams">Parameters for initializing the edit session</param>
+        /// <param name="connector">Delegate that will return a DbConnection when executed</param>
+        /// <param name="queryRunner">
+        /// Delegate that will run the requested query and return a
+        /// <see cref="EditSessionQueryExecutionState"/> object on execution
+        /// </param>
+        /// <param name="successHandler">Func to call when initialization has completed successfully</param>
+        /// <param name="errorHandler">Func to call when initialization has completed with errors</param>
+        /// <exception cref="InvalidOperationException">
+        /// When session is already initialized or in progress of initializing
+        /// </exception>
+        public void Initialize(EditInitializeParams initParams, Connector connector, QueryRunner queryRunner, Func<Task> successHandler, Func<Exception, Task> errorHandler)
         {
             if (IsInitialized)
             {
@@ -98,13 +106,17 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
                 throw new InvalidOperationException(SR.EditDataSessionAlreadyInitializing);
             }
 
+            Validate.IsNotNullOrWhitespaceString(nameof(initParams.ObjectName), initParams.ObjectName);
+            Validate.IsNotNullOrWhitespaceString(nameof(initParams.ObjectType), initParams.ObjectType);
+            Validate.IsNotNull(nameof(initParams.Filters), initParams.Filters);
+
             Validate.IsNotNull(nameof(connector), connector);
             Validate.IsNotNull(nameof(queryRunner), queryRunner);
             Validate.IsNotNull(nameof(successHandler), successHandler);
             Validate.IsNotNull(nameof(errorHandler), errorHandler);
 
             // Start up the initialize process
-            InitializeTask = InitializeInternal(connector, queryRunner, successHandler, errorHandler);
+            InitializeTask = InitializeInternal(initParams, connector, queryRunner, successHandler, errorHandler);
         }
 
         /// <summary>
@@ -403,16 +415,19 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
 
         #endregion
 
-        private async Task InitializeInternal(Connector connector, QueryRunner queryRunner,
+        #region Private Helpers
-            Func<Task> successHandler, Func<Exception, Task> failureHandler)
+
+        private async Task InitializeInternal(EditInitializeParams initParams, Connector connector,
+            QueryRunner queryRunner, Func<Task> successHandler, Func<Exception, Task> failureHandler)
         {
             try
             {
                 // Step 1) Look up the SMO metadata
-                objectMetadata = metadataFactory.GetObjectMetadata(await connector(), objectName, objectType);
+                objectMetadata = metadataFactory.GetObjectMetadata(await connector(), initParams.ObjectName,
+                    initParams.ObjectType);
 
                 // Step 2) Get and execute a query for the rows in the object we're looking up
-                EditSessionQueryExecutionState state = await queryRunner(ConstructInitializeQuery());
+                EditSessionQueryExecutionState state = await queryRunner(ConstructInitializeQuery(objectMetadata, initParams.Filters));
                 if (state.Query == null)
                 {
                     // TODO: Move to SR file
@@ -467,13 +482,33 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
             }
         }
 
-        private string ConstructInitializeQuery()
+        /// <summary>
+        /// Constructs a query for selecting rows in a table based on the filters provided.
+        /// Internal for unit testing purposes only.
+        /// </summary>
+        internal static string ConstructInitializeQuery(EditTableMetadata metadata, EditInitializeFiltering initFilters)
         {
-            // Using the columns we know, put together a query for the rows in the table
+            StringBuilder queryBuilder = new StringBuilder("SELECT ");
-            var columns = objectMetadata.Columns.Select(col => col.EscapedName);
-            var columnClause = string.Join(", ", columns);
 
-            return $"SELECT ${columnClause} FROM ${objectMetadata.EscapedMultipartName}";
+            // If there is a filter for top n rows, then apply it
+            if (initFilters.LimitResults.HasValue)
+            {
+                if (initFilters.LimitResults < 0)
+                {
+                    throw new ArgumentOutOfRangeException(nameof(initFilters.LimitResults), SR.EditDataFilteringNegativeLimit);
+                }
+                queryBuilder.AppendFormat("TOP {0} ", initFilters.LimitResults.Value);
+            }
+
+            // Using the columns we know, add them to the query
+            var columns = metadata.Columns.Select(col => col.EscapedName);
+            var columnClause = string.Join(", ", columns);
+            queryBuilder.Append(columnClause);
+
+            // Add the FROM
+            queryBuilder.AppendFormat(" FROM {0}", metadata.EscapedMultipartName);
+
+            return queryBuilder.ToString();
         }
 
         private void ThrowIfNotInitialized()
@@ -484,6 +519,8 @@ namespace Microsoft.SqlTools.ServiceLayer.EditData
             }
         }
 
+        #endregion
+
         /// <summary>
         /// State object to return upon completion of an edit session intialization query
         /// </summary>

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.cs

@@ -413,6 +413,14 @@ namespace Microsoft.SqlTools.ServiceLayer
             }
         }
 
+        public static string EditDataFilteringNegativeLimit
+        {
+            get
+            {
+                return Keys.GetString(Keys.EditDataFilteringNegativeLimit);
+            }
+        }
+
         public static string EditDataQueryFailed
         {
             get
@@ -1098,6 +1106,9 @@ namespace Microsoft.SqlTools.ServiceLayer
             public const string EditDataSessionAlreadyInitializing = "EditDataSessionAlreadyInitializing";
 
 
+            public const string EditDataFilteringNegativeLimit = "EditDataFilteringNegativeLimit";
+
+
             public const string EditDataUnsupportedObjectType = "EditDataUnsupportedObjectType";
 
 

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.resx

@@ -374,6 +374,10 @@
     <value>Edit session has already been initialized or is in the process of initializing</value>
     <comment></comment>
   </data>
+  <data name="EditDataFilteringNegativeLimit" xml:space="preserve">
+    <value>Result limit cannot be negative</value>
+    <comment></comment>
+  </data>
   <data name="EditDataUnsupportedObjectType" xml:space="preserve">
     <value>Database object {0} cannot be used for editing.</value>
     <comment>.

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.strings

@@ -176,6 +176,8 @@ EditDataSessionAlreadyInitialized = Edit session has already been initialized
 
 EditDataSessionAlreadyInitializing = Edit session has already been initialized or is in the process of initializing
 
+EditDataFilteringNegativeLimit = Result limit cannot be negative
+
 EditDataUnsupportedObjectType(string typeName) = Database object {0} cannot be used for editing.
 
 EditDataQueryFailed = Query execution failed, see messages for details

src/Microsoft.SqlTools.ServiceLayer/Localization/sr.xlf

@@ -591,6 +591,11 @@
         <target state="new">Query execution failed, see messages for details</target>
         <note></note>
       </trans-unit>
+      <trans-unit id="EditDataFilteringNegativeLimit">
+        <source>Result limit cannot be negative</source>
+        <target state="new">Result limit cannot be negative</target>
+        <note></note>
+      </trans-unit>
     </body>
   </file>
 </xliff>

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

@@ -9,6 +9,7 @@ using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 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;
@@ -22,21 +23,42 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
     {
         public const string OwnerUri = "testFile";
 
+        public static EditInitializeParams BasicInitializeParameters
+        {
+            get
+            {
+                return new EditInitializeParams
+                {
+                    Filters = new EditInitializeFiltering(),
+                    ObjectName = "tbl",
+                    ObjectType = "tbl"
+                };
+            }
+        }
+
         public static async Task<EditSession> GetCustomSession(Query q, EditTableMetadata etm)
         {
+            // Step 1) Create the Session object
             // Mock metadata factory
             Mock<IEditMetadataFactory> metaFactory = new Mock<IEditMetadataFactory>();
             metaFactory
                 .Setup(f => f.GetObjectMetadata(It.IsAny<DbConnection>(), It.IsAny<string>(), It.IsAny<string>()))
                 .Returns(etm);
 
-            EditSession session = new EditSession(metaFactory.Object, "tbl", "tbl");
+            EditSession session = new EditSession(metaFactory.Object);
 
+
+            // Step 2) Initialize the Session
+            // Mock connector that does nothing
             EditSession.Connector connector = () => Task.FromResult<DbConnection>(null);
+
+            // Mock query runner that returns the query we were provided
             EditSession.QueryRunner queryRunner = (s) => Task.FromResult(new EditSession.EditSessionQueryExecutionState(q));
 
-            session.Initialize(connector, queryRunner, () => Task.FromResult(0), (e) => Task.FromResult(0));
+            // Initialize
+            session.Initialize(BasicInitializeParameters, connector, queryRunner, () => Task.FromResult(0), (e) => Task.FromResult(0));
             await session.InitializeTask;
+
             return session;
         }
 

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

@@ -8,6 +8,7 @@ using System.Collections.Generic;
 using System.Data.Common;
 using System.IO;
 using System.Linq;
+using System.Text.RegularExpressions;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.EditData;
@@ -32,31 +33,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
         {
             // If: I create a session object with a null metadata factory
             // Then: It should throw an exception
-            Assert.Throws<ArgumentNullException>(() => new EditSession(null, Constants.OwnerUri, Constants.OwnerUri));
+            Assert.Throws<ArgumentNullException>(() => new EditSession(null));
-        }
-
-        [Theory]
-        [InlineData(null)]
-        [InlineData("")]
-        [InlineData(" \t\r\n")]
-        public void SessionConstructionNullObjectName(string objName)
-        {
-            // If: I create a session object with a null or whitespace object name
-            // Then: It should throw an exception
-            Mock<IEditMetadataFactory> mockFactory = new Mock<IEditMetadataFactory>();
-            Assert.Throws<ArgumentException>(() => new EditSession(mockFactory.Object, objName, Constants.OwnerUri));
-        }
-
-        [Theory]
-        [InlineData(null)]
-        [InlineData("")]
-        [InlineData(" \t\r\n")]
-        public void SessionConstructionNullObjectType(string objType)
-        {
-            // If: I create a session object with a null or whitespace object type
-            // Then: It should throw an exception
-            Mock<IEditMetadataFactory> mockFactory = new Mock<IEditMetadataFactory>();
-            Assert.Throws<ArgumentException>(() => new EditSession(mockFactory.Object, Constants.OwnerUri, objType));
         }
 
         [Fact]
@@ -64,7 +41,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
         {
             // If: I create a session object with a proper arguments
             Mock<IEditMetadataFactory> mockFactory = new Mock<IEditMetadataFactory>();
-            EditSession s = new EditSession(mockFactory.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(mockFactory.Object);
 
             // Then:
             // ... The edit cache should not exist
@@ -138,7 +115,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             // Setup:
             // ... Create a session without initializing
             Mock<IEditMetadataFactory> emf = new Mock<IEditMetadataFactory>();
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object);
 
             // If: I ask to create a row without initializing
             // Then: I should get an exception
@@ -311,12 +288,11 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             // Setup:
             // ... Create a session and fake that it has been initialized
             Mock<IEditMetadataFactory> emf = new Mock<IEditMetadataFactory>();
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object) {IsInitialized = true};
-            s.IsInitialized = true;
 
             // If: I initialize it
             // Then: I should get an exception
-            Assert.Throws<InvalidOperationException>(() => s.Initialize(null, null, null, null));
+            Assert.Throws<InvalidOperationException>(() => s.Initialize(null, null, null, null, null));
         }
 
         [Fact]
@@ -325,37 +301,75 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             // Setup:
             // ... Create a session and fake that it is in progress of initializing
             Mock<IEditMetadataFactory> emf = new Mock<IEditMetadataFactory>();
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object) {InitializeTask = new Task(() => {})};
-            s.InitializeTask = new Task(() => { });
 
             // If: I initialize it
             // Then: I should get an exception
-            Assert.Throws<InvalidOperationException>(() => s.Initialize(null, null, null, null));
+            Assert.Throws<InvalidOperationException>(() => s.Initialize(null, null, null, null, null));
         }
 
         [Theory]
         [MemberData(nameof(InitializeNullParamsData))]
-        public void InitializeNullParams(EditSession.Connector c, EditSession.QueryRunner qr,
+        public void InitializeNullParams(EditInitializeParams initParams, EditSession.Connector c,
-            Func<Task> sh, Func<Exception, Task> fh)
+            EditSession.QueryRunner qr, Func<Task> sh, Func<Exception, Task> fh)
         {
             // Setup:
             // ... Create a session that hasn't been initialized
             Mock<IEditMetadataFactory> emf = new Mock<IEditMetadataFactory>();
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object);
 
             // If: I initialize it with a missing parameter
             // Then: It should throw an exception
-            Assert.ThrowsAny<ArgumentException>(() => s.Initialize(c, qr, sh, fh));
+            Assert.ThrowsAny<ArgumentException>(() => s.Initialize(initParams, c, qr, sh, fh));
         }
 
         public static IEnumerable<object> InitializeNullParamsData
         {
             get
             {
-                yield return new object[] {null, DoNothingQueryRunner, DoNothingSuccessHandler, DoNothingFailureHandler};
+                yield return new object[] {Common.BasicInitializeParameters, null, DoNothingQueryRunner, DoNothingSuccessHandler, DoNothingFailureHandler};
-                yield return new object[] {DoNothingConnector, null, DoNothingSuccessHandler, DoNothingFailureHandler};
+                yield return new object[] {Common.BasicInitializeParameters, DoNothingConnector, null, DoNothingSuccessHandler, DoNothingFailureHandler};
-                yield return new object[] {DoNothingConnector, DoNothingQueryRunner, null, DoNothingFailureHandler};
+                yield return new object[] {Common.BasicInitializeParameters, DoNothingConnector, DoNothingQueryRunner, null, DoNothingFailureHandler};
-                yield return new object[] {DoNothingConnector, DoNothingQueryRunner, DoNothingSuccessHandler, null};
+                yield return new object[] {Common.BasicInitializeParameters, DoNothingConnector, DoNothingQueryRunner, DoNothingSuccessHandler, null};
+
+                string[] nullOrWhitespace = {null, string.Empty, " \t\r\n"};
+
+                // Tests with invalid object name or type
+                foreach (string str in nullOrWhitespace)
+                {
+                    // Invalid object name
+                    var eip1 = new EditInitializeParams
+                    {
+                        ObjectName = str,
+                        ObjectType = "table",
+                        Filters = new EditInitializeFiltering()
+                    };
+                    yield return new object[] {eip1, DoNothingConnector, DoNothingQueryRunner, DoNothingSuccessHandler, DoNothingFailureHandler};
+
+                    // Invalid object type
+                    var eip2 = new EditInitializeParams
+                    {
+                        ObjectName = "tbl",
+                        ObjectType = str,
+                        Filters = new EditInitializeFiltering()
+                    };
+                    yield return new object[] {eip2, DoNothingConnector, DoNothingQueryRunner, DoNothingSuccessHandler, DoNothingFailureHandler};
+                }
+
+                // Test with null init filters
+                yield return new object[]
+                {
+                    new EditInitializeParams
+                    {
+                        ObjectName = "tbl",
+                        ObjectType = "table",
+                        Filters = null
+                    },
+                    DoNothingConnector,
+                    DoNothingQueryRunner,
+                    DoNothingSuccessHandler,
+                    DoNothingFailureHandler
+                };
             }
         }
 
@@ -369,14 +383,14 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
                 .Throws<Exception>();
 
             // ... Create a session that hasn't been initialized
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object);
 
             // ... Create a mock for verifying the failure handler will be called
             var successHandler = DoNothingSuccessMock;
             var failureHandler = DoNothingFailureMock;
 
             // If: I initalize the session with a metadata factory that will fail
-            s.Initialize(DoNothingConnector, DoNothingQueryRunner, successHandler.Object, failureHandler.Object);
+            s.Initialize(Common.BasicInitializeParameters, DoNothingConnector, DoNothingQueryRunner, successHandler.Object, failureHandler.Object);
             await s.InitializeTask;
 
             // Then:
@@ -402,7 +416,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
                 .Returns(etm);
 
             // ... Create a session that hasn't been initialized
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object);
 
             // ... Create a query runner that will fail via exception
             Mock<EditSession.QueryRunner> qr = new Mock<EditSession.QueryRunner>();
@@ -413,7 +427,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             var failureHandler = DoNothingFailureMock;
 
             // If: I initialize the session with a query runner that will fail
-            s.Initialize(DoNothingConnector, qr.Object, successHandler.Object, failureHandler.Object);
+            s.Initialize(Common.BasicInitializeParameters, DoNothingConnector, qr.Object, successHandler.Object, failureHandler.Object);
             await s.InitializeTask;
 
             // Then:
@@ -441,7 +455,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
                 .Returns(etm);
 
             // ... Create a session that hasn't been initialized
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object);
 
             // ... Create a query runner that will fail via returning a null query
             Mock<EditSession.QueryRunner> qr = new Mock<EditSession.QueryRunner>();
@@ -453,7 +467,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             var failureHandler = DoNothingFailureMock;
 
             // If: I initialize the session with a query runner that will fail
-            s.Initialize(DoNothingConnector, qr.Object, successHandler.Object, failureHandler.Object);
+            s.Initialize(Common.BasicInitializeParameters, DoNothingConnector, qr.Object, successHandler.Object, failureHandler.Object);
             await s.InitializeTask;
 
             // Then:
@@ -480,7 +494,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
                 .Returns(etm);
 
             // ... Create a session that hasn't been initialized
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object);
 
             // ... Create a query runner that will return a successful query
             Mock<EditSession.QueryRunner> qr = new Mock<EditSession.QueryRunner>();
@@ -492,7 +506,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             var failureHandler = DoNothingFailureMock;
 
             // If: I initialize the session with a query runner that will fail
-            s.Initialize(DoNothingConnector, qr.Object, successHandler.Object, failureHandler.Object);
+            s.Initialize(Common.BasicInitializeParameters, DoNothingConnector, qr.Object, successHandler.Object, failureHandler.Object);
             await s.InitializeTask;
 
             // Then:
@@ -519,7 +533,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             // Setup:
             // ... Create a session without initializing
             Mock<IEditMetadataFactory> emf = new Mock<IEditMetadataFactory>();
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object);
 
             // If: I ask to delete a row without initializing
             // Then: I should get an exception
@@ -570,7 +584,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             // Setup:
             // ... Create a session without initializing
             Mock<IEditMetadataFactory> emf = new Mock<IEditMetadataFactory>();
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object);
 
             // If: I ask to revert a row without initializing
             // Then: I should get an exception
@@ -606,7 +620,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             // Setup:
             // ... Create a session without initializing
             Mock<IEditMetadataFactory> emf = new Mock<IEditMetadataFactory>();
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object);
 
             // If: I ask to revert a cell without initializing
             // Then: I should get an exception
@@ -623,7 +637,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             // Setup:
             // ... Create a session without initializing
             Mock<IEditMetadataFactory> emf = new Mock<IEditMetadataFactory>();
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object);
 
             // If: I ask to update a cell without initializing
             // Then: I should get an exception
@@ -677,7 +691,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             // Setup:
             // ... Create a session without initializing
             Mock<IEditMetadataFactory> emf = new Mock<IEditMetadataFactory>();
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object);
 
             // If: I ask to update a cell without initializing
             // Then: I should get an exception
@@ -844,7 +858,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             // Setup:
             // ... Create a session without initializing
             Mock<IEditMetadataFactory> emf = new Mock<IEditMetadataFactory>();
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object);
 
             // If: I ask to script edits without initializing
             // Then: I should get an exception
@@ -902,7 +916,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             // Setup:
             // ... Create a session without initializing
             Mock<IEditMetadataFactory> emf = new Mock<IEditMetadataFactory>();
-            EditSession s = new EditSession(emf.Object, Constants.OwnerUri, Constants.OwnerUri);
+            EditSession s = new EditSession(emf.Object);
 
             // If: I ask to script edits without initializing
             // Then: I should get an exception
@@ -1057,6 +1071,91 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
 
         #endregion
 
+        #region Construct Initialize Query Tests
+
+        [Fact]
+        public void ConstructQueryWithoutLimit()
+        {
+            // Setup: Create a metadata provider for some basic columns
+            var cols = Common.GetColumns(false);
+            var etm = Common.GetStandardMetadata(cols);
+
+            // If: I generate a query for selecting rows without a limit
+            EditInitializeFiltering eif = new EditInitializeFiltering
+            {
+                LimitResults = null
+            };
+            string query = EditSession.ConstructInitializeQuery(etm, eif);
+
+            // Then:
+            // ... The query should look like a select statement
+            Regex selectRegex = new Regex("SELECT (.+) FROM (.+)", RegexOptions.IgnoreCase);
+            var match = selectRegex.Match(query);
+            Assert.True(match.Success);
+
+            // ... There should be columns in it
+            Assert.Equal(etm.Columns.Length, match.Groups[1].Value.Split(',').Length);
+
+            // ... The table name should be in it
+            Assert.Equal(etm.EscapedMultipartName, match.Groups[2].Value);
+
+            // ... It should NOT have a TOP clause in it
+            Assert.DoesNotContain("TOP", query);
+        }
+
+        [Fact]
+        public void ConstructQueryNegativeLimit()
+        {
+            // Setup: Create a metadata provider for some basic columns
+            var cols = Common.GetColumns(false);
+            var etm = Common.GetStandardMetadata(cols);
+
+            // If: I generate a query for selecting rows with a negative limit
+            // Then: An exception should be thrown
+            EditInitializeFiltering eif = new EditInitializeFiltering
+            {
+                LimitResults = -1
+            };
+            Assert.Throws<ArgumentOutOfRangeException>(() => EditSession.ConstructInitializeQuery(etm, eif));
+        }
+
+        [Theory]
+        [InlineData(0)]        // Yes, zero is valid
+        [InlineData(10)]
+        [InlineData(1000)]
+        public void ConstructQueryWithLimit(int limit)
+        {
+            // Setup: Create a metadata provider for some basic columns
+            var cols = Common.GetColumns(false);
+            var etm = Common.GetStandardMetadata(cols);
+
+            // If: I generate a query for selecting rows without a limit
+            EditInitializeFiltering eif = new EditInitializeFiltering
+            {
+                LimitResults = limit
+            };
+            string query = EditSession.ConstructInitializeQuery(etm, eif);
+
+            // Then:
+            // ... The query should look like a select statement
+            Regex selectRegex = new Regex(@"SELECT TOP (\d+) (.+) FROM (.+)", RegexOptions.IgnoreCase);
+            var match = selectRegex.Match(query);
+            Assert.True(match.Success);
+
+            // ... There should be columns in it
+            Assert.Equal(etm.Columns.Length, match.Groups[2].Value.Split(',').Length);
+
+            // ... The table name should be in it
+            Assert.Equal(etm.EscapedMultipartName, match.Groups[3].Value);
+
+            // ... The top count should be equal to what we provided
+            int limitFromQuery;
+            Assert.True(int.TryParse(match.Groups[1].Value, out limitFromQuery));
+            Assert.Equal(limit, limitFromQuery);
+        }
+
+        #endregion
+
         private static EditSession.Connector DoNothingConnector
         {
             get { return () => Task.FromResult<DbConnection>(null); }