// // 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.Text; using Microsoft.SqlServer.Management.Common; using Microsoft.SqlServer.Management.Smo; namespace Microsoft.SqlTools.ServiceLayer.TaskServices { /// /// Any SMO operation that supports scripting should implement this class. /// It provides all of the configuration needed to choose between scripting or execution mode, /// hook into the Task manager framework, and send success / completion notifications to the caller. /// public abstract class SmoScriptableTaskOperation : IScriptableTaskOperation { /// /// Script content /// public string ScriptContent { get; set; } /// /// If an error occurred during task execution, this field contains the error message text /// public abstract string ErrorMessage { get; } /// /// SMO Server instance used for the operation /// public abstract Server Server { get; } /// /// Cancels the operation /// public abstract void Cancel(); /// /// Updates messages in sql task given new progress message /// /// public void OnMessageAdded(TaskMessage message) { if (this.SqlTask != null) { this.SqlTask.AddMessage(message.Description, message.Status); } } /// /// Updates scripts in sql task given new script /// /// public void OnScriptAdded(TaskScript script) { this.SqlTask.AddScript(script.Status, script.Script, script.ErrorMessage); } /// /// Executes the operations /// public abstract void Execute(); /// /// Execute the operation for given execution mode /// /// public virtual void Execute(TaskExecutionMode mode) { var currentExecutionMode = Server.ConnectionContext.SqlExecutionModes; try { if (Server != null) { Server.ConnectionContext.CapturedSql.Clear(); SetExecutionMode(mode); } Execute(); GenerateScript(mode); } catch { throw; } finally { Server.ConnectionContext.CapturedSql.Clear(); Server.ConnectionContext.SqlExecutionModes = currentExecutionMode; } } protected void GenerateScript(TaskExecutionMode mode) { if (mode == TaskExecutionMode.Script || mode == TaskExecutionMode.ExecuteAndScript) { this.ScriptContent = GetScriptContent(); if (SqlTask != null) { OnScriptAdded(new TaskScript { Status = SqlTaskStatus.Succeeded, Script = this.ScriptContent }); } } } protected void SetExecutionMode(TaskExecutionMode mode) { switch (mode) { case TaskExecutionMode.Execute: { Server.ConnectionContext.SqlExecutionModes = SqlExecutionModes.ExecuteSql; break; } case TaskExecutionMode.ExecuteAndScript: { Server.ConnectionContext.SqlExecutionModes = SqlExecutionModes.ExecuteAndCaptureSql; break; } case TaskExecutionMode.Script: { Server.ConnectionContext.SqlExecutionModes = SqlExecutionModes.CaptureSql; break; } } } private string GetScriptContent() { StringBuilder sb = new StringBuilder(); foreach (String s in this.Server.ConnectionContext.CapturedSql.Text) { sb.Append(s); sb.Append(Environment.NewLine); } return sb.ToString(); } /// /// The sql task to run the operations /// public SqlTask SqlTask { get; set; } } }