Feature/batch line info (#56)

* inital pipe of line numbers and getting text from workspace services * tests compile * Fixed bug regarding tests using connections on mac * updated tests * fixed workspace service and fixed tests * integrated feedback
2026-01-17 01:25:40 -05:00 · 2016-09-22 17:58:45 -07:00
parent 93a75f1ff4
commit f22c8a7283
16 changed files with 321 additions and 93 deletions
--- a/0
+++ b/0
--- a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
@@ -197,20 +197,24 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
            // try to get information about the connected SQL Server instance
            try
            {
-                ReliableConnectionHelper.ServerInfo serverInfo = ReliableConnectionHelper.GetServerVersion(connectionInfo.SqlConnection);
-                response.ServerInfo = new Contracts.ServerInfo()
+                var connection = connectionInfo.SqlConnection as ReliableSqlConnection;
+                if (connection != null)
                {
-                    ServerMajorVersion = serverInfo.ServerMajorVersion,
-                    ServerMinorVersion = serverInfo.ServerMinorVersion,
-                    ServerReleaseVersion = serverInfo.ServerReleaseVersion,
-                    EngineEditionId = serverInfo.EngineEditionId,
-                    ServerVersion = serverInfo.ServerVersion,
-                    ServerLevel = serverInfo.ServerLevel,
-                    ServerEdition = serverInfo.ServerEdition,
-                    IsCloud = serverInfo.IsCloud,
-                    AzureVersion = serverInfo.AzureVersion,
-                    OsVersion = serverInfo.OsVersion
-                };
+                    ReliableConnectionHelper.ServerInfo serverInfo = ReliableConnectionHelper.GetServerVersion(connection.GetUnderlyingConnection());
+                    response.ServerInfo = new Contracts.ServerInfo()
+                    {
+                        ServerMajorVersion = serverInfo.ServerMajorVersion,
+                        ServerMinorVersion = serverInfo.ServerMinorVersion,
+                        ServerReleaseVersion = serverInfo.ServerReleaseVersion,
+                        EngineEditionId = serverInfo.EngineEditionId,
+                        ServerVersion = serverInfo.ServerVersion,
+                        ServerLevel = serverInfo.ServerLevel,
+                        ServerEdition = serverInfo.ServerEdition,
+                        IsCloud = serverInfo.IsCloud,
+                        AzureVersion = serverInfo.AzureVersion,
+                        OsVersion = serverInfo.OsVersion
+                    };
+                }
            }
            catch(Exception ex)
            {
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
@@ -46,7 +46,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution

        #endregion

-        internal Batch(string batchText, int startLine, IFileStreamFactory outputFileFactory)
+        internal Batch(string batchText, int startLine, int startColumn, int endLine, int endColumn, IFileStreamFactory outputFileFactory)
        {
            // Sanity check for input
            Validate.IsNotNullOrEmptyString(nameof(batchText), batchText);
@@ -54,7 +54,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution

            // Initialize the internal state
            BatchText = batchText;
-            StartLine = startLine - 1;  // -1 to make sure that the line number of the batch is 0-indexed, since SqlParser gives 1-indexed line numbers
+            Selection = new SelectionData(startLine, startColumn, endLine, endColumn);
            HasExecuted = false;
            resultSets = new List<ResultSet>();
            resultMessages = new List<string>();
@@ -111,9 +111,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
        }

        /// <summary>
-        /// The 0-indexed line number that this batch started on
+        /// The range from the file that is this batch
        /// </summary>
-        internal int StartLine { get; set; }
+        internal SelectionData Selection { get; set; }

        #endregion

@@ -253,9 +253,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                    SqlError sqlError = error as SqlError;
                    if (sqlError != null)
                    {
-                        int lineNumber = sqlError.LineNumber + StartLine;
-                        string message = SR.QueryServiceErrorFormat(sqlError.Number, sqlError.Class, sqlError.State,
-                            lineNumber, Environment.NewLine, sqlError.Message);
+                        int lineNumber = sqlError.LineNumber + Selection.StartLine;
+                        string message = string.Format("Msg {0}, Level {1}, State {2}, Line {3}{4}{5}",
+                            sqlError.Number, sqlError.Class, sqlError.State, lineNumber,
+                            Environment.NewLine, sqlError.Message);
                        resultMessages.Add(message);
                    }
                }
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/BatchSummary.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/BatchSummary.cs
@@ -20,6 +20,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
        /// </summary>
        public int Id { get; set; }

