ckaczor/sqltoolsservice
1
0
Fork 0
mirror of https://github.com/ckaczor/sqltoolsservice.git synced 2026-01-16 09:35:36 -05:00
Code Issues Packages Projects Releases Wiki Activity

SqlCmd Connect/On Error/Include commands support (#898)

Browse Source 
* Initial Investigation

* Working code with include, connect, on error and tests

* Adding some loc strings

* Some cleanup and more tests

* Some dummy change to trigger build

* Adding PR comments

* Addressing PR comments
This commit is contained in:
Udeesha Gautam
2020-01-10 17:54:39 -08:00
committed by GitHub
parent d512c101c0
commit fe17962ac9
25 changed files with 925 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
View File

@@ -18,6 +18,7 @@ using Microsoft.SqlTools.Utility;
using System.Globalization;
using System.Collections.ObjectModel;
using Microsoft.SqlTools.ServiceLayer.Connection;
using Microsoft.SqlTools.ServiceLayer.BatchParser;
namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
{
@@ -71,7 +72,14 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
#endregion
internal Batch(string batchText, SelectionData selection, int ordinalId,
IFileStreamFactory outputFileFactory, int executionCount = 1, bool getFullColumnSchema = false)
IFileStreamFactory outputFileFactory, SqlCmdCommand sqlCmdCommand, int executionCount = 1, bool getFullColumnSchema = false) : this(batchText, selection, ordinalId,
outputFileFactory, executionCount, getFullColumnSchema)
{
this.SqlCmdCommand = sqlCmdCommand;
}
internal Batch(string batchText, SelectionData selection, int ordinalId,
IFileStreamFactory outputFileFactory, int executionCount = 1, bool getFullColumnSchema = false)
{
// Sanity check for input
Validate.IsNotNullOrEmptyString(nameof(batchText), batchText);
@@ -138,6 +146,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
/// called from the Batch but from the ResultSet instance.
/// </summary>
public event ResultSet.ResultSetAsyncEventHandler ResultSetUpdated;
/// <summary>
/// Event that will be called when additional rows in the result set are available (rowCount available has increased). It will not be
/// called from the Batch but from the ResultSet instance.
/// </summary>
public event EventHandler<bool> HandleOnErrorAction;
#endregion
#region Properties
@@ -147,6 +161,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
/// </summary>
public string BatchText { get; set; }
public SqlCmdCommand SqlCmdCommand { get; set; }
public int BatchExecutionCount { get; private set; }
/// <summary>
/// Localized timestamp for when the execution completed.
@@ -243,13 +259,24 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
#endregion
#region Public Methods
/// <summary>
/// Executes this batch and captures any server messages that are returned.
/// </summary>
/// <param name="conn">The connection to use to execute the batch</param>
/// <param name="cancellationToken">Token for cancelling the execution</param>
public async Task Execute(DbConnection conn, CancellationToken cancellationToken)
{
await Execute(conn, cancellationToken, OnErrorAction.Ignore);
}
/// <summary>
/// Executes this batch and captures any server messages that are returned.
/// </summary>
/// <param name="conn">The connection to use to execute the batch</param>
/// <param name="cancellationToken">Token for cancelling the execution</param>
/// <param name="onErrorAction">Continue (Ignore) or Exit on Error. This comes only in SQLCMD mode</param>
public async Task Execute(DbConnection conn, CancellationToken cancellationToken, OnErrorAction onErrorAction = OnErrorAction.Ignore)
{
// Sanity check to make sure we haven't already run this batch
if (HasExecuted)
@@ -270,10 +297,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
{
sqlConn.GetUnderlyingConnection().InfoMessage += ServerMessageHandler;
}
try
{
await DoExecute(conn, cancellationToken);
await DoExecute(conn, cancellationToken, onErrorAction);
}
catch (TaskCanceledException)
{
@@ -308,7 +335,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
}
private async Task DoExecute(DbConnection conn, CancellationToken cancellationToken)
private async Task DoExecute(DbConnection conn, CancellationToken cancellationToken, OnErrorAction onErrorAction = OnErrorAction.Ignore)
{
bool canContinue = true;
int timesLoop = this.BatchExecutionCount;
@@ -326,6 +353,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
catch (DbException dbe)
{
HasError = true;
if (onErrorAction == OnErrorAction.Exit)
{
throw new SqlCmdException(dbe.Message);
}
canContinue = await UnwrapDbException(dbe);
if (canContinue)
{
@@ -445,7 +476,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
private void ExtendResultMetadata(List<DbColumn[]> columnSchemas, List<ResultSet> results)
{
if (columnSchemas.Count != results.Count)
{
{
return;
}
@@ -679,6 +710,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
{
this.HasError = true;
}
if (this.HandleOnErrorAction != null)
{
HandleOnErrorAction(this, isError);
}
}
/// <summary>
@@ -693,6 +729,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
{
bool canIgnore = true;
SqlException se = dbe as SqlException;
if (se != null)
{
var errors = se.Errors.Cast<SqlError>().ToList();
@@ -729,7 +766,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
/// </summary>
private SpecialAction ProcessResultSetSpecialActions()
{
foreach (ResultSet resultSet in resultSets)
foreach (ResultSet resultSet in resultSets)
{
specialAction.CombineSpecialAction(resultSet.Summary.SpecialAction);
}

146
src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs
View File

@@ -30,12 +30,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
public class Query : IDisposable
{
#region Constants
/// <summary>
/// "Error" code produced by SQL Server when the database context (name) for a connection changes.
/// </summary>
private const int DatabaseContextChangeErrorNumber = 5701;
/// <summary>
/// ON keyword
/// </summary>
@@ -45,7 +45,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
/// OFF keyword
/// </summary>
private const string Off = "OFF";
/// <summary>
/// showplan_xml statement
/// </summary>
@@ -55,7 +55,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
/// statistics xml statement
/// </summary>
private const string SetStatisticsXml = "SET STATISTICS XML {0}";
#endregion
#region Member Variables
@@ -84,7 +84,23 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
/// <summary>
/// Name of the new database if the database name was changed in the query
/// </summary>
private string newDatabaseName;
private string newDatabaseName;
/// <summary>
/// On Error Action for the query in SQLCMD Mode- Ignore or Exit
/// </summary>
private OnErrorAction onErrorAction;
/// <summary>
/// Connection that is used for query to run
/// This is always initialized from editor connection but might be different in case of SQLCMD mode
/// </summary>
private DbConnection queryConnection;
/// <summary>
/// Cancelled but not user but by SQLCMD settings
/// </summary>
private bool CancelledBySqlCmd;
#endregion
@@ -96,9 +112,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
/// <param name="settings">Settings for how to execute the query, from the user</param>
/// <param name="outputFactory">Factory for creating output files</param>
public Query(
string queryText,
ConnectionInfo connection,
QueryExecutionSettings settings,
string queryText,
ConnectionInfo connection,
QueryExecutionSettings settings,
IFileStreamFactory outputFactory,
bool getFullColumnSchema = false,
bool applyExecutionSettings = false)
@@ -125,13 +141,14 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
var batchSelection = parserResult
.Select((batchDefinition, index) =>
new Batch(batchDefinition.BatchText,
new Batch(batchDefinition.BatchText,
new SelectionData(
batchDefinition.StartLine-1,
batchDefinition.StartColumn-1,
batchDefinition.EndLine-1,
batchDefinition.EndColumn-1),
batchDefinition.EndColumn-1),
index, outputFactory,
batchDefinition.SqlCmdCommand,
batchDefinition.BatchExecutionCount,
getFullColumnSchema));
@@ -154,14 +171,14 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
/// </summary>
/// <param name="query">The query that completed</param>
public delegate Task QueryAsyncEventHandler(Query query);
/// <summary>
/// Delegate type for callback when a query fails
/// </summary>
/// <param name="query">Query that raised the event</param>
/// <param name="exception">Exception that caused the query to fail</param>
public delegate Task QueryAsyncErrorEventHandler(Query query, Exception exception);
/// <summary>
/// Event to be called when a batch is completed.
/// </summary>
@@ -356,7 +373,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
/// </param>
/// <param name="successHandler">Delegate to call when the request completes successfully</param>
/// <param name="failureHandler">Delegate to call if the request fails</param>
public void SaveAs(SaveResultsRequestParams saveParams, IFileStreamFactory fileFactory,
public void SaveAs(SaveResultsRequestParams saveParams, IFileStreamFactory fileFactory,
ResultSet.SaveAsAsyncEventHandler successHandler, ResultSet.SaveAsFailureAsyncEventHandler failureHandler)
{
// Sanity check to make sure that the batch is within bounds
@@ -385,7 +402,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
// Mark that we've internally executed
hasExecuteBeenCalled = true;
// Don't actually execute if there aren't any batches to execute
if (Batches.Length == 0)
{
@@ -399,9 +416,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
}
return;
}
// Locate and setup the connection
DbConnection queryConnection = await ConnectionService.Instance.GetOrOpenConnection(editorConnection.OwnerUri, ConnectionType.Query);
queryConnection = await ConnectionService.Instance.GetOrOpenConnection(editorConnection.OwnerUri, ConnectionType.Query);
onErrorAction = OnErrorAction.Ignore;
sqlConn = queryConnection as ReliableSqlConnection;
if (sqlConn != null)
{
@@ -410,7 +428,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
sqlConn.GetUnderlyingConnection().InfoMessage += OnInfoMessage;
}
// Execute beforeBatches synchronously, before the user defined batches
foreach (Batch b in BeforeBatches)
{
@@ -427,7 +445,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
b.ResultSetCompletion += ResultSetCompleted;
b.ResultSetAvailable += ResultSetAvailable;
b.ResultSetUpdated += ResultSetUpdated;
await b.Execute(queryConnection, cancellationSource.Token);
await ExecuteBatch(b);
}
// Execute afterBatches synchronously, after the user defined batches
@@ -445,7 +464,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
catch (Exception e)
{
HasErrored = true;
if (e is OperationCanceledException)
if (e is SqlCmdException || CancelledBySqlCmd)
{
await BatchMessageSent(new ResultMessage(SR.SqlCmdExitOnError, false, null));
}
else if (e is OperationCanceledException)
{
await BatchMessageSent(new ResultMessage(SR.QueryServiceQueryCancelled, false, null));
}
@@ -463,7 +486,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
// Subscribe to database informational messages
sqlConn.GetUnderlyingConnection().InfoMessage -= OnInfoMessage;
}
// If any message notified us we had changed databases, then we must let the connection service know
if (newDatabaseName != null)
{
@@ -480,6 +503,63 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
}
}
}
private async Task ExecuteBatch(Batch b)
{
if (b.SqlCmdCommand != null)
{
b.HandleOnErrorAction += HandleOnErrorAction;
await PerformInplaceSqlCmdAction(b);
}
await b.Execute(queryConnection, cancellationSource.Token, this.onErrorAction);
if (CancelledBySqlCmd)
{
throw new SqlCmdException(SR.SqlCmdExitOnError);
}
}
public async Task PerformInplaceSqlCmdAction(Batch b)
{
try
{
switch (b.SqlCmdCommand.LexerTokenType)
{
case LexerTokenType.Connect:
var qc = (b.SqlCmdCommand as ConnectSqlCmdCommand)?.Connect();
queryConnection = qc ?? queryConnection;
break;
case LexerTokenType.OnError:
onErrorAction = (b.SqlCmdCommand as OnErrorSqlCmdCommand).Action;
break;
default:
throw new SqlCmdException(string.Format(SR.SqlCmdUnsupportedToken, b.SqlCmdCommand.LexerTokenType));
}
}
catch (Exception ex)
{
b.HasError = true;
await BatchMessageSent(new ResultMessage(ex.Message, true, null));
if (this.onErrorAction == OnErrorAction.Exit)
{
HasCancelled = true;
CancelledBySqlCmd = true;
cancellationSource.Cancel();
throw new SqlCmdException(SR.SqlCmdExitOnError);
}
}
}
private void HandleOnErrorAction(object sender, bool iserror)
{
if (iserror && this.onErrorAction == OnErrorAction.Exit)
{
HasCancelled = true;
CancelledBySqlCmd = true;
cancellationSource.Cancel();
}
}
/// <summary>
/// Handler for database messages during query execution
@@ -511,8 +591,8 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
}
private void ApplyExecutionSettings(
ConnectionInfo connection,
QueryExecutionSettings settings,
ConnectionInfo connection,
QueryExecutionSettings settings,
IFileStreamFactory outputFactory)
{
outputFactory.QueryExecutionSettings = settings;
@@ -549,7 +629,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
}
if (settings.StatisticsIO)
{
{
builderBefore.AppendFormat("{0} ", helper.GetSetStatisticsIOString(true));
builderAfter.AppendFormat("{0} ", helper.GetSetStatisticsIOString (false));
}
@@ -562,35 +642,35 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
}
if (settings.ParseOnly)
{
{
builderBefore.AppendFormat("{0} ", helper.GetSetParseOnlyString(true));
builderAfter.AppendFormat("{0} ", helper.GetSetParseOnlyString(false));
}
// append first part of exec options
builderBefore.AppendFormat("{0} {1} {2}",
builderBefore.AppendFormat("{0} {1} {2}",
helper.SetRowCountString, helper.SetTextSizeString, helper.SetNoCountString);
if (!connection.IsSqlDW)
{
// append second part of exec options
builderBefore.AppendFormat(" {0} {1} {2} {3} {4} {5} {6}",
helper.SetConcatenationNullString,
helper.SetArithAbortString,
helper.SetLockTimeoutString,
helper.SetQueryGovernorCostString,
helper.SetDeadlockPriorityString,
helper.SetConcatenationNullString,
helper.SetArithAbortString,
helper.SetLockTimeoutString,
helper.SetQueryGovernorCostString,
helper.SetDeadlockPriorityString,
helper.SetTransactionIsolationLevelString,
// We treat XACT_ABORT special in that we don't add anything if the option
// isn't checked. This is because we don't want to be overwriting the server
// if it has a default of ON since that's something people would specifically
// set and having a client change it could be dangerous (the reverse is much
// less risky)
// The full fix would probably be to make the options tri-state instead of
// just on/off, where the default is to use the servers default. Until that
// happens though this is the best solution we came up with. See TFS#7937925
// Note that users can always specifically add SET XACT_ABORT OFF to their
// queries if they do truly want to set it off. We just don't want to
// do it silently (since the default is going to be off)
@@ -655,7 +735,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
/// <summary>
/// Does this connection support XML Execution plans
/// </summary>
private bool DoesSupportExecutionPlan(ConnectionInfo connectionInfo)
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);