Feature/sqlcmd : Enable running scripts with SQLCMD variables - Part 1

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

@@ -14,7 +14,7 @@ public class GeneralRequestDetails
     {
         public GeneralRequestDetails()
         {
-            Options = new Dictionary<string, object>();
+            Options = new Dictionary<string, object>(new OptionsKeyComparer());
         }
 
         public T GetOptionValue<T>(string name, T defaultValue = default(T))
@@ -111,4 +111,21 @@ protected void SetOptionValue<T>(string name, T value)
         /// </summary>
         public Dictionary<string, object> Options { get; set; }
     }
+
+    /// <summary>
+    /// Compare class for Execution Setting options to ensure ignoring case in dictionary keys
+    /// </summary>
+    public class OptionsKeyComparer : IEqualityComparer<string>
+    {
+        public bool Equals(string x, string y)
+        {
+            return string.Equals(x, y, StringComparison.InvariantCultureIgnoreCase);
+        }
+
+        public int GetHashCode(string obj)
+        {
+            return obj.ToLowerInvariant().GetHashCode();
+        }
+    }
 }
+

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

@@ -51,7 +51,7 @@ private static List<BatchDefinition> ConvertToBatchDefinitionList(string content
                     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; // content.Substring(offset, batchInfos[0].length);
 
                     // if there's only one batch then the line difference is just startLine
                     if (count > 1)
@@ -82,15 +82,15 @@ private static List<BatchDefinition> ConvertToBatchDefinitionList(string content
                     }
 
                     // 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);
-                        startLine = batchInfos[index].startLine;
+                        batchText = batchInfos[index].batchText;
+                        startLine = batchInfos[index].startLine + 1; // this is done for first batch but was missing for following batches
                         startColumn = batchInfos[index].startColumn;
 
                         // make a new batch definition for each batch
@@ -109,7 +109,7 @@ private static List<BatchDefinition> ConvertToBatchDefinitionList(string content
                     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 @@ private static List<int> GetOffsets(string content, IList<BatchInfo> batchInfos)
         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 @@ public BatchParserWrapper()
         /// <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 @@ private void OnBatchParserExecutionFinished(object sender, BatchParserExecutionF
                     }
 
                     // 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, batchTextLength, batchText, args.Batch.ExpectedExecutionCount));
                 }
             }
             catch (NotImplementedException)
@@ -474,17 +474,19 @@ private void Dispose(bool disposing)
 
         private class BatchInfo
         {
-            public BatchInfo(int startLine, int startColumn, int length, int repeatCount = 1)
+            public BatchInfo(int startLine, int startColumn, int length, 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 @@ private void ConfigureBatchParser()
             batchParser.HaltParser = new BatchParser.HaltParserDelegate(OnHaltParser);
             batchParser.StartingLine = startingLine;
 
-            if (isLocalParse)
+            if (isLocalParse && !sqlCmdMode)
             {
                 batchParser.DisableVariableSubstitution();
             }
@@ -1045,23 +1045,27 @@ public void ExecuteBatch(ScriptExecutionArgs scriptExecutionArgs)
         /// 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 @@ public static LanguageService Instance
 
         #endregion
 
-        #region Private / internal instance fields and constructor
-        private const int PrepopulateBindTimeout = 60000;
+        #region Instance fields and constructor
 
         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 @@ public static LanguageService Instance
 
         internal const int CompletionExtTimeout = 200;
 
+        // For testability only
+        internal Task DelayedDiagnosticsTask = null;
+
         private ConnectionService connectionService = null;
 
         private WorkspaceService<SqlToolsSettings> workspaceServiceInstance;
@@ -836,7 +843,10 @@ public async Task HandleDidChangeTextDocumentNotification(ScriptFile[] changedFi
                 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 @@ public async Task HandleDidChangeTextDocumentNotification(ScriptFile[] changedFi
                 {
                     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 @@ private Task RunScriptDiagnostics(ScriptFile[] filesToAnalyze, EventContext even
             existingRequestCancellation = new CancellationTokenSource();
             Task.Factory.StartNew(
                 () =>
-                    DelayThenInvokeDiagnostics(
+                    this.DelayedDiagnosticsTask = DelayThenInvokeDiagnostics(
                         LanguageService.DiagnosticParseDelay,
                         filesToAnalyze,
                         eventContext,

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

@@ -116,7 +116,13 @@ public class Query : IDisposable
 
             // Process the query into batches 
             BatchParserWrapper parser = new BatchParserWrapper();
-            List<BatchDefinition> parserResult = parser.GetBatches(queryText);
+            ExecutionEngineConditions conditions = null;
+            if (settings.IsSqlCmdMode)
+            {
+                conditions = new ExecutionEngineConditions();
+                conditions.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 @@ public class QueryExecutionSettings : GeneralRequestDetails
         /// </summary>
         private const int DefaultQueryGovernorCostLimit = 0;
 
+        /// <summary>
+        /// Flag to run query in sqlcmd mode
+        /// </summary>
+        private bool DefaultSqlCmdMode = false;
+
         #endregion
 
         #region Member Variables
@@ -633,6 +638,21 @@ public bool IncludeEstimatedExecutionPlanXml
             }
         }
 
+        /// <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 @@ public void Update(QueryExecutionSettings newSettings)
             AnsiPadding = newSettings.AnsiPadding;
             AnsiWarnings = newSettings.AnsiWarnings;
             AnsiNulls = newSettings.AnsiNulls;
+            IsSqlCmdMode = newSettings.IsSqlCmdMode;
         }
 
         #endregion

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

@@ -338,6 +338,41 @@ public void VerifyIsSqlCmd()
             }
         }
 
+        // Verify whether the batchParser execute SqlCmd successfully
+        [Fact]
+        public void VerifyRunSqlCmd()
+        {
+            using (ExecutionEngine executionEngine = new ExecutionEngine())
+            {
+                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();
+                    condition.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 @@ public TestExecutor(string batch, SqlConnection conn, ExecutionEngineConditions
             conditions.IsNoExec = exeCondition.IsNoExec;
             conditions.IsStatisticsIO = exeCondition.IsStatisticsIO;
             conditions.IsStatisticsTime = exeCondition.IsStatisticsTime;
+            conditions.IsSqlCmd = exeCondition.IsSqlCmd;
 
             _cancel = cancelExecution;
             connection = conn;