+        /// <summary>
+        /// The selection from the file for this batch
+        /// </summary>
+        public SelectionData Selection { get; set; }
+
        /// <summary>
        /// Any messages that came back from the server during execution of the batch
        /// </summary>
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/QueryExecuteRequest.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/QueryExecuteRequest.cs
@@ -7,15 +7,30 @@ using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;

 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
 {
+    /// <summary>
+    /// Container class for a selection range from file
+    /// </summary>
+    public class SelectionData {
+        public int StartLine { get; set; }
+        public int StartColumn { get; set; }
+        public int EndLine { get; set; }
+        public int EndColumn { get; set; }
+        public SelectionData(int startLine, int startColumn, int endLine, int endColumn) {
+            StartLine = startLine;
+            StartColumn = startColumn;
+            EndLine = endLine;
+            EndColumn = endColumn;
+        }
+    }
    /// <summary>
    /// Parameters for the query execute request
    /// </summary>
    public class QueryExecuteParams
    {
        /// <summary>
-        /// The text of the query to execute
+        /// The selection from the document
        /// </summary>
-        public string QueryText { get; set; }
+        public SelectionData QuerySelection { get; set; }

        /// <summary>
        /// URI for the editor that is asking for the query execute
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
@@ -86,7 +86,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
            });
            // NOTE: We only want to process batches that have statements (ie, ignore comments and empty lines)
            Batches = parseResult.Script.Batches.Where(b => b.Statements.Count > 0)
-                .Select(b => new Batch(b.Sql, b.StartLocation.LineNumber, outputFileFactory)).ToArray();
+                .Select(b => new Batch(b.Sql, 
+                                       b.StartLocation.LineNumber - 1, 
+                                       b.StartLocation.ColumnNumber - 1, 
+                                       b.EndLocation.LineNumber - 1, 
+                                       b.EndLocation.ColumnNumber - 1, 
+                                       outputFileFactory)).ToArray();
        }

        #region Properties
@@ -113,7 +118,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                    Id = index,
                    HasError = batch.HasError,
                    Messages = batch.ResultMessages.ToArray(),
-                    ResultSetSummaries = batch.ResultSummaries
+                    ResultSetSummaries = batch.ResultSummaries,
+                    Selection = batch.Selection
                }).ToArray();
            }
        }
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
@@ -14,6 +14,7 @@ using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage;
 using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Workspace;
+using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 using Newtonsoft.Json;

 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
@@ -38,11 +39,13 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
        private QueryExecutionService()
        {
            ConnectionService = ConnectionService.Instance;
+            WorkspaceService = WorkspaceService<SqlToolsSettings>.Instance;
        }

-        internal QueryExecutionService(ConnectionService connService)
+        internal QueryExecutionService(ConnectionService connService, WorkspaceService<SqlToolsSettings> workspaceService)
        {
            ConnectionService = connService;
+            WorkspaceService = workspaceService;
        }

        #endregion
@@ -78,6 +81,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
        /// </summary>
        private ConnectionService ConnectionService { get; set; }

+        private WorkspaceService<SqlToolsSettings> WorkspaceService { get; set; }
+
        /// <summary>
        /// Internal storage of active queries, lazily constructed as a threadsafe dictionary
        /// </summary>
@@ -111,7 +116,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
            });

            // Register a handler for when the configuration changes
