Feature/sqlcmd : Enable running scripts with SQLCMD variables - Part 1 (#839)

* Part1 : Changes to make cmdcmd script to work with parameters in script

* Stop SQL intellisense for SQLCMD

* Adding test for Intellisense handling of SQLCMD page

* Removing unintentional spacing changes caused by formatting

* Updating with smaller CR comments. Will discuss regarding script vs other options in batch info

* Removing unintentional change

* Adding latest PR comments

src/Microsoft.SqlTools.Hosting/Utility/GeneralRequestDetails.cs

@@ -14,7 +14,7 @@ namespace Microsoft.SqlTools.Utility
     {
         public GeneralRequestDetails()
         {
-            Options = new Dictionary<string, object>();
+            Options = new Dictionary<string, object>(StringComparer.InvariantCultureIgnoreCase);
         }
 
         public T GetOptionValue<T>(string name, T defaultValue = default(T))
@@ -112,3 +112,4 @@ namespace Microsoft.SqlTools.Utility
         public Dictionary<string, object> Options { get; set; }
     }
 }
+

src/Microsoft.SqlTools.ManagedBatchParser/BatchParser/BatchParserWrapper.cs

@@ -51,7 +51,7 @@ namespace Microsoft.SqlTools.ServiceLayer.BatchParser
                     int offset = offsets[0];
                     int startColumn = batchInfos[0].startColumn;
                     int count = batchInfos.Count;
-                    string batchText = content.Substring(offset, batchInfos[0].length);
+                    string batchText = batchInfos[0].batchText;
 
                     // if there's only one batch then the line difference is just startLine
                     if (count > 1)
@@ -82,15 +82,15 @@ namespace Microsoft.SqlTools.ServiceLayer.BatchParser
                     }
 
                     // Generate the rest batch definitions
-                    for (int index = 1; index < count - 1; index++)
+                    for (int index = 1; index < count - 1 ; index++)
                     {
                         lineDifference = batchInfos[index + 1].startLine - batchInfos[index].startLine;
                         position = ReadLines(reader, lineDifference, endLine);
                         endLine = position.Item1;
                         endColumn = position.Item2;
                         offset = offsets[index];
-                        batchText = content.Substring(offset, batchInfos[index].length);
+                        batchText = batchInfos[index].batchText;
-                        startLine = batchInfos[index].startLine;
+                        startLine = batchInfos[index].startLine + 1; //positions is 0 index based
                         startColumn = batchInfos[index].startColumn;
 
                         // make a new batch definition for each batch
@@ -109,7 +109,7 @@ namespace Microsoft.SqlTools.ServiceLayer.BatchParser
                     if (count > 1)
                     {
 
-                        batchText = content.Substring(offsets[count-1], batchInfos[count - 1].length);
+                        batchText = batchInfos[count - 1].batchText;
                         BatchDefinition lastBatchDef = GetLastBatchDefinition(reader, batchInfos[count - 1], batchText);
                         batchDefinitionList.Add(lastBatchDef);
                     }
@@ -209,7 +209,7 @@ namespace Microsoft.SqlTools.ServiceLayer.BatchParser
         private static BatchDefinition GetLastBatchDefinition(StringReader reader,
             BatchInfo batchInfo, string batchText)
         {
-            int startLine = batchInfo.startLine;
+            int startLine = batchInfo.startLine + 1;
             int startColumn = batchInfo.startColumn;
             string prevLine = null;
             string line = reader.ReadLine();
@@ -328,12 +328,12 @@ namespace Microsoft.SqlTools.ServiceLayer.BatchParser
         /// <summary>
         /// Takes in a query string and returns a list of BatchDefinitions
         /// </summary>
-        public List<BatchDefinition> GetBatches(string sqlScript)
+        public List<BatchDefinition> GetBatches(string sqlScript, ExecutionEngineConditions conditions = null)
         {
             batchInfos = new List<BatchInfo>();
 
             // execute the script - all communication / integration after here happen via event handlers
-            executionEngine.ParseScript(sqlScript, notificationHandler);
+            executionEngine.ParseScript(sqlScript, notificationHandler, conditions);
 
             // retrieve a list of BatchDefinitions 
             List<BatchDefinition> batchDefinitionList = ConvertToBatchDefinitionList(sqlScript, batchInfos);
@@ -381,7 +381,7 @@ namespace Microsoft.SqlTools.ServiceLayer.BatchParser
                     }
 
                     // Add the script info
-                    batchInfos.Add(new BatchInfo(args.Batch.TextSpan.iStartLine, args.Batch.TextSpan.iStartIndex, batchTextLength, args.Batch.ExpectedExecutionCount));
+                    batchInfos.Add(new BatchInfo(args.Batch.TextSpan.iStartLine, args.Batch.TextSpan.iStartIndex, batchText, args.Batch.ExpectedExecutionCount));
                 }
             }
             catch (NotImplementedException)
