generic way to support scriptable operations (#438)

* implemented a generic way to support scriptable operations

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/DisasterRecovery/RestoreDatabaseServiceTests.cs

@@ -82,7 +82,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
             //Verify that all backupsets are restored
             int[] expectedTable = new int[] { };
 
-            await VerifyRestoreMultipleBackupSets(backupFiles, indexToDelete, expectedTable);
+            await VerifyRestoreMultipleBackupSets(backupFiles, indexToDelete, expectedTable, TaskExecutionModeFlag.Execute);
         }
 
         [Fact]
@@ -95,7 +95,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
             {
                 Dictionary<string, object> options = new Dictionary<string, object>();
                 options.Add(RestoreOptionsHelper.ReplaceDatabase, true);
-                await VerifyRestore(null, databaseNameToRestoreFrom, true, true, testDb.DatabaseName, null, options, (database) =>
+                await VerifyRestore(null, databaseNameToRestoreFrom, true, TaskExecutionModeFlag.ExecuteAndScript, testDb.DatabaseName, null, options, (database) =>
                 {
                     return database.Tables.Contains("tb1", "test");
                 });
@@ -129,14 +129,14 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
             await VerifyRestoreMultipleBackupSets(backupFiles, indexToDelete, expectedTable);
         }
 
-        private async Task VerifyRestoreMultipleBackupSets(string[] backupFiles, int backupSetIndexToDelete, int[] expectedSelectedIndexes)
+        private async Task VerifyRestoreMultipleBackupSets(string[] backupFiles, int backupSetIndexToDelete, int[] expectedSelectedIndexes, TaskExecutionModeFlag executionMode = TaskExecutionModeFlag.ExecuteAndScript)
         {
             var testDb = await SqlTestDb.CreateNewAsync(TestServerType.OnPrem, false, null, null, "RestoreTest");
             try
             {
                 string targetDbName = testDb.DatabaseName;
                 bool canRestore = true;
-                var response = await VerifyRestore(backupFiles, null, canRestore, false, targetDbName, null, null);
+                var response = await VerifyRestore(backupFiles, null, canRestore, TaskExecutionModeFlag.None, targetDbName, null, null);
                 Assert.True(response.BackupSetsToRestore.Count() >= 2);
                 var allIds = response.BackupSetsToRestore.Select(x => x.Id).ToList();
                 if (backupSetIndexToDelete >= 0)
@@ -146,20 +146,24 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
                 string[] selectedIds = allIds.ToArray();
                 Dictionary<string, object> options = new Dictionary<string, object>();
                 options.Add(RestoreOptionsHelper.ReplaceDatabase, true);
-                response = await VerifyRestore(backupFiles, null, canRestore, true, targetDbName, selectedIds, options, (database) =>
+                response = await VerifyRestore(backupFiles, null, canRestore, executionMode, targetDbName, selectedIds, options, (database) =>
                 {
-                    bool tablesFound = true;
-                    for (int i = 0; i < tableNames.Length; i++)
+                    if (executionMode.HasFlag(TaskExecutionModeFlag.Execute))
                     {
-                        string tableName = tableNames[i];
-                        if (!database.Tables.Contains(tableName, "test") && expectedSelectedIndexes.Contains(i))
+                        bool tablesFound = true;
+                        for (int i = 0; i < tableNames.Length; i++)
                         {
-                            tablesFound = false;
-                            break;
+                            string tableName = tableNames[i];
+                            if (!database.Tables.Contains(tableName, "test") && expectedSelectedIndexes.Contains(i))
+                            {
+                                tablesFound = false;
+                                break;
+                            }
                         }
+                        bool numberOfTableCreatedIsCorrect = database.Tables.Count == expectedSelectedIndexes.Length;
+                        return numberOfTableCreatedIsCorrect && tablesFound;
                     }
-                    bool numberOfTableCreatedIsCorrect = database.Tables.Count == expectedSelectedIndexes.Length;
-                    return numberOfTableCreatedIsCorrect && tablesFound;
+                    return true;
                 });
 
                 for (int i = 0; i < response.BackupSetsToRestore.Count(); i++)
@@ -190,7 +194,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
                 Dictionary<string, object> options = new Dictionary<string, object>();
                 options.Add(RestoreOptionsHelper.ReplaceDatabase, true);
 
-                await VerifyRestore(new string[] { fullBackupFilePath }, null, canRestore, true, testDb.DatabaseName, null, options);
+                await VerifyRestore(new string[] { fullBackupFilePath }, null, canRestore, TaskExecutionModeFlag.ExecuteAndScript, testDb.DatabaseName, null, options);
             }
             finally
             {
@@ -212,7 +216,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
                 await VerifyBackupFileCreated();
                 bool canRestore = true;
 
-                await VerifyRestore(new string[] { fullBackupFilePath }, null, canRestore, false, testDb.DatabaseName, null, null);
+                await VerifyRestore(new string[] { fullBackupFilePath }, null, canRestore, TaskExecutionModeFlag.None, testDb.DatabaseName, null, null);
             }
             finally
             {
@@ -229,7 +233,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
 
             string[] backupFileNames = new string[] { "FullBackup.bak", "DiffBackup.bak" };
             bool canRestore = true;
-            var response = await VerifyRestore(backupFileNames, null, canRestore, false, "RestoredFromTwoBackupFile");
+            var response = await VerifyRestore(backupFileNames, null, canRestore, TaskExecutionModeFlag.None, "RestoredFromTwoBackupFile");
             Assert.True(response.BackupSetsToRestore.Count() == 2);
         }
 
@@ -239,13 +243,13 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
 
             string[] backupFileNames = new string[] { "FullBackup.bak", "DiffBackup.bak" };
             bool canRestore = true;
-            var response = await VerifyRestore(backupFileNames, null, canRestore, false, "RestoredFromTwoBackupFile");
+            var response = await VerifyRestore(backupFileNames, null, canRestore, TaskExecutionModeFlag.None, "RestoredFromTwoBackupFile");
             Assert.True(response.BackupSetsToRestore.Count() == 2);
             var fileInfo = response.BackupSetsToRestore.FirstOrDefault(x => x.GetPropertyValueAsString(BackupSetInfo.BackupTypePropertyName) != RestoreConstants.TypeFull);
             if(fileInfo != null)
             {
                 var selectedBackupSets = new string[] { fileInfo.Id };
-                await VerifyRestore(backupFileNames, null, true, false, "RestoredFromTwoBackupFile", selectedBackupSets);
+                await VerifyRestore(backupFileNames, null, true, TaskExecutionModeFlag.None, "RestoredFromTwoBackupFile", selectedBackupSets);
             }
         }
 
@@ -255,13 +259,13 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
 
             string[] backupFileNames = new string[] { "FullBackup.bak", "DiffBackup.bak" };
             bool canRestore = true;
-            var response = await VerifyRestore(backupFileNames, null, canRestore, false, "RestoredFromTwoBackupFile");
+            var response = await VerifyRestore(backupFileNames, null, canRestore, TaskExecutionModeFlag.None, "RestoredFromTwoBackupFile");
             Assert.True(response.BackupSetsToRestore.Count() == 2);
             var fileInfo = response.BackupSetsToRestore.FirstOrDefault(x => x.GetPropertyValueAsString(BackupSetInfo.BackupTypePropertyName) == RestoreConstants.TypeFull);
             if (fileInfo != null)
             {
                 var selectedBackupSets = new string[] { fileInfo.Id };
-                await VerifyRestore(backupFileNames, null, true, false, "RestoredFromTwoBackupFile2", selectedBackupSets);
+                await VerifyRestore(backupFileNames, null, true, TaskExecutionModeFlag.None, "RestoredFromTwoBackupFile2", selectedBackupSets);
             }
         }
 
@@ -272,7 +276,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
 
             string backupFileName = fullBackupFilePath;
             bool canRestore = true;
-            var restorePlan = await VerifyRestore(backupFileName, canRestore, true);
+            var restorePlan = await VerifyRestore(backupFileName, canRestore, TaskExecutionModeFlag.Execute);
             Assert.NotNull(restorePlan.BackupSetsToRestore);
         }
 