-            WorkspaceService<SqlToolsSettings>.Instance.RegisterConfigChangeCallback((oldSettings, newSettings, eventContext) =>
+            WorkspaceService.RegisterConfigChangeCallback((oldSettings, newSettings, eventContext) =>
            {
                Settings.QueryExecutionSettings.Update(newSettings.QueryExecutionSettings);
                return Task.FromResult(0);
@@ -403,10 +408,36 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                }

                // Retrieve the current settings for executing the query with
-                QueryExecutionSettings settings = WorkspaceService<SqlToolsSettings>.Instance.CurrentSettings.QueryExecutionSettings;
+                QueryExecutionSettings settings = WorkspaceService.CurrentSettings.QueryExecutionSettings;

+                // Get query text from the workspace.
+                ScriptFile queryFile = WorkspaceService.Workspace.GetFile(executeParams.OwnerUri);
+
+                string queryText;
+
+                if (executeParams.QuerySelection != null) 
+                {
+                    string[] queryTextArray = queryFile.GetLinesInRange(
+                        new BufferRange(
+                            new BufferPosition(
+                                executeParams.QuerySelection.StartLine + 1, 
+                                executeParams.QuerySelection.StartColumn + 1
+                            ), 
+                            new BufferPosition(
+                                executeParams.QuerySelection.EndLine + 1, 
+                                executeParams.QuerySelection.EndColumn + 1
+                            )
+                        )
+                    );
+                    queryText = queryTextArray.Aggregate((a, b) => a + '\r' + '\n' + b);
+                } 
+                else 
+                {
+                    queryText = queryFile.Contents;
+                }
+                
                // If we can't add the query now, it's assumed the query is in progress
-                Query newQuery = new Query(executeParams.QueryText, connectionInfo, settings, BufferFileFactory);
+                Query newQuery = new Query(queryText, connectionInfo, settings, BufferFileFactory);
                if (!ActiveQueries.TryAdd(executeParams.OwnerUri, newQuery))
                {
                    await requestContext.SendResult(new QueryExecuteResult
--- a/src/Microsoft.SqlTools.ServiceLayer/Workspace/Contracts/ScriptFile.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/Workspace/Contracts/ScriptFile.cs
@@ -56,7 +56,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Workspace.Contracts
        /// Gets or sets a string containing the full contents of the file.
        /// Setter for testing purposes only
        /// </summary>
-        public string Contents 
+        public virtual string Contents 
        {
            get
            {
@@ -109,7 +109,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Workspace.Contracts
        /// <summary>
        /// Add a default constructor for testing
        /// </summary>
-        internal ScriptFile()
+        public ScriptFile()
        {
            ClientFilePath = "test.sql";
        }
@@ -171,11 +171,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Workspace.Contracts
        }

        /// <summary>
-        /// Gets a range of lines from the file's contents.
+        /// Gets a range of lines from the file's contents. Virtual method to allow for
+        /// mocking.
        /// </summary>
        /// <param name="bufferRange">The buffer range from which lines will be extracted.</param>
        /// <returns>An array of strings from the specified range of the file.</returns>
-        public string[] GetLinesInRange(BufferRange bufferRange)
+        public virtual string[] GetLinesInRange(BufferRange bufferRange)
        {
            this.ValidatePosition(bufferRange.Start);
            this.ValidatePosition(bufferRange.End);
--- a/src/Microsoft.SqlTools.ServiceLayer/Workspace/Workspace.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/Workspace/Workspace.cs
@@ -50,7 +50,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Workspace

        /// <summary>
        /// Gets an open file in the workspace.  If the file isn't open but
-        /// exists on the filesystem, load and return it.
+        /// exists on the filesystem, load and return it. Virtual method to
+        /// allow for mocking
        /// </summary>
        /// <param name="filePath">The file path at which the script resides.</param>
        /// <exception cref="FileNotFoundException">
@@ -59,7 +60,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Workspace
        /// <exception cref="ArgumentException">
        /// <paramref name="filePath"/> contains a null or empty string.
        /// </exception>
-        public ScriptFile GetFile(string filePath)
+        public virtual ScriptFile GetFile(string filePath)
        {
            Validate.IsNotNullOrWhitespaceString("filePath", filePath);

--- a/src/Microsoft.SqlTools.ServiceLayer/Workspace/WorkspaceService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/Workspace/WorkspaceService.cs
@@ -28,7 +28,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Workspace

        #region Singleton Instance Implementation

-        private static readonly Lazy<WorkspaceService<TConfig>> instance = new Lazy<WorkspaceService<TConfig>>(() => new WorkspaceService<TConfig>());
+        private static Lazy<WorkspaceService<TConfig>> instance = new Lazy<WorkspaceService<TConfig>>(() => new WorkspaceService<TConfig>());

        public static WorkspaceService<TConfig> Instance
        {
@@ -52,8 +52,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Workspace

        #region Properties

-        public Workspace Workspace { get; private set; }
+        /// <summary>
+        /// Workspace object for the service. Virtual to allow for mocking
+        /// </summary>
+        public virtual Workspace Workspace { get; private set; }

+        /// <summary>
+        /// Current settings for the workspace
+        /// </summary>
        public TConfig CurrentSettings { get; internal set; }

        /// <summary>
--- a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/CancelTests.cs
+++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/CancelTests.cs
@@ -7,7 +7,10 @@ using System;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Test.Utility;
+using Microsoft.SqlTools.ServiceLayer.Workspace;
+using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 using Moq;
 using Xunit;

@@ -18,10 +21,19 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void CancelInProgressQueryTest()
        {
+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.Setup(file => file.GetLinesInRange(It.IsAny<BufferRange>()))
+                .Returns(new string[] { Common.StandardQuery });
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
+
            // If:
            // ... I request a query (doesn't matter what kind) and execute it
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
-            var executeParams = new QueryExecuteParams { QueryText = Common.StandardQuery, OwnerUri = Common.OwnerUri };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+            var executeParams = new QueryExecuteParams { QuerySelection = Common.GetSubSectionDocument(), OwnerUri = Common.OwnerUri };
            var executeRequest = 
                RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(null, QueryExecuteCompleteEvent.Type, null, null);
            queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();
@@ -45,10 +57,18 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void CancelExecutedQueryTest()
        {
+            
+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
            // If:
            // ... I request a query (doesn't matter what kind) and wait for execution
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
-            var executeParams = new QueryExecuteParams {QueryText = Common.StandardQuery, OwnerUri = Common.OwnerUri};
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+            var executeParams = new QueryExecuteParams {QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri};
            var executeRequest =
                RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(null, QueryExecuteCompleteEvent.Type, null, null);
            queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();
@@ -71,9 +91,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void CancelNonExistantTest()
        {
+
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
            // If:
            // ... I request to cancel a query that doesn't exist
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), false);
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), false, workspaceService.Object);
            var cancelParams = new QueryCancelParams {OwnerUri = "Doesn't Exist"};
            QueryCancelResult result = null;
            var cancelRequest = GetQueryCancelResultContextMock(qcr => result = qcr, null);
--- a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Common.cs
+++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Common.cs
@@ -18,9 +18,11 @@ using Microsoft.SqlServer.Management.SqlParser.Binder;
 using Microsoft.SqlServer.Management.SqlParser.MetadataProvider;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage;
 using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Test.Utility;
+using Microsoft.SqlTools.ServiceLayer.Workspace;
 using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 using Moq;
 using Moq.Protected;
@@ -29,6 +31,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
 {
    public class Common
    {
+        public const SelectionData WholeDocument = null;
+        
        public const string StandardQuery = "SELECT * FROM sys.objects";

        public const string InvalidQuery = "SELECT *** FROM sys.objects";
@@ -72,9 +76,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
            return output;
        }

+        public static SelectionData GetSubSectionDocument() 
+        {
+            return new SelectionData(0, 0, 2, 2);
+        }
+
        public static Batch GetBasicExecutedBatch()
        {
-            Batch batch = new Batch(StandardQuery, 1, GetFileStreamFactory());
+            Batch batch = new Batch(StandardQuery, 0, 0, 2, 2, GetFileStreamFactory());
            batch.Execute(CreateTestConnection(new[] {StandardTestData}, false), CancellationToken.None).Wait();
            return batch;
        }
@@ -272,7 +281,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
            };
        }

-        public static QueryExecutionService GetPrimedExecutionService(ISqlConnectionFactory factory, bool isConnected)
+        public static QueryExecutionService GetPrimedExecutionService(ISqlConnectionFactory factory, bool isConnected, WorkspaceService<SqlToolsSettings> workspaceService)
        {
            var connectionService = new ConnectionService(factory);
            if (isConnected)
@@ -283,7 +292,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
                    OwnerUri = OwnerUri
                });
            }
-            return new QueryExecutionService(connectionService) {BufferFileStreamFactory = GetFileStreamFactory()};
+            return new QueryExecutionService(connectionService, workspaceService) {BufferFileStreamFactory = GetFileStreamFactory()};
        }

        #endregion
--- a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DisposeTests.cs
+++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DisposeTests.cs
@@ -7,7 +7,10 @@ using System;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Test.Utility;
+using Microsoft.SqlTools.ServiceLayer.Workspace;
+using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 using Moq;
 using Xunit;

@@ -18,17 +21,26 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void DisposeExecutedQuery()
        {
+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns("doesn't matter");
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
            // If:
            // ... I request a query (doesn't matter what kind)
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
-            var executeParams = new QueryExecuteParams {QueryText = "Doesn'tMatter", OwnerUri = Common.OwnerUri};
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+            var executeParams = new QueryExecuteParams {QuerySelection = null, OwnerUri = Common.OwnerUri};
            var executeRequest = RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(null, QueryExecuteCompleteEvent.Type, null, null);
            queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();

            // ... And then I dispose of the query
            var disposeParams = new QueryDisposeParams {OwnerUri = Common.OwnerUri};
            QueryDisposeResult result = null;
-            var disposeRequest = GetQueryDisposeResultContextMock(qdr => result = qdr, null);
+            var disposeRequest = GetQueryDisposeResultContextMock(qdr => {
+                result = qdr;
+            }, null);
            queryService.HandleDisposeRequest(disposeParams, disposeRequest.Object).Wait();

            // Then:
@@ -42,9 +54,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void QueryDisposeMissingQuery()
        {
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
            // If:
            // ... I attempt to dispose a query that doesn't exist
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), false);
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), false, workspaceService.Object);
            var disposeParams = new QueryDisposeParams {OwnerUri = Common.OwnerUri};
            QueryDisposeResult result = null;
            var disposeRequest = GetQueryDisposeResultContextMock(qdr => result = qdr, null);
--- a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/ExecuteTests.cs
+++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/ExecuteTests.cs
@@ -16,6 +16,7 @@ using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Test.Utility;
 using Microsoft.SqlTools.ServiceLayer.Workspace;
+using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 using Moq;
 using Xunit;

@@ -29,7 +30,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        public void BatchCreationTest()
        {
            // If I create a new batch...
-            Batch batch = new Batch(Common.StandardQuery, 1, Common.GetFileStreamFactory());
+            Batch batch = new Batch(Common.StandardQuery, 0, 0, 2, 2, Common.GetFileStreamFactory());

            // Then: 
            // ... The text of the batch should be stored
@@ -45,14 +46,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
            Assert.Empty(batch.ResultMessages);

            // ... The start line of the batch should be 0
-            Assert.Equal(0, batch.StartLine);
+            Assert.Equal(0, batch.Selection.StartLine);
        }

        [Fact]
        public void BatchExecuteNoResultSets()
        {
            // If I execute a query that should get no result sets
-            Batch batch = new Batch(Common.StandardQuery, 1, Common.GetFileStreamFactory());
+            Batch batch = new Batch(Common.StandardQuery, 0, 0, 2, 2, Common.GetFileStreamFactory());
            batch.Execute(GetConnection(Common.CreateTestConnectionInfo(null, false)), CancellationToken.None).Wait();

            // Then:
@@ -82,7 +83,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
            ConnectionInfo ci = Common.CreateTestConnectionInfo(new[] { Common.StandardTestData }, false);

            // If I execute a query that should get one result set
-            Batch batch = new Batch(Common.StandardQuery, 1, Common.GetFileStreamFactory());
+            Batch batch = new Batch(Common.StandardQuery, 0, 0, 2, 2, Common.GetFileStreamFactory());
            batch.Execute(GetConnection(ci), CancellationToken.None).Wait();

            // Then:
@@ -115,7 +116,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
            ConnectionInfo ci = Common.CreateTestConnectionInfo(dataset, false);

            // If I execute a query that should get two result sets
-            Batch batch = new Batch(Common.StandardQuery, 1, Common.GetFileStreamFactory());
+            Batch batch = new Batch(Common.StandardQuery, 0, 0, 1, 1, Common.GetFileStreamFactory());
            batch.Execute(GetConnection(ci), CancellationToken.None).Wait();

            // Then:
@@ -161,7 +162,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
            ConnectionInfo ci = Common.CreateTestConnectionInfo(null, true);

            // If I execute a batch that is invalid
-            Batch batch = new Batch(Common.StandardQuery, 1, Common.GetFileStreamFactory());
+            Batch batch = new Batch(Common.StandardQuery, 0, 0, 2, 2, Common.GetFileStreamFactory());
            batch.Execute(GetConnection(ci), CancellationToken.None).Wait();

            // Then:
@@ -183,7 +184,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
            ConnectionInfo ci = Common.CreateTestConnectionInfo(new[] { Common.StandardTestData }, false);

            // If I execute a batch
-            Batch batch = new Batch(Common.StandardQuery, 1, Common.GetFileStreamFactory());
+            Batch batch = new Batch(Common.StandardQuery, 0, 0, 2, 2, Common.GetFileStreamFactory());
            batch.Execute(GetConnection(ci), CancellationToken.None).Wait();

            // Then:
@@ -213,7 +214,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
            // ... I create a batch that has an empty query
            // Then:
            // ... It should throw an exception
-            Assert.Throws<ArgumentException>(() => new Batch(query, 1, Common.GetFileStreamFactory()));
+            Assert.Throws<ArgumentException>(() => new Batch(query, 0, 0, 2, 2, Common.GetFileStreamFactory()));
        }

        [Fact]
@@ -223,7 +224,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
            // ... I create a batch that has no file stream factory
            // Then:
            // ... It should throw an exception
-            Assert.Throws<ArgumentNullException>(() => new Batch("stuff", 1, null));
+            Assert.Throws<ArgumentNullException>(() => new Batch("stuff", 0, 0, 2, 2, null));
        }

        #endregion