@@ -474,17 +474,17 @@ namespace Microsoft.SqlTools.ServiceLayer.BatchParser
 
         private class BatchInfo
         {
-            public BatchInfo(int startLine, int startColumn, int length, int repeatCount = 1)
+            public BatchInfo(int startLine, int startColumn, string batchText, int repeatCount = 1)
             {
                 this.startLine = startLine;
                 this.startColumn = startColumn;
-                this.length = length;
                 this.executionCount = repeatCount;
+                this.batchText = batchText;
             }
             public int startLine;
             public int startColumn;
-            public int length;
             public int executionCount;
+            public string batchText;
         }
 
     }

src/Microsoft.SqlTools.ManagedBatchParser/BatchParser/ExecutionEngineCode/ExecutionEngine.cs

@@ -200,7 +200,7 @@ namespace Microsoft.SqlTools.ServiceLayer.BatchParser.ExecutionEngineCode
             batchParser.HaltParser = new BatchParser.HaltParserDelegate(OnHaltParser);
             batchParser.StartingLine = startingLine;
 
-            if (isLocalParse)
+            if (isLocalParse && !sqlCmdMode)
             {
                 batchParser.DisableVariableSubstitution();
             }
@@ -1045,23 +1045,27 @@ namespace Microsoft.SqlTools.ServiceLayer.BatchParser.ExecutionEngineCode
         /// Parses the script locally
         /// </summary>
         /// <param name="script">script to parse</param>
-        /// <param name="batchEventsHandler">batch handler</param>        
+        /// <param name="batchEventsHandler">batch handler</param>   
+        /// <param name="conditions">execution engine conditions if specified</param>
         /// <remarks>
         /// The batch parser functionality is used in this case
         /// </remarks>
-        public void ParseScript(string script, IBatchEventsHandler batchEventsHandler)
+        public void ParseScript(string script, IBatchEventsHandler batchEventsHandler, ExecutionEngineConditions conditions = null)
         {
             Validate.IsNotNull(nameof(script), script);
             Validate.IsNotNull(nameof(batchEventsHandler), batchEventsHandler);
 
-
+            if (conditions != null)
+            {
+                this.conditions = conditions;
+            }
             this.script = script;
             batchEventHandlers = batchEventsHandler;
             isLocalParse = true;
 
             DoExecute(/* isBatchParser */ true);
         }
-                        
+
         /// <summary>
         /// Close the current connection
         /// </summary>

src/Microsoft.SqlTools.ServiceLayer/LanguageServices/LanguageService.cs

@@ -58,12 +58,16 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 
         #endregion
 
-        #region Private / internal instance fields and constructor
+        #region Instance fields and constructor
-        private const int PrepopulateBindTimeout = 60000;
 
         public const string SQL_LANG = "SQL";
+
+        public const string SQL_CMD_LANG = "SQLCMD";
+
         private const int OneSecond = 1000;
 
+        private const int PrepopulateBindTimeout = 60000;
+
         internal const string DefaultBatchSeperator = "GO";
 
         internal const int DiagnosticParseDelay = 750;
@@ -80,6 +84,9 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
 
         internal const int CompletionExtTimeout = 200;
 
+        // For testability only
+        internal Task DelayedDiagnosticsTask = null;
+
         private ConnectionService connectionService = null;
 
         private WorkspaceService<SqlToolsSettings> workspaceServiceInstance;
@@ -836,7 +843,10 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                 if (SQL_LANG.Equals(changeParams.Language, StringComparison.OrdinalIgnoreCase)) {
                     shouldBlock = !ServiceHost.ProviderName.Equals(changeParams.Flavor, StringComparison.OrdinalIgnoreCase);
                 }