@@ -283,7 +287,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
 
             string backupFileName = fullBackupFilePath;
             bool canRestore = true;
-            var restorePlan = await VerifyRestore(backupFileName, canRestore, true, "NewRestoredDatabase");
+            var restorePlan = await VerifyRestore(backupFileName, canRestore, TaskExecutionModeFlag.ExecuteAndScript, "NewRestoredDatabase");
         }
 
         [Fact]
@@ -414,16 +418,16 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
             await TestServiceProvider.Instance.RunQueryAsync(TestServerType.OnPrem, "master", dropDatabaseQuery);
         }
 
-        private async Task<RestorePlanResponse> VerifyRestore(string backupFileName, bool canRestore, bool execute = false, string targetDatabase = null)
+        private async Task<RestorePlanResponse> VerifyRestore(string backupFileName, bool canRestore, TaskExecutionModeFlag executionMode = TaskExecutionModeFlag.None, string targetDatabase = null)
         {
-            return await VerifyRestore(new string[] { backupFileName }, null, canRestore, execute, targetDatabase);
+            return await VerifyRestore(new string[] { backupFileName }, null, canRestore, executionMode, targetDatabase);
         }
 
         private async Task<RestorePlanResponse> VerifyRestore(
             string[] backupFileNames = null, 
             string sourceDbName = null,
-            bool canRestore = true, 
-            bool execute = false, 
+            bool canRestore = true,
+            TaskExecutionModeFlag executionMode = TaskExecutionModeFlag.None, 
             string targetDatabase = null, 
             string[] selectedBackupSets = null,
             Dictionary<string, object> options = null,
@@ -496,7 +500,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
                     Assert.NotNull(response.PlanDetails[RestoreOptionsHelper.StandbyFile]);
                     Assert.NotNull(response.PlanDetails[RestoreOptionsHelper.StandbyFile]);
                    
-                    if(execute)
+                    if(executionMode != TaskExecutionModeFlag.None)
                     {
                         try
                         {
@@ -504,21 +508,28 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.DisasterRecovery
                             restoreDataObject = service.CreateRestoreDatabaseTaskDataObject(request);
                             Assert.Equal(response.SessionId, restoreDataObject.SessionId);
                             request.RelocateDbFiles = !restoreDataObject.DbFilesLocationAreValid();
-                            service.ExecuteRestore(restoreDataObject);
-                            Assert.True(restoreDataObject.Server.Databases.Contains(targetDatabase));
+                            restoreDataObject.Execute((TaskExecutionMode)Enum.Parse(typeof(TaskExecutionMode), executionMode.ToString()));
 
-                            if (verifyDatabase != null)
+                            if (executionMode.HasFlag(TaskExecutionModeFlag.Execute))
                             {
-                                Assert.True(verifyDatabase(restoreDataObject.Server.Databases[targetDatabase]));
-                            }
+                                Assert.True(restoreDataObject.Server.Databases.Contains(targetDatabase));
 
-                            //To verify the backupset that are restored, verifying the database is a better options.
-                            //Some tests still verify the number of backup sets that are executed which in some cases can be less than the selected list
-                            if (verifyDatabase == null && selectedBackupSets != null)
+                                if (verifyDatabase != null)
+                                {
+                                    Assert.True(verifyDatabase(restoreDataObject.Server.Databases[targetDatabase]));
+                                }
+
+                                //To verify the backupset that are restored, verifying the database is a better options.
+                                //Some tests still verify the number of backup sets that are executed which in some cases can be less than the selected list
+                                if (verifyDatabase == null && selectedBackupSets != null)
+                                {
+                                    Assert.Equal(selectedBackupSets.Count(), restoreDataObject.RestorePlanToExecute.RestoreOperations.Count());
+                                }
+                            }
+                            if(executionMode.HasFlag(TaskExecutionModeFlag.Script))
                             {
-                                Assert.Equal(selectedBackupSets.Count(), restoreDataObject.RestorePlanToExecute.RestoreOperations.Count());
+                                Assert.False(string.IsNullOrEmpty(restoreDataObject.ScriptContent));
                             }
-
                         }
                         catch(Exception ex)
                         {

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/DisasterRecovery/TaskExecutionModeFlag.cs

@@ -0,0 +1,32 @@
+//
+// 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.IntegrationTests.DisasterRecovery
+{
+    [Flags]
+    public enum TaskExecutionModeFlag
+    {
+        None = 0x00,
+
+        /// <summary>
+        /// Execute task
+        /// </summary>
+        Execute = 0x01,
+
+        /// <summary>
+        /// Script task
+        /// </summary>
+        Script = 0x02,
+
+        /// <summary>
+        /// Execute and script task
+        /// Needed for tasks that will show the script when execution completes
+        /// </summary>
+        ExecuteAndScript = Execute | Script
+    }
+}

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/LanguageServer/LanguageServiceTests.cs

@@ -63,12 +63,15 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.LanguageServer
         [Fact]
         public void PrepopulateCommonMetadata()
         {
-            var result = LiveConnectionHelper.InitLiveConnectionInfo();
-            var connInfo = result.ConnectionInfo;
+            using (SelfCleaningTempFile queryTempFile = new SelfCleaningTempFile())
+            {
+                var result = LiveConnectionHelper.InitLiveConnectionInfo("master", queryTempFile.FilePath);
+                var connInfo = result.ConnectionInfo;
 
-            ScriptParseInfo scriptInfo = new ScriptParseInfo { IsConnected = true };
+                ScriptParseInfo scriptInfo = new ScriptParseInfo { IsConnected = true };
 
-            LanguageService.Instance.PrepopulateCommonMetadata(connInfo, scriptInfo, null);
+                LanguageService.Instance.PrepopulateCommonMetadata(connInfo, scriptInfo, null);
+            }
         }
 
         // This test currently requires a live database connection to initialize 

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/TaskServices/RequstParamStub.cs

@@ -0,0 +1,15 @@
+using System;
+using System.Collections.Generic;
+using System.Text;
+using Microsoft.SqlTools.ServiceLayer.TaskServices;
+using Microsoft.SqlTools.ServiceLayer.Utility;
+
+namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.TaskServices
+{
+    public class RequstParamStub : IScriptableRequestParams
+    {
+        public TaskExecutionMode TaskExecutionMode { get; set; }
+        public string OwnerUri { get; set; }
+        public string DatabaseName { get; set; }
+    }
+}

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/TaskServices/SmoScriptableTaskOperationStub.cs

@@ -0,0 +1,52 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.SqlServer.Management.Smo;
+using Microsoft.SqlTools.ServiceLayer.TaskServices;
+
+namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.TaskServices
+{
+    public class SmoScriptableTaskOperationStub : SmoScriptableTaskOperation
+    {
+        private Server server;
+        public string DatabaseName { get; set; }
+        public SmoScriptableTaskOperationStub(Server server)
+        {
+            this.server = server;
+        }
+        public override string ErrorMessage
+        {
+            get
+            {
+                return string.Empty;
+            }
+        }
+
+        public override Server Server
+        {
+            get
+            {
+                return server;
+            }
+        }
+
+        public override void Cancel()
+        {
+        }
+
+        public string TableName { get; set; }
+
+        public override void Execute()
+        {
+            var database = server.Databases[DatabaseName];
+            Table table = new Table(database, TableName, "test");
+            Column column = new Column(table, "c1");
+            column.DataType = DataType.Int;
+            table.Columns.Add(column);
+            database.Tables.Add(table);
+            table.Create();
+        }
+    }
+}

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/TaskServices/TaskServiceTests.cs

@@ -0,0 +1,156 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System.Data.Common;
+using System.Data.SqlClient;
+using System.Linq;
+using System.Threading.Tasks;
+using Microsoft.SqlServer.Management.Common;
+using Microsoft.SqlServer.Management.Smo;
+using Microsoft.SqlTools.Extensibility;
+using Microsoft.SqlTools.Hosting.Protocol;
+using Microsoft.SqlTools.ServiceLayer.Connection;
+using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
+using Microsoft.SqlTools.ServiceLayer.IntegrationTests.Utility;
+using Microsoft.SqlTools.ServiceLayer.TaskServices;
+using Microsoft.SqlTools.ServiceLayer.TaskServices.Contracts;
+using Microsoft.SqlTools.ServiceLayer.Test.Common;
+using Microsoft.SqlTools.ServiceLayer.UnitTests;
+using Microsoft.SqlTools.ServiceLayer.UnitTests.Utility;
+using Moq;
+using Xunit;
+using static Microsoft.SqlTools.ServiceLayer.IntegrationTests.Utility.LiveConnectionHelper;
+
+namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.TaskServices
+{
+    public class TaskServiceTests : ServiceTestBase
+    {
+        private TaskService service;
+        private Mock<IProtocolEndpoint> serviceHostMock;
+
+        public TaskServiceTests()
+        {
+            serviceHostMock = new Mock<IProtocolEndpoint>();
+            service = CreateService();
+            service.InitializeService(serviceHostMock.Object);
+        }
+
+        [Fact]
+        public async Task VerifyTaskExecuteTheQueryGivenExecutionModeExecute()
+        {
+            await VerifyTaskWithExecutionMode(TaskExecutionMode.Execute);
+        }
+
+        [Fact]
+        public async Task VerifyTaskGenerateScriptOnlyGivenExecutionModeScript()
+        {
+            await VerifyTaskWithExecutionMode(TaskExecutionMode.Script);
+        }
+
+        [Fact]
+        public async Task VerifyTaskNotExecuteAndGenerateScriptGivenExecutionModeExecuteAndScript()
+        {
+            await VerifyTaskWithExecutionMode(TaskExecutionMode.ExecuteAndScript);
+        }
+
+        [Fact]
+        public async Task VerifyTaskSendsFailureNotificationGivenInvalidQuery()
+        {
+            await VerifyTaskWithExecutionMode(TaskExecutionMode.ExecuteAndScript, true);
+        }
+
+        public async Task VerifyTaskWithExecutionMode(TaskExecutionMode executionMode, bool makeTaskFail = false)
+        {
+            serviceHostMock.AddEventHandling(TaskStatusChangedNotification.Type, null);
+
+            using (SelfCleaningTempFile queryTempFile = new SelfCleaningTempFile())
+            {
+                //To make the task fail don't create the schema so create table fails
+                string query = string.Empty;
+                if (!makeTaskFail)
+                {
+                    query = $"CREATE SCHEMA [test]";
+                }
+                SqlTestDb testDb = await SqlTestDb.CreateNewAsync(TestServerType.OnPrem, false, null, query, "TaskService");
+                try
+                {
+                    TestConnectionResult connectionResult = await LiveConnectionHelper.InitLiveConnectionInfoAsync(testDb.DatabaseName, queryTempFile.FilePath);
+                    string taskName = "task name";
+                    Server server = CreateServerObject(connectionResult.ConnectionInfo);
+                    RequstParamStub requstParam = new RequstParamStub
+                    {
+                        TaskExecutionMode = executionMode,
+                        OwnerUri = queryTempFile.FilePath
+                    };
+                    SmoScriptableTaskOperationStub taskOperation = new SmoScriptableTaskOperationStub(server);
+                    taskOperation.DatabaseName = testDb.DatabaseName;
+                    taskOperation.TableName = "newTable";
+                    TaskMetadata taskMetadata = TaskMetadata.Create(requstParam, taskName, taskOperation, ConnectionService.Instance);
+                    SqlTask sqlTask = service.TaskManager.CreateTask<SqlTask>(taskMetadata);
+                    Task taskToVerify = sqlTask.RunAsync().ContinueWith(task =>
+                    {
+                        if (!makeTaskFail)
+                        {
+                            if (executionMode == TaskExecutionMode.Script || executionMode == TaskExecutionMode.ExecuteAndScript)
+                            {
+                                serviceHostMock.Verify(x => x.SendEvent(TaskStatusChangedNotification.Type,
+                                           It.Is<TaskProgressInfo>(t => !string.IsNullOrEmpty(t.Script))), Times.AtLeastOnce());
+                            }
+
+                            //Verify if the table created if execution mode includes execute
+                            bool expected = executionMode == TaskExecutionMode.Execute || executionMode == TaskExecutionMode.ExecuteAndScript;
+                            Server serverToverfiy = CreateServerObject(connectionResult.ConnectionInfo);
+                            bool actual = serverToverfiy.Databases[testDb.DatabaseName].Tables.Contains(taskOperation.TableName, "test");
+                            Assert.Equal(expected, actual);
+                        }
+                        else
+                        {
+                            serviceHostMock.Verify(x => x.SendEvent(TaskStatusChangedNotification.Type,
+                                           It.Is<TaskProgressInfo>(t => t.Status == SqlTaskStatus.Failed)), Times.AtLeastOnce());
+                        }
+                    });
+                    await taskToVerify;
+                }
+                finally
+                {
+                    testDb.Cleanup();
+                }
+
+            }
+        }
+        protected TaskService CreateService()
+        {
+            CreateServiceProviderWithMinServices();
+
+            // Create the service using the service provider, which will initialize dependencies
+            return ServiceProvider.GetService<TaskService>();
+        }
+
+        protected override RegisteredServiceProvider CreateServiceProviderWithMinServices()
+        {
+            TaskService service = new TaskService();
+            service.TaskManager = new SqlTaskManager();
+            return CreateProvider()
+                .RegisterSingleService(service);
+        }
+
+        private Server CreateServerObject(ConnectionInfo connInfo   )
+        {
+            SqlConnection connection = null;
+            DbConnection dbConnection = connInfo.AllConnections.First();
+            ReliableSqlConnection reliableSqlConnection = dbConnection as ReliableSqlConnection;
+            SqlConnection sqlConnection = dbConnection as SqlConnection;
+            if (reliableSqlConnection != null)
+            {
+                connection = reliableSqlConnection.GetUnderlyingConnection();
+            }
+            else if (sqlConnection != null)
+            {
+                connection = sqlConnection;
+            }
+            return new Server(new ServerConnection(connection));
+        }
+    }
+}