@@ -420,10 +421,17 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
            // ... Default settings are stored in the workspace service
            WorkspaceService<SqlToolsSettings>.Instance.CurrentSettings = new SqlToolsSettings();

+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
            // If:
            // ... I request to execute a valid query with no results
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
-            var queryParams = new QueryExecuteParams { QueryText = Common.StandardQuery, OwnerUri = Common.OwnerUri };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+            var queryParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri };

            QueryExecuteResult result = null;
            QueryExecuteCompleteParams completeParams = null;
@@ -452,10 +460,19 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void QueryExecuteValidResultsTest()
        {
+            
+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
            // If:
            // ... I request to execute a valid query with results
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(new[] { Common.StandardTestData }, false), true);
-            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QueryText = Common.StandardQuery };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(new[] { Common.StandardTestData }, false), true,
+                workspaceService.Object);
+            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument };

            QueryExecuteResult result = null;
            QueryExecuteCompleteParams completeParams = null;
@@ -485,10 +502,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void QueryExecuteUnconnectedUriTest()
        {
+
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
            // If:
            // ... I request to execute a query using a file URI that isn't connected
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), false);
-            var queryParams = new QueryExecuteParams { OwnerUri = "notConnected", QueryText = Common.StandardQuery };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), false, workspaceService.Object);
+            var queryParams = new QueryExecuteParams { OwnerUri = "notConnected", QuerySelection = Common.WholeDocument };

            QueryExecuteResult result = null;
            var requestContext = RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(qer => result = qer, QueryExecuteCompleteEvent.Type, null, null);