-
+                if (SQL_CMD_LANG.Equals(changeParams.Language, StringComparison.OrdinalIgnoreCase))
+                {
+                    shouldBlock = true; // the provider will continue to be mssql
+                }
                 if (shouldBlock) {
                     this.nonMssqlUriMap.AddOrUpdate(changeParams.Uri, true, (k, oldValue) => true);
                     if (CurrentWorkspace.ContainsFile(changeParams.Uri))
@@ -848,7 +858,10 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
                 {
                     bool value;
                     this.nonMssqlUriMap.TryRemove(changeParams.Uri, out value);
-                }
+                    // should rebuild intellisense when re-considering as sql 
+                    RebuildIntelliSenseParams param = new RebuildIntelliSenseParams { OwnerUri = changeParams.Uri };
+                    await HandleRebuildIntelliSenseNotification(param, eventContext);              
+                }                
             }
             catch (Exception ex)
             {
@@ -1733,7 +1746,7 @@ namespace Microsoft.SqlTools.ServiceLayer.LanguageServices
             existingRequestCancellation = new CancellationTokenSource();
             Task.Factory.StartNew(
                 () =>
-                    DelayThenInvokeDiagnostics(
+                    this.DelayedDiagnosticsTask = DelayThenInvokeDiagnostics(
                         LanguageService.DiagnosticParseDelay,
                         filesToAnalyze,
                         eventContext,

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

@@ -116,7 +116,12 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
             // Process the query into batches 
             BatchParserWrapper parser = new BatchParserWrapper();
-            List<BatchDefinition> parserResult = parser.GetBatches(queryText);
+            ExecutionEngineConditions conditions = null;
+            if (settings.IsSqlCmdMode)
+            {
+                conditions = new ExecutionEngineConditions() { IsSqlCmd = settings.IsSqlCmdMode };
+            }
+            List<BatchDefinition> parserResult = parser.GetBatches(queryText, conditions);
 
             var batchSelection = parserResult
                 .Select((batchDefinition, index) =>

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

@@ -163,6 +163,11 @@ namespace Microsoft.SqlTools.ServiceLayer.SqlContext
         /// </summary>
         private const int DefaultQueryGovernorCostLimit = 0;
 
+        /// <summary>
+        /// Default value for flag to run query in sqlcmd mode
+        /// </summary>
+        private bool DefaultSqlCmdMode = false;
+
         #endregion
 
         #region Member Variables
@@ -633,6 +638,21 @@ namespace Microsoft.SqlTools.ServiceLayer.SqlContext
             }
         }
 
+        /// <summary>
+        /// Set sqlCmd Mode
+        /// </summary>
+        public bool IsSqlCmdMode
+        {
+            get
+            {
+                return GetOptionValue<bool>("isSqlCmdMode", DefaultSqlCmdMode);
+            }
+            set
+            {
+                SetOptionValue("isSqlCmdMode", value);
+            }
+        }
+
         #endregion
 
         #region Public Methods
@@ -670,6 +690,7 @@ namespace Microsoft.SqlTools.ServiceLayer.SqlContext
             AnsiPadding = newSettings.AnsiPadding;
             AnsiWarnings = newSettings.AnsiWarnings;
             AnsiNulls = newSettings.AnsiNulls;
+            IsSqlCmdMode = newSettings.IsSqlCmdMode;
         }
 
         #endregion

test/Microsoft.SqlTools.ManagedBatchParser.IntegrationTests/BatchParser/BatchParserTests.cs

@@ -338,6 +338,42 @@ namespace Microsoft.SqlTools.ManagedBatchParser.UnitTests.BatchParser
             }
         }
 
+        /// <summary>
+        /// Verify whether the batchParser execute SqlCmd successfully
+        /// </summary>
+        [Fact]
+        public void VerifyRunSqlCmd()
+        {
+            using (ExecutionEngine executionEngine = new ExecutionEngine())
+            {
+                const string sqlCmdQuery = @"
+:setvar __var1 1
+:setvar __var2 2
+:setvar __IsSqlCmdEnabled " + "\"True\"" + @"
+GO
+IF N'$(__IsSqlCmdEnabled)' NOT LIKE N'True'
+    BEGIN
+        PRINT N'SQLCMD mode must be enabled to successfully execute this script.';
+                SET NOEXEC ON;
+                END
+GO
+select $(__var1) + $(__var2) as col
+GO";
+
+                using (SqlConnection con = new SqlConnection(CONNECTION_STRING))
+                {
+                    con.Open();
+                    var condition = new ExecutionEngineConditions() { IsSqlCmd = true };
+                    TestExecutor testExecutor = new TestExecutor(sqlCmdQuery, con, condition);
+                    testExecutor.Run();
+                    
+                    Assert.True(testExecutor.ResultCountQueue.Count >= 1);
+                    Assert.True(testExecutor.ErrorMessageQueue.Count == 0);
+                    
+                }
+            }
+        }
+
         // Verify whether the executionEngine execute Batch
         [Fact]
         public void VerifyExecuteBatch()

test/Microsoft.SqlTools.ManagedBatchParser.IntegrationTests/TSQLExecutionEngine/TestExecutor.cs

@@ -231,6 +231,7 @@ namespace Microsoft.SqlTools.ManagedBatchParser.IntegrationTests.TSQLExecutionEn
             conditions.IsNoExec = exeCondition.IsNoExec;
             conditions.IsStatisticsIO = exeCondition.IsStatisticsIO;
             conditions.IsStatisticsTime = exeCondition.IsStatisticsTime;
