diff --git a/     b/    
new file mode 100755
index 00000000..e69de29b
diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
index bc685eff..6d7fdb3e 100644
--- 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)
             {
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
index acdd70a9..c81362dd 100644
--- 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);
                     }
                 }
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/BatchSummary.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/BatchSummary.cs
index 73d1d4c8..4205fecb 100644
--- 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>
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/QueryExecuteRequest.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/QueryExecuteRequest.cs
index cac98c1a..6079bf51 100644
--- 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
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
index c4f49b2a..2c1c2ce2 100644
--- 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();
             }
         }
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
index 562a3a61..20bc4347 100644
--- 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
diff --git a/src/Microsoft.SqlTools.ServiceLayer/Workspace/Contracts/ScriptFile.cs b/src/Microsoft.SqlTools.ServiceLayer/Workspace/Contracts/ScriptFile.cs
index 43f064e6..4eaff942 100644
--- 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);
diff --git a/src/Microsoft.SqlTools.ServiceLayer/Workspace/Workspace.cs b/src/Microsoft.SqlTools.ServiceLayer/Workspace/Workspace.cs
index 80fe6a9f..1c798a5d 100644
--- 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);
 
diff --git a/src/Microsoft.SqlTools.ServiceLayer/Workspace/WorkspaceService.cs b/src/Microsoft.SqlTools.ServiceLayer/Workspace/WorkspaceService.cs
index c1fe5080..40dc1574 100644
--- 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>
diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/CancelTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/CancelTests.cs
index 7e0e5a4d..087a87cb 100644
--- 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);
diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Common.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Common.cs
index 56273be4..4aed41d8 100644
--- 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
diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DisposeTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/DisposeTests.cs
index 8c79296d..226afcc8 100644
--- 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);
diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/ExecuteTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/ExecuteTests.cs
index c81f5ee6..b44b3618 100644
--- 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;
diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResultsTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SaveResultsTests.cs
index 99e77f5d..96e5742b 100644
--- 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();
 
diff --git a/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SubsetTests.cs b/test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SubsetTests.cs
index 8212eba3..8fcc9386 100644
--- 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();