@@ -508,10 +527,19 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void QueryExecuteInProgressTest()
        {
+
+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
+
            // If:
            // ... I request to execute a query
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
-            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QueryText = Common.StandardQuery };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument };

            // Note, we don't care about the results of the first request
            var firstRequestContext = RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(null, QueryExecuteCompleteEvent.Type, null, null);
@@ -537,10 +565,19 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void QueryExecuteCompletedTest()
        {
+            
+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
+                
            // If:
            // ... I request to execute a query
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
-            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QueryText = Common.StandardQuery };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument };

            // Note, we don't care about the results of the first request
            var firstRequestContext = RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(null, QueryExecuteCompleteEvent.Type, null, null);
@@ -565,14 +602,21 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        }

        [Theory]
-        [InlineData("")]
        [InlineData(null)]
-        public void QueryExecuteMissingQueryTest(string query)
+        public void QueryExecuteMissingSelectionTest(SelectionData selection)
        {
+
+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns("");
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
            // If:
            // ... I request to execute a query with a missing query string
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
-            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QueryText = query };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = selection };

            QueryExecuteResult result = null;
            var requestContext =
@@ -594,10 +638,17 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void QueryExecuteInvalidQueryTest()
        {
+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
            // If:
            // ... I request to execute a query that is invalid
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, true), true);
-            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QueryText = Common.StandardQuery };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, true), true, workspaceService.Object);
+            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument };

            QueryExecuteResult result = null;
            QueryExecuteCompleteParams complete = null;