+            conditions.IsSqlCmd = exeCondition.IsSqlCmd;
 
             _cancel = cancelExecution;
             connection = conn;

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

@@ -17,6 +17,10 @@ using Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Test.Common;
 using Microsoft.SqlTools.ServiceLayer.UnitTests.ServiceHost;
 using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
+using Microsoft.SqlTools.ServiceLayer.Workspace;
+using Microsoft.SqlTools.ServiceLayer.SqlContext;
+using Microsoft.SqlTools.ServiceLayer.Connection;
+using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
 using Moq;
 using Xunit;
 
@@ -329,5 +333,73 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.LanguageServer
                 testDb.Cleanup();
             }
         }
+
+        /// <summary>
+        // This test validates switching off editor intellisesnse for now. 
+        // Will change to better handling once we have specific SQLCMD intellisense in Language Service
+        /// </summary>
+        [Fact]
+        public async Task HandleRequestToChangeToSqlcmdFile()
+        {
+
+            var scriptFile = new ScriptFile() { ClientFilePath = "HandleRequestToChangeToSqlcmdFile_" + DateTime.Now.ToLongDateString() + "_.sql" };
+
+            try
+            {
+                // Prepare a script file
+                scriptFile.SetFileContents("koko wants a bananas");
+                File.WriteAllText(scriptFile.ClientFilePath, scriptFile.Contents);
+
+                // Create a workspace and add file to it so that its found for intellense building
+                var workspace = new ServiceLayer.Workspace.Workspace();
+                var workspaceService = new WorkspaceService<SqlToolsSettings> { Workspace = workspace };
+                var langService = new LanguageService() { WorkspaceServiceInstance = workspaceService }; 
+                langService.CurrentWorkspace.GetFile(scriptFile.ClientFilePath);
+                langService.CurrentWorkspaceSettings.SqlTools.IntelliSense.EnableIntellisense = true;
+
+                // Add a connection to ensure the intellisense building works            
+                ConnectionInfo connectionInfo = GetLiveAutoCompleteTestObjects().ConnectionInfo;
+                langService.ConnectionServiceInstance.OwnerToConnectionMap.Add(scriptFile.ClientFilePath, connectionInfo);
+
+                // Test SQL
+                int countOfValidationCalls = 0;
+                var eventContextSql = new Mock<SqlTools.Hosting.Protocol.EventContext>();
+                eventContextSql.Setup(x => x.SendEvent(PublishDiagnosticsNotification.Type, It.Is<PublishDiagnosticsNotification>((notif) => ValidateNotification(notif, 2, ref countOfValidationCalls)))).Returns(Task.FromResult(new object()));
+                await langService.HandleDidChangeLanguageFlavorNotification(new LanguageFlavorChangeParams
+                {
+                    Uri = scriptFile.ClientFilePath,
+                    Language = LanguageService.SQL_LANG.ToLower(),
+                    Flavor = "MSSQL"
+                }, eventContextSql.Object);
+                await langService.DelayedDiagnosticsTask; // to ensure completion and validation before moveing to next step
+
+                // Test SQL CMD
+                var eventContextSqlCmd = new Mock<SqlTools.Hosting.Protocol.EventContext>();
+                eventContextSqlCmd.Setup(x => x.SendEvent(PublishDiagnosticsNotification.Type, It.Is<PublishDiagnosticsNotification>((notif) => ValidateNotification(notif, 0, ref countOfValidationCalls)))).Returns(Task.FromResult(new object()));
+                await langService.HandleDidChangeLanguageFlavorNotification(new LanguageFlavorChangeParams
+                {
+                    Uri = scriptFile.ClientFilePath,
+                    Language = LanguageService.SQL_CMD_LANG.ToLower(),
+                    Flavor = "MSSQL"
+                }, eventContextSqlCmd.Object);
+                await langService.DelayedDiagnosticsTask;
+
+                Assert.True(countOfValidationCalls == 2, $"Validation should be called 2 time but is called {countOfValidationCalls} times");
+            }
+            finally
+            {
+                if (File.Exists(scriptFile.ClientFilePath))
+                {
+                    File.Delete(scriptFile.ClientFilePath);
+                }
+            }
+        }
+
+        private bool ValidateNotification(PublishDiagnosticsNotification notif, int errors, ref int countOfValidationCalls)
+        {
+            countOfValidationCalls++;
+            Assert.True(notif.Diagnostics.Length == errors, $"Notification errors {notif.Diagnostics.Length} are not as expected {errors}");
+            return true;
+        }
     }
 }