Feature execution plan settings and request implementation (#213)

* experimental showplan implementation (tools side only)

* fix for redundant massages from showplan executions

* moved showplan batches to new variables to make it less confusing

* returns showplan as part of batch summary with in each result summary

* cleaned up showplan resultsets

* cleaning up code and making showplan var optional

* changes all var names to showplan

* adding estimated support

* small fixes

* updatin var names and adding EPOptions struct

* adding ssms execution plan logic based on server types

* adding special actions logic

* removing redundant name changes

* execution plan query handler added

* cleaning up functions and data structures

* seperated special actions into its own class

* cleaning up special actions

* cleaning up code

* small new line fixes

* commenting out pre-yukon code

* removing all pre yukon code

* last yukon code commented out

* fixes broken tests

* adding related unit tests; integration tests incoming

* finishing tests and cleaning up code

* semantic changes

* cleaning up semantics

* changes and test fixes, also adding new exceptions into RS

* fixing special actions and cleaning up request logic

* fixing comment to trigger new build

* triggering another  build

* fixed up specialaction and added tests for it

src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionInfo.cs

@@ -60,5 +60,16 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// Returns true is the db connection is to a SQL db
         /// </summary>
         public bool IsAzure { get; set; }
+
+        /// <summary>
+        /// Returns true if the sql connection is to a DW instance
+        /// </summary>
+        public bool IsSqlDW { get; set; }
+
+        /// <summary>
+        /// Returns the major version number of the db we are connected to 
+        /// </summary>
+        public int MajorVersion { get; set; }
+
     }
 }

src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs

@@ -19,6 +19,7 @@ using Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts;
 using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Utility;
 using Microsoft.SqlTools.ServiceLayer.Workspace;
+using Microsoft.SqlServer.Management.Common;
 
 namespace Microsoft.SqlTools.ServiceLayer.Connection
 {
@@ -292,6 +293,8 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
                     OsVersion = serverInfo.OsVersion
                 };
                 connectionInfo.IsAzure = serverInfo.IsCloud;
+                connectionInfo.MajorVersion = serverInfo.ServerMajorVersion;
+                connectionInfo.IsSqlDW = (serverInfo.EngineEditionId == (int)DatabaseEngineEdition.SqlDataWarehouse);
             }
             catch(Exception ex)
             {

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs

@@ -55,6 +55,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         private readonly List<ResultSet> resultSets;
 
+        /// <summary>
+        /// Special action which this batch performed 
+        /// </summary>
+        private SpecialAction specialAction;
+
         #endregion
 
         internal Batch(string batchText, SelectionData selection, int ordinalId, IFileStreamFactory outputFileFactory)
@@ -72,6 +77,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             Id = ordinalId;
             resultSets = new List<ResultSet>();
             this.outputFileFactory = outputFileFactory;
+            specialAction = new SpecialAction();
         }
 
         #region Events
@@ -201,6 +207,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     summary.ResultSetSummaries = ResultSummaries;
                     summary.ExecutionEnd = ExecutionEndTimeStamp;
                     summary.ExecutionElapsed = ExecutionElapsedTime;
+                    summary.SpecialAction = ProcessResultSetSpecialActions();
                 }
 
                 return summary;
@@ -370,6 +377,30 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             return targetResultSet.GetSubset(startRow, rowCount);
         }
 
+        /// <summary>
+        /// Generates an execution plan
+        /// </summary>
+        /// <param name="resultSetIndex">The index for selecting the result set</param>
+        /// <returns>An exeuction plan object</returns>
+        public Task<ExecutionPlan> GetExecutionPlan(int resultSetIndex)
+        {
+            ResultSet targetResultSet;
+            lock (resultSets)
+            {
+                // Sanity check to make sure we have valid numbers
+                if (resultSetIndex < 0 || resultSetIndex >= resultSets.Count)
+                {
+                    throw new ArgumentOutOfRangeException(nameof(resultSetIndex),
+                        SR.QueryServiceSubsetResultSetOutOfRange);
+                }
+
+                targetResultSet = resultSets[resultSetIndex];
+            }
+
+            // Retrieve the result set
+            return targetResultSet.GetExecutionPlan();
+        }
+
         /// <summary>
         /// Saves a result to a file format selected by the user
         /// </summary>
@@ -482,6 +513,19 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             }
         }
 
+        /// <summary>
+        /// Aggregates all result sets in the batch into a single special action 
+        /// </summary>
+        private SpecialAction ProcessResultSetSpecialActions()
+        {
+            foreach (ResultSet resultSet in resultSets) 
+            {
+                specialAction.CombineSpecialAction(resultSet.Summary.SpecialAction);
+            }
+
+            return specialAction;
+        }
+
         #endregion
 
         #region IDisposable Implementation

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/BatchSummary.cs

@@ -44,5 +44,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
         /// The summaries of the result sets inside the batch
         /// </summary>
         public ResultSetSummary[] ResultSetSummaries { get; set; }
+
+        /// <summary>
+        /// The special action of the batch 
+        /// </summary>
+        public SpecialAction SpecialAction { get; set; }
     }
 }

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/ExecutionPlan.cs