--- a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResultsTests.cs
+++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResultsTests.cs
@@ -9,6 +9,9 @@ using System.Runtime.InteropServices;
 using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+using Microsoft.SqlTools.ServiceLayer.SqlContext;
+using Microsoft.SqlTools.ServiceLayer.Workspace;
+using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 using Moq;
 using Xunit;

@@ -25,9 +28,17 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void SaveResultsAsCsvSuccessTest()
        {
+
+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
            // Execute a query
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
-            var executeParams = new QueryExecuteParams { QueryText = Common.StandardQuery, OwnerUri = Common.OwnerUri };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+            var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri };
            var executeRequest = GetQueryExecuteResultContextMock(null, null, null);
            queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();

@@ -63,9 +74,18 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void SaveResultsAsCsvExceptionTest()
        {
+
+             // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
+                
            // Execute a query
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
-            var executeParams = new QueryExecuteParams { QueryText = Common.StandardQuery, OwnerUri = Common.OwnerUri };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+            var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri };
            var executeRequest = GetQueryExecuteResultContextMock(null, null, null);
            queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();

@@ -95,9 +115,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void SaveResultsAsCsvQueryNotFoundTest()
        {
+
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
            // Execute a query
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
-            var executeParams = new QueryExecuteParams { QueryText = Common.StandardQuery, OwnerUri = Common.OwnerUri };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+            var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri };
            var executeRequest = GetQueryExecuteResultContextMock(null, null, null);
            queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();

@@ -125,9 +147,17 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void SaveResultsAsJsonSuccessTest()
        {
+
+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
            // Execute a query
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
-            var executeParams = new QueryExecuteParams { QueryText = Common.StandardQuery, OwnerUri = Common.OwnerUri };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+            var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri };
            var executeRequest = GetQueryExecuteResultContextMock(null, null, null);
            queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();

@@ -162,9 +192,16 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void SaveResultsAsJsonExceptionTest()
        {
+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
            // Execute a query
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
-            var executeParams = new QueryExecuteParams { QueryText = Common.StandardQuery, OwnerUri = Common.OwnerUri };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+            var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri };
            var executeRequest = GetQueryExecuteResultContextMock(null, null, null);
            queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();

@@ -194,9 +231,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void SaveResultsAsJsonQueryNotFoundTest()
        {
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
            // Execute a query
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
-            var executeParams = new QueryExecuteParams { QueryText = Common.StandardQuery, OwnerUri = Common.OwnerUri };
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
+            var executeParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri };
            var executeRequest = GetQueryExecuteResultContextMock(null, null, null);
            queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();

--- a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SubsetTests.cs
+++ b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SubsetTests.cs
@@ -11,6 +11,8 @@ using Microsoft.SqlTools.ServiceLayer.QueryExecution;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Test.Utility;
+using Microsoft.SqlTools.ServiceLayer.Workspace;
+using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 using Moq;
 using Xunit;

@@ -130,11 +132,20 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public async Task SubsetServiceValidTest()
        {
+
+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
            // If:
            // ... I have a query that has results (doesn't matter what)
-            var queryService =Common.GetPrimedExecutionService(
-                Common.CreateMockFactory(new[] {Common.StandardTestData}, false), true);
-            var executeParams = new QueryExecuteParams {QueryText = Common.StandardQuery, OwnerUri = Common.OwnerUri};
+            var queryService = Common.GetPrimedExecutionService(
+                Common.CreateMockFactory(new[] {Common.StandardTestData}, false), true,
+                workspaceService.Object);
+            var executeParams = new QueryExecuteParams {QuerySelection = null, OwnerUri = Common.OwnerUri};
            var executeRequest = RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(null, QueryExecuteCompleteEvent.Type, null, null);
            await queryService.HandleExecuteRequest(executeParams, executeRequest.Object);

@@ -156,9 +167,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void SubsetServiceMissingQueryTest()
        {
+
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
            // If:
            // ... I ask for a set of results for a file that hasn't executed a query
-            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true);
+            var queryService = Common.GetPrimedExecutionService(Common.CreateMockFactory(null, false), true, workspaceService.Object);
            var subsetParams = new QueryExecuteSubsetParams { OwnerUri = Common.OwnerUri, RowsCount = 1, ResultSetIndex = 0, RowsStartIndex = 0 };
            QueryExecuteSubsetResult result = null;
            var subsetRequest = GetQuerySubsetResultContextMock(qesr => result = qesr, null);
@@ -176,11 +189,20 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void SubsetServiceUnexecutedQueryTest()
        {
+            
+            // Set up file for returning the query
+            var fileMock = new Mock<ScriptFile>();
+            fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+            // Set up workspace mock
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
+            workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
+                .Returns(fileMock.Object);
            // If:
            // ... I have a query that hasn't finished executing (doesn't matter what)
            var queryService = Common.GetPrimedExecutionService(
-                Common.CreateMockFactory(new[] { Common.StandardTestData }, false), true);
-            var executeParams = new QueryExecuteParams { QueryText = Common.StandardQuery, OwnerUri = Common.OwnerUri };
+                Common.CreateMockFactory(new[] { Common.StandardTestData }, false), true,
+                workspaceService.Object);
+            var executeParams = new QueryExecuteParams { QuerySelection = null, OwnerUri = Common.OwnerUri };
            var executeRequest = RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(null, QueryExecuteCompleteEvent.Type, null, null);
            queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();
            queryService.ActiveQueries[Common.OwnerUri].HasExecuted = false;
@@ -203,11 +225,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
        [Fact]
        public void SubsetServiceOutOfRangeSubsetTest()
        {
+
+            var workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
            // If:
            // ... I have a query that doesn't have any result sets
            var queryService = Common.GetPrimedExecutionService(
-                Common.CreateMockFactory(null, false), true);
-            var executeParams = new QueryExecuteParams { QueryText = Common.StandardQuery, OwnerUri = Common.OwnerUri };
+                Common.CreateMockFactory(null, false), true,
+                workspaceService.Object);
+            var executeParams = new QueryExecuteParams { QuerySelection = null, OwnerUri = Common.OwnerUri };
            var executeRequest = RequestContextMocks.SetupRequestContextMock<QueryExecuteResult, QueryExecuteCompleteParams>(null, QueryExecuteCompleteEvent.Type, null, null);
            queryService.HandleExecuteRequest(executeParams, executeRequest.Object).Wait();