@@ -0,0 +1,23 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    /// <summary>
+    /// Class used to represent an execution plan from a query for transmission across JSON RPC
+    /// </summary>
+    public class ExecutionPlan
+    {
+        /// <summary>
+        /// The format of the execution plan 
+        /// </summary>
+        public string Format { get; set; }
+
+        /// <summary>
+        /// The execution plan content
+        /// </summary>
+        public string Content { get; set; }
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/ExecutionPlanOptions.cs

@@ -0,0 +1,23 @@
+// 
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    /// <summary> 
+    /// Incoming execution plan options from the extension
+    /// </summary>
+    public struct ExecutionPlanOptions
+    {
+
+        /// <summary>
+        /// Setting to return the actual execution plan as XML
+        /// </summary>
+        public bool IncludeActualExecutionPlanXml { get; set; }
+
+        /// <summary>
+        /// Setting to return the estimated execution plan as XML
+        /// </summary>
+        public bool IncludeEstimatedExecutionPlanXml { get; set; }
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/QueryExecuteRequest.cs

@@ -21,6 +21,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
         /// URI for the editor that is asking for the query execute
         /// </summary>
         public string OwnerUri { get; set; }
+
+        /// <summary>
+        /// Execution plan options
+        /// </summary>
+        public ExecutionPlanOptions ExecutionPlanOptions { get; set; }
+
     }
 
     /// <summary>

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/QueryExecutionPlanRequest.cs

@@ -0,0 +1,49 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    /// <summary>
+    /// Parameters for query execution plan request
+    /// </summary>
+    public class QueryExecutionPlanParams
+    {
+        /// <summary>
+        /// URI for the file that owns the query to look up the results for
+        /// </summary>
+        public string OwnerUri { get; set; }
+
+        /// <summary>
+        /// Index of the batch to get the results from
+        /// </summary>
+        public int BatchIndex { get; set; }
+
+        /// <summary>
+        /// Index of the result set to get the results from
+        /// </summary>
+        public int ResultSetIndex { get; set; }
+
+    }
+
+    /// <summary>
+    /// Parameters for the query execution plan request
+    /// </summary>
+    public class QueryExecutionPlanResult
+    {
+        /// <summary>
+        /// The requested execution plan. Optional, can be set to null to indicate an error
+        /// </summary>
+        public ExecutionPlan ExecutionPlan { get; set; }
+    }
+
+    public class QueryExecutionPlanRequest
+    {
+        public static readonly
+            RequestType<QueryExecutionPlanParams, QueryExecutionPlanResult> Type =
+            RequestType<QueryExecutionPlanParams, QueryExecutionPlanResult>.Create("query/executionPlan");
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/ResultSetSummary.cs

@@ -29,5 +29,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
         /// Details about the columns that are provided as solutions
         /// </summary>
         public DbColumnWrapper[] ColumnInfo { get; set; }
+
+        /// <summary>
+        /// The special action definition of the result set 
+        /// </summary>
+        public SpecialAction SpecialAction { get; set; }
+
     }
 }

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs

@@ -15,6 +15,7 @@ using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage;
 using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Utility;
+using System.Collections.Generic;
 
 namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 {
@@ -51,6 +52,36 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         private bool hasExecuteBeenCalled;
 
+        /// <summary>
+        /// Settings for query runtime 
+        /// </summary>
+        private QueryExecutionSettings querySettings;
+
+        /// <summary>
+        /// Streaming output factory for the query 
+        /// </summary>
+        private IFileStreamFactory streamOutputFactory;
+
+        /// <summary>
+        /// ON keyword
+        /// </summary>
+        private const string On = "ON";
+
+        /// <summary>
+        /// OFF keyword
+        /// </summary>
+        private const string Off = "OFF";
+
+        /// <summary>
+        /// showplan_xml statement
+        /// </summary>
+        private const string SetShowPlanXml = "SET SHOWPLAN_XML {0}";
+
+        /// <summary>
+        /// statistics xml statement
+        /// </summary>
+        private const string SetStatisticsXml = "SET STATISTICS XML {0}";
+
         #endregion
 
         /// <summary>
@@ -72,6 +103,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             QueryText = queryText;
             editorConnection = connection;
             cancellationSource = new CancellationTokenSource();
+            querySettings = settings;
+            streamOutputFactory = outputFactory;
 
             // Process the query into batches
             ParseResult parseResult = Parser.Parse(queryText, new ParseOptions
@@ -89,7 +122,28 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                             batch.EndLocation.LineNumber - 1,
                             batch.EndLocation.ColumnNumber - 1),
                         index, outputFactory));
+                        
             Batches = batchSelection.ToArray();
+
+            // Create our batch lists
+            BeforeBatches = new List<Batch>();
+            AfterBatches = new List<Batch>();
+
+            if (DoesSupportExecutionPlan(connection))
+            {
+                // Checking settings for execution plan options 
+                if (querySettings.ExecutionPlanOptions.IncludeEstimatedExecutionPlanXml)
+                {
+                    // Enable set showplan xml
+                    addBatch(string.Format(SetShowPlanXml, On), BeforeBatches, streamOutputFactory);
+                    addBatch(string.Format(SetShowPlanXml, Off), AfterBatches, streamOutputFactory);
+                }
+                else if (querySettings.ExecutionPlanOptions.IncludeActualExecutionPlanXml)
+                {
+                    addBatch(string.Format(SetStatisticsXml, On), BeforeBatches, streamOutputFactory);
+                    addBatch(string.Format(SetStatisticsXml, Off), AfterBatches, streamOutputFactory);
+                }
+            }
         }
 
         #region Events
@@ -145,11 +199,21 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// <param name="q">The query that completed</param>
         public delegate Task QueryAsyncEventHandler(Query q);
 
+        /// <summary>
+        /// The batches which should run before the user batches 
+        /// </summary>
+        internal List<Batch> BeforeBatches { get; set; }
+
         /// <summary>
         /// The batches underneath this query
         /// </summary>
         internal Batch[] Batches { get; set; }
 
+        /// <summary>
+        /// The batches which should run after the user batches 
+        /// </summary>
+        internal List<Batch> AfterBatches { get; set; }
+
         /// <summary>
         /// The summaries of the batches underneath this query
         /// </summary>
@@ -241,6 +305,23 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             return Batches[batchIndex].GetSubset(resultSetIndex, startRow, rowCount);
         }
 
+        /// <summary>
+        /// Retrieves a subset of the result sets
+        /// </summary>
+        /// <param name="batchIndex">The index for selecting the batch item</param>
+        /// <param name="resultSetIndex">The index for selecting the result set</param>
+        /// <returns>The Execution Plan, if the result set has one</returns>
+        public Task<ExecutionPlan> GetExecutionPlan(int batchIndex, int resultSetIndex)
+        {
+            // Sanity check to make sure that the batch is within bounds
+            if (batchIndex < 0 || batchIndex >= Batches.Length)
+            {
+                throw new ArgumentOutOfRangeException(nameof(batchIndex), SR.QueryServiceSubsetBatchOutOfRange);
+            }
+
+            return Batches[batchIndex].GetExecutionPlan(resultSetIndex);
+        }
+
         /// <summary>
         /// Saves the requested results to a file format of the user's choice
         /// </summary>
@@ -316,9 +397,16 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
                 try
                 {
+                    // Execute beforeBatches synchronously, before the user defined batches 
+                    foreach (Batch b in BeforeBatches)
+                    {
+                        await b.Execute(conn, cancellationSource.Token);
+                    }
+
                     // We need these to execute synchronously, otherwise the user will be very unhappy
                     foreach (Batch b in Batches)
                     {
+                        // Add completion callbacks 
                         b.BatchStart += BatchStarted;
                         b.BatchCompletion += BatchCompleted;
                         b.BatchMessageSent += BatchMessageSent;
@@ -326,6 +414,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                         await b.Execute(conn, cancellationSource.Token);
                     }
 
+                    // Execute afterBatches synchronously, after the user defined batches
+                    foreach (Batch b in AfterBatches)
+                    {
+                        await b.Execute(conn, cancellationSource.Token);
+                    }
+
                     // Call the query execution callback
                     if (QueryCompleted != null)
                     {
@@ -374,6 +468,14 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             }
         }
 
+        /// <summary>
+        /// Function to add a new batch to a Batch set
+        /// </summary>
+        private void addBatch(string query, List<Batch> batchSet, IFileStreamFactory outputFactory)
+        {
+            batchSet.Add(new Batch(query, null, batchSet.Count, outputFactory));
+        }
+
         #endregion
 
         #region IDisposable Implementation
@@ -403,6 +505,14 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             disposed = true;
         }
 
+        /// <summary>
+        /// Does this connection support XML Execution plans
+        /// </summary>
+        private bool DoesSupportExecutionPlan(ConnectionInfo connectionInfo) {
+            // Determining which execution plan options may be applied (may be added to for pre-yukon support)
+            return (!connectionInfo.IsSqlDW && connectionInfo.MajorVersion >= 9);
+        }
+
         #endregion
     }
 }

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs

@@ -126,6 +126,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             serviceHost.SetRequestHandler(QueryCancelRequest.Type, HandleCancelRequest);
             serviceHost.SetRequestHandler(SaveResultsAsCsvRequest.Type, HandleSaveResultsAsCsvRequest);
             serviceHost.SetRequestHandler(SaveResultsAsJsonRequest.Type, HandleSaveResultsAsJsonRequest);
+            serviceHost.SetRequestHandler(QueryExecutionPlanRequest.Type, HandleExecutionPlanRequest);
 
             // Register handler for shutdown event
             serviceHost.RegisterShutdownTask((shutdownParams, requestContext) =>
@@ -208,6 +209,36 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             }
         }
 
+         /// <summary>
+        /// Handles a request to get an execution plan
+        /// </summary>
+        public async Task HandleExecutionPlanRequest(QueryExecutionPlanParams planParams,
+            RequestContext<QueryExecutionPlanResult> requestContext)
+        {
+            try
+            {
+                // Attempt to load the query
+                Query query;
+                if (!ActiveQueries.TryGetValue(planParams.OwnerUri, out query))
+                {
+                    await requestContext.SendError(SR.QueryServiceRequestsNoQuery);
+                    return;
+                }
+
+                // Retrieve the requested execution plan and return it
+                var result = new QueryExecutionPlanResult
+                {
+                    ExecutionPlan = await query.GetExecutionPlan(planParams.BatchIndex, planParams.ResultSetIndex)
+                };
+                await requestContext.SendResult(result);
+            }
+            catch (Exception e)
+            {
+                // This was unexpected, so send back as error
+                await requestContext.SendError(e.Message);
+            }
+        }
+
         /// <summary>
         /// Handles a request to dispose of this query
         /// </summary>
@@ -334,6 +365,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 // Retrieve the current settings for executing the query with
                 QueryExecutionSettings settings = WorkspaceService.CurrentSettings.QueryExecutionSettings;
 
+                // Apply execution parameter settings 
+                settings.ExecutionPlanOptions = executeParams.ExecutionPlanOptions;
+
                 // Get query text from the workspace.
                 ScriptFile queryFile = WorkspaceService.Workspace.GetFile(executeParams.OwnerUri);
 
@@ -425,6 +459,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     BatchSummary = b.Summary,
                     OwnerUri = executeParams.OwnerUri
                 };
+
                 await requestContext.SendEvent(QueryExecuteBatchStartEvent.Type, eventParams);
             };
             query.BatchStarted += batchStartCallback;
@@ -436,6 +471,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     BatchSummary = b.Summary,
                     OwnerUri = executeParams.OwnerUri
                 };
+
                 await requestContext.SendEvent(QueryExecuteBatchCompleteEvent.Type, eventParams);
             };
             query.BatchCompleted += batchCompleteCallback;

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultSet.cs

@@ -27,6 +27,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         // Column names of 'for xml' and 'for json' queries
         private const string NameOfForXMLColumn = "XML_F52E2B61-18A1-11d1-B105-00805F49916B";
         private const string NameOfForJSONColumn = "JSON_F52E2B61-18A1-11d1-B105-00805F49916B";
+        private const string YukonXmlShowPlanColumn = "Microsoft SQL Server 2005 XML Showplan";
 
         #endregion
 
@@ -68,6 +69,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         private readonly string outputFileName;
 
+        /// <summary>
+        /// The special action which applied to this result set
+        /// </summary>
+        private SpecialAction specialAction;
+
         #endregion
 
         /// <summary>
@@ -89,6 +95,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             // Initialize the storage
             outputFileName = factory.CreateFile();
             fileOffsets = new LongList<long>();
+            specialAction = new SpecialAction();
 
             // Store the factory
             fileStreamFactory = factory;
@@ -165,7 +172,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     ColumnInfo = Columns,
                     Id = Id,
                     BatchId = BatchId,
-                    RowCount = RowCount
+                    RowCount = RowCount,
+                    SpecialAction = ProcessSpecialAction()
+                    
                 };
             }
         }
@@ -236,6 +245,51 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             });
         }
 
+        /// <summary>
+        /// Generates the execution plan from the table returned 
+        /// </summary>
+        /// <returns>An execution plan object</returns>
+        public Task<ExecutionPlan> GetExecutionPlan()
+        {
+            // Proccess the action just incase is hasn't been yet 
+            ProcessSpecialAction();
+
+            // Sanity check to make sure that the results have been read beforehand
+            if (!hasBeenRead)
+            {
+                throw new InvalidOperationException(SR.QueryServiceResultSetNotRead);
+            }
+            // Check that we this result set contains a showplan 
+            else if (!specialAction.ExpectYukonXMLShowPlan)
+            {
+                throw new Exception(SR.QueryServiceExecutionPlanNotFound);
+            }
+
+
+            return Task.Factory.StartNew(() =>
+            { 
+                string content = null;
+                string format = null;
+
+                using (IFileStreamReader fileStreamReader = fileStreamFactory.GetReader(outputFileName))
+                {
+                    // Determine the format and get the first col/row of XML
+                    content = fileStreamReader.ReadRow(0, Columns)[0].DisplayValue;
+
+                    if (specialAction.ExpectYukonXMLShowPlan) 
+                    {
+                        format = "xml";
+                    }
+                }
+                    
+                return new ExecutionPlan
+                {
+                    Format = format,
+                    Content = content
+                };
+            });
+        }
+
         /// <summary>
         /// Reads from the reader until there are no more results to read
         /// </summary>
@@ -436,6 +490,21 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             }
         }
 
+        /// <summary>
+        /// Determine the special action, if any, for this result set
+        /// </summary>
+        private SpecialAction ProcessSpecialAction() 
+        {           
+
+            // Check if this result set is a showplan 
+            if (dataReader.Columns.Length == 1 && string.Compare(dataReader.Columns[0].ColumnName, YukonXmlShowPlanColumn, StringComparison.OrdinalIgnoreCase) == 0)
+            {
+                specialAction.ExpectYukonXMLShowPlan = true;
+            }
+
+            return specialAction;
+        }
+
         #endregion
     }
 }

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/SpecialAction.cs

@@ -0,0 +1,80 @@
+// 
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+using System;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
+{
+    /// <summary>
+    /// Class that represents a Special Action which occured by user request during the query 
+    /// </summary>
+    public class SpecialAction {
+        
+        #region Private Class variables 
+
+        // Underlying representation as bitwise flags to simplify logic
+        [Flags]
+        private enum ActionFlags 
+        {
+            None    = 0,
+            // All added options must be powers of 2
+            ExpectYukonXmlShowPlan = 1
+        }
+
+        private ActionFlags flags; 
+
+        #endregion
+
+        /// <summary>
+        /// The type of XML execution plan that is contained with in a result set  
+        /// </summary>
+        public SpecialAction()
+        {
+            flags = ActionFlags.None;
+        }
+
+        #region Public Functions
+        /// <summary>
+        /// No Special action performed 
+        /// </summary>
+        public bool None 
+        { 
+            get { return flags == ActionFlags.None; } 
+            set 
+            {
+                flags = ActionFlags.None;
+            }
+        }
+
+        /// <summary>
+        /// Contains an XML execution plan result set  
+        /// </summary>
+        public bool ExpectYukonXMLShowPlan 
+        {
+            get { return flags.HasFlag(ActionFlags.ExpectYukonXmlShowPlan); }
+            set 
+            {
+                if (value)
+                {
+                    // OR flags with value to apply 
+                    flags |= ActionFlags.ExpectYukonXmlShowPlan;
+                } 
+                else
+                {
+                    // AND flags with the inverse of the value we want to remove
+                    flags &= ~(ActionFlags.ExpectYukonXmlShowPlan);
+                }
+            }
+        }
+
+        /// <summary>
+        /// Aggregate this special action with the input
+        /// </summary>
+        public void CombineSpecialAction(SpecialAction action)
+        {
+            flags |= action.flags;
+        }
+        
+        #endregion
+    };
+}

src/Microsoft.SqlTools.ServiceLayer/SqlContext/QueryExecutionSettings.cs

@@ -1,6 +1,9 @@
 // 
 // Copyright (c) Microsoft. All rights reserved.
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+
 namespace Microsoft.SqlTools.ServiceLayer.SqlContext
 {
     /// <summary>
@@ -26,6 +29,16 @@ namespace Microsoft.SqlTools.ServiceLayer.SqlContext
         /// </summary>
         private const int DefaultMaxXmlCharsToStore = 2097152; // 2 MB - QE default
 
+        /// <summary>
+        /// Default selection of returning an actual XML showplan with all batches
+        /// Do not return any execution plan by default
+        /// </summary>
+        private ExecutionPlanOptions DefaultExecutionPlanOptions = new ExecutionPlanOptions() 
+        { 
+            IncludeActualExecutionPlanXml = false,
+            IncludeEstimatedExecutionPlanXml = false
+        };
+
         #endregion
 
         #region Member Variables
@@ -36,6 +49,8 @@ namespace Microsoft.SqlTools.ServiceLayer.SqlContext
 
         private int? maxXmlCharsToStore;
 
+        private ExecutionPlanOptions? executionPlanOptions;
+
         #endregion
 
         #region Properties
@@ -61,6 +76,12 @@ namespace Microsoft.SqlTools.ServiceLayer.SqlContext
             set { maxXmlCharsToStore = value; }
         }
 
+        public ExecutionPlanOptions ExecutionPlanOptions
+        {
+            get { return executionPlanOptions ?? DefaultExecutionPlanOptions; }
+            set { executionPlanOptions = value; }
+        }
+
         #endregion
 
         #region Public Methods
@@ -74,6 +95,7 @@ namespace Microsoft.SqlTools.ServiceLayer.SqlContext
             BatchSeparator = newSettings.BatchSeparator;
             MaxCharsToStore = newSettings.MaxCharsToStore;
             MaxXmlCharsToStore = newSettings.MaxXmlCharsToStore;
+            ExecutionPlanOptions = newSettings.ExecutionPlanOptions;
         }
 
         #endregion

src/Microsoft.SqlTools.ServiceLayer/sr.cs

@@ -365,6 +365,14 @@ namespace Microsoft.SqlTools.ServiceLayer
             }
         }
 
+        public static string QueryServiceExecutionPlanNotFound
+        {
+            get
+            {
+                return Keys.GetString(Keys.QueryServiceExecutionPlanNotFound);
+            }
+        }
+
         public static string PeekDefinitionNoResultsError
         {
             get
@@ -623,6 +631,9 @@ namespace Microsoft.SqlTools.ServiceLayer
             public const string QueryServiceResultSetNoColumnSchema = "QueryServiceResultSetNoColumnSchema";
 
 
+            public const string QueryServiceExecutionPlanNotFound = "QueryServiceExecutionPlanNotFound";
+
+
             public const string PeekDefinitionAzureError = "PeekDefinitionAzureError";
 
 

src/Microsoft.SqlTools.ServiceLayer/sr.resx

@@ -325,6 +325,10 @@
     <value>Could not retrieve column schema for result set</value>
     <comment></comment>
   </data>
+  <data name="QueryServiceExecutionPlanNotFound" xml:space="preserve">
+    <value>Could not retrieve an execution plan from the result set </value>
+    <comment></comment>
+  </data>
   <data name="PeekDefinitionAzureError" xml:space="preserve">
     <value>This feature is currently not supported on Azure SQL DB and Data Warehouse: {0}</value>
     <comment>.

src/Microsoft.SqlTools.ServiceLayer/sr.strings

@@ -149,6 +149,8 @@ QueryServiceResultSetRowCountOutOfRange = Row count must be a positive integer
 
 QueryServiceResultSetNoColumnSchema = Could not retrieve column schema for result set
 
+QueryServiceExecutionPlanNotFound = Could not retrieve an execution plan from the result set 
+
 ############################################################################
 # Language Service
 

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Common.cs

@@ -72,6 +72,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             return batch;
         }
 
+        public static Batch GetExecutedBatchWithExecutionPlan()
+        {
+            Batch batch = new Batch(StandardQuery, SubsectionDocument, 1, GetFileStreamFactory(new Dictionary<string, byte[]>()));
+            batch.Execute(CreateTestConnection(new[] {GetExecutionPlanTestData()}, false), CancellationToken.None).Wait();
+            return batch;
+        }
+
         public static Query GetBasicExecutedQuery()
         {
             ConnectionInfo ci = CreateTestConnectionInfo(new[] {StandardTestData}, false);
@@ -81,6 +88,15 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             return query;
         }
 
+        public static Query GetBasicExecutedQuery(QueryExecutionSettings querySettings)
+        {
+            ConnectionInfo ci = CreateTestConnectionInfo(new[] {StandardTestData}, false);
+            Query query = new Query(StandardQuery, ci, querySettings, GetFileStreamFactory(new Dictionary<string, byte[]>()));
+            query.Execute();
+            query.ExecutionTask.Wait();
+            return query;
+        }
+
         public static Dictionary<string, string>[] GetTestData(int columns, int rows)
         {
             Dictionary<string, string>[] output = new Dictionary<string, string>[rows];
@@ -97,6 +113,19 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             return output;
         }
 
+
+        public static Dictionary<string, string>[] GetExecutionPlanTestData()
+        {
+            Dictionary<string, string>[] output = new Dictionary<string, string>[1];
+            int col = 0;
+            int row = 0;
+            Dictionary<string, string> rowDictionary = new Dictionary<string, string>();
+            rowDictionary.Add(string.Format("Microsoft SQL Server 2005 XML Showplan", col), string.Format("Execution Plan", col, row));
+            output[row] = rowDictionary;
+            
+            return output;
+        }
+      
         public static Dictionary<string, string>[][] GetTestDataSet(int dataSets)
         {
             List<Dictionary<string, string>[]> output = new List<Dictionary<string, string>[]>();

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/ExecutionPlanTests.cs

@@ -0,0 +1,242 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+using Microsoft.SqlTools.ServiceLayer.Test.Utility;
+using Microsoft.SqlTools.ServiceLayer.SqlContext;
+using Xunit;
+
+namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
+{
+    public class ExecutionPlanTests
+    {
+        #region ResultSet Class Tests
+
+        [Fact]
+        public void ExecutionPlanValid()
+        {
+            // Setup:
+            // ... I have a batch that has been executed with a execution plan 
+            Batch b = Common.GetExecutedBatchWithExecutionPlan();
+
+            // If:
+            // ... I have a result set and I ask for a valid execution plan
+            ResultSet planResultSet = b.ResultSets.First();
+            ExecutionPlan plan = planResultSet.GetExecutionPlan().Result;
+
+            // Then:
+            // ... I should get the execution plan back
+            Assert.Equal("xml", plan.Format);
+            Assert.Contains("Execution Plan", plan.Content);
+        }
+
+        [Fact]
+        public void ExecutionPlanInvalid()
+        {
+            // Setup:
+            // ... I have a batch that has been executed
+            Batch b = Common.GetBasicExecutedBatch();
+
+            // If:
+            // ... I have a result set and I ask for an execution plan that doesn't exist
+            ResultSet planResultSet = b.ResultSets.First();
+
+            // Then:
+            // ... It should throw an exception
+            Assert.ThrowsAsync<Exception>(async () => await planResultSet.GetExecutionPlan());
+        }
+
+        #endregion
+
+        #region Batch Class Tests
+
+        [Fact]
+        public void BatchExecutionPlanValidTest()
+        {
+            // If I have an executed batch which has an execution plan 
+            Batch b = Common.GetExecutedBatchWithExecutionPlan();
+
+            // ... And I ask for a valid execution plan 
+            ExecutionPlan plan = b.GetExecutionPlan(0).Result;
+
+            // Then:
+            // ... I should get the execution plan back
+            Assert.Equal("xml", plan.Format);
+            Assert.Contains("Execution Plan", plan.Content);
+        }
+
+        [Fact]
+        public void BatchExecutionPlanInvalidTest()
+        {
+            // Setup:
+            // ... I have a batch that has been executed without an execution plan 
+            Batch b = Common.GetBasicExecutedBatch();
+
+            // If: 
+            // ... I ask for an invalid execution plan 
+            Assert.ThrowsAsync<Exception>(async () => await b.GetExecutionPlan(0));
+        }
+
+        [Theory]
+        [InlineData(-1)]  // Invalid result set, too low
+        [InlineData(2)]   // Invalid result set, too high
+        public void BatchExecutionPlanInvalidParamsTest(int resultSetIndex)
+        {
+            // If I have an executed batch which has an execution plan 
+            Batch b = Common.GetExecutedBatchWithExecutionPlan();
+
+            // ... And I ask for an execution plan with an invalid result set index
+            // Then: 
+            // ... It should throw an exception
+            Assert.ThrowsAsync<ArgumentOutOfRangeException>(async () => await b.GetExecutionPlan(resultSetIndex));
+        }
+
+        #endregion
+
+        #region Query Class Tests
+
+        [Theory]
+        [InlineData(-1)]  // Invalid batch, too low
+        [InlineData(2)]   // Invalid batch, too high
+        public void QueryExecutionPlanInvalidParamsTest(int batchIndex)
+        {
+            // Setup query settings
+            QueryExecutionSettings querySettings = new QueryExecutionSettings(); 
+            querySettings.ExecutionPlanOptions = new ExecutionPlanOptions()
+            { 
+                IncludeActualExecutionPlanXml = false,
+                IncludeEstimatedExecutionPlanXml = true
+            };
+
+            // If I have an executed query
+            Query q = Common.GetBasicExecutedQuery(querySettings);
+
+            // ... And I ask for a subset with an invalid result set index
+            // Then: 
+            // ... It should throw an exception
+            Assert.ThrowsAsync<ArgumentOutOfRangeException>(async () => await q.GetExecutionPlan(batchIndex, 0));
+        }
+
+        #endregion
+
+
+        #region Service Intergration Tests
+
+        [Fact]
+        public async Task ExecutionPlanServiceValidTest()
+        {
+            // If:
+            // ... I have a query that has results in the form of an execution plan 
+            var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery);
+            var queryService = Common.GetPrimedExecutionService(new[] {Common.GetExecutionPlanTestData()}, true, false, workspaceService);
+            var executeParams = new QueryExecuteParams {QuerySelection = null, OwnerUri = Common.OwnerUri};
+            executeParams.ExecutionPlanOptions = new ExecutionPlanOptions()
+            { 
+                IncludeActualExecutionPlanXml = false,
+                IncludeEstimatedExecutionPlanXml = true
+            };
+            var executeRequest = RequestContextMocks.Create<QueryExecuteResult>(null);
+            await queryService.HandleExecuteRequest(executeParams, executeRequest.Object);
+            await queryService.ActiveQueries[Common.OwnerUri].ExecutionTask;
+
+            // ... And I then ask for a valid execution plan 
+            var executionPlanParams = new QueryExecutionPlanParams { OwnerUri = Common.OwnerUri, BatchIndex = 0, ResultSetIndex = 0 };
+            var executionPlanRequest = new EventFlowValidator<QueryExecutionPlanResult>()
+                .AddResultValidation(r =>
+                {
+                    // Then: Messages should be null and execution plan should not be null
+                    Assert.NotNull(r.ExecutionPlan);
+                }).Complete();
+            await queryService.HandleExecutionPlanRequest(executionPlanParams, executionPlanRequest.Object);
+            executionPlanRequest.Validate();
+        }
+
+        
+        [Fact]
+        public async Task ExecutionPlanServiceMissingQueryTest()
+        {
+            // If:
+            // ... I ask for an execution plan for a file that hasn't executed a query
+            var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery);
+            var queryService = Common.GetPrimedExecutionService(null, true, false, workspaceService);
+            var executionPlanParams = new QueryExecutionPlanParams { OwnerUri = Common.OwnerUri, ResultSetIndex = 0, BatchIndex = 0 };
+            var executionPlanRequest = new EventFlowValidator<QueryExecutionPlanResult>()
+                .AddErrorValidation<string>(r =>
+                {
+                    // Then: It should return a populated error 
+                    Assert.NotNull(r);
+                }).Complete();
+            await queryService.HandleExecutionPlanRequest(executionPlanParams, executionPlanRequest.Object);
+            executionPlanRequest.Validate();
+        }
+
+        [Fact]
+        public async Task ExecutionPlanServiceUnexecutedQueryTest()
+        {
+            // If:
+            // ... I have a query that hasn't finished executing (doesn't matter what)
+            var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery);
+            var queryService = Common.GetPrimedExecutionService(new[] { Common.GetExecutionPlanTestData() }, true, false, workspaceService);
+            var executeParams = new QueryExecuteParams { QuerySelection = null, OwnerUri = Common.OwnerUri };
+            executeParams.ExecutionPlanOptions = new ExecutionPlanOptions()
+            { 
+                IncludeActualExecutionPlanXml = false,
+                IncludeEstimatedExecutionPlanXml = true
+            };
+            var executeRequest = RequestContextMocks.Create<QueryExecuteResult>(null);
+            await queryService.HandleExecuteRequest(executeParams, executeRequest.Object);
+            await queryService.ActiveQueries[Common.OwnerUri].ExecutionTask;
+            queryService.ActiveQueries[Common.OwnerUri].Batches[0].ResultSets[0].hasBeenRead = false;
+
+            // ... And I then ask for a valid execution plan from it 
+            var executionPlanParams = new QueryExecutionPlanParams { OwnerUri = Common.OwnerUri, ResultSetIndex = 0, BatchIndex = 0 };
+            var executionPlanRequest = new EventFlowValidator<QueryExecutionPlanResult>()
+                .AddErrorValidation<string>(r =>
+                {
+                    // Then: It should return a populated error 
+                    Assert.NotNull(r);
+                }).Complete();
+            await queryService.HandleExecutionPlanRequest(executionPlanParams, executionPlanRequest.Object);
+            executionPlanRequest.Validate();
+        }
+
+        [Fact]
+        public async Task ExecutionPlanServiceOutOfRangeSubsetTest()
+        {
+            // If:
+            // ... I have a query that doesn't have any result sets
+            var workspaceService = Common.GetPrimedWorkspaceService(Common.StandardQuery);
+            var queryService = Common.GetPrimedExecutionService(null, true, false, workspaceService);
+            var executeParams = new QueryExecuteParams { QuerySelection = null, OwnerUri = Common.OwnerUri };
+            executeParams.ExecutionPlanOptions = new ExecutionPlanOptions()
+            { 
+                IncludeActualExecutionPlanXml = false,
+                IncludeEstimatedExecutionPlanXml = true
+            };
+            var executeRequest = RequestContextMocks.Create<QueryExecuteResult>(null);
+            await queryService.HandleExecuteRequest(executeParams, executeRequest.Object);
+            await queryService.ActiveQueries[Common.OwnerUri].ExecutionTask;
+
+            // ... And I then ask for an execution plan from a result set 
+            var executionPlanParams = new QueryExecutionPlanParams { OwnerUri = Common.OwnerUri, ResultSetIndex = 0, BatchIndex = 0 };
+            var executionPlanRequest = new EventFlowValidator<QueryExecutionPlanResult>()
+                .AddErrorValidation<string>(r =>
+                {
+                    // Then: It should return a populated error 
+                    Assert.NotNull(r);
+                }).Complete();
+            await queryService.HandleExecutionPlanRequest(executionPlanParams, executionPlanRequest.Object);
+            executionPlanRequest.Validate();
+        }
+        
+        #endregion
+        
+    }
+}

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/SpecialActionTests.cs

@@ -0,0 +1,79 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.SqlTools.ServiceLayer.QueryExecution;
+using Microsoft.SqlTools.ServiceLayer.Test.Utility;
+using Xunit;
+
+namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
+{
+    public class SpecialActionTests
+    {
+
+        [Fact]
+        public void SpecialActionInstantiation()
+        {
+            // If:
+            // ... I create a special action object 
+            var specialAction = new SpecialAction();
+
+            // Then:
+            // ... The special action should be set to none and only none
+            Assert.Equal(true, specialAction.None);
+            Assert.Equal(false, specialAction.ExpectYukonXMLShowPlan);
+        }
+
+        [Fact]
+        public void SpecialActionNoneProperty()
+        {
+            // If:
+            // ... I create a special action object and add properties but set it back to none 
+            var specialAction = new SpecialAction();
+            specialAction.ExpectYukonXMLShowPlan = true;
+            specialAction.None = true;
+
+            // Then:
+            // ... The special action should be set to none and only none
+            Assert.Equal(true, specialAction.None);
+            Assert.Equal(false, specialAction.ExpectYukonXMLShowPlan);
+        }
+
+        [Fact]
+        public void SpecialActionExpectYukonXmlShowPlanReset()
+        {
+            // If:
+            // ... I create a special action object and add properties but set the property back to false 
+            var specialAction = new SpecialAction();
+            specialAction.ExpectYukonXMLShowPlan = true;
+            specialAction.ExpectYukonXMLShowPlan = false;
+
+            // Then:
+            // ... The special action should be set to none and only none
+            Assert.Equal(true, specialAction.None);
+            Assert.Equal(false, specialAction.ExpectYukonXMLShowPlan);
+        }
+
+        [Fact]
+        public void SpecialActionCombiningProperties()
+        {
+            // If:
+            // ... I create a special action object and add properties and combine with the same property 
+            var specialAction = new SpecialAction();
+            specialAction.ExpectYukonXMLShowPlan = true;
+
+            var specialAction2 = new SpecialAction();
+            specialAction2.ExpectYukonXMLShowPlan = true;
+
+            specialAction.CombineSpecialAction(specialAction2);
+
+            // Then:
+            // ... The special action should be set to none and only none
+            Assert.Equal(false, specialAction.None);
+            Assert.Equal(true, specialAction.ExpectYukonXMLShowPlan);
+        }
+
+
+    }
+}