Kusto Auth Refactor Tests (#1148)

* Refactored Kusto.ServiceLayer to pass ConnectionDetails to DataSourceFactory instead of connection string. Created KustoConnectionDetails to map needed details to KustoClient.

* Removed unused ScriptingScriptOperation from KustoServiceLayer.

* Created DstsAuthenticationManager and moved logic for getting DstsToken. Updated error message for failing to create KustoConnection.

* Removed DstsAuthenticationManager.cs. Refactored DataSourceFactory to retrieve UserToken from ConnectionDetails.

* Renamed AzureAccountToken in ConnectionDetails to AccountToken. Changed mapping to KustoConnectionDetails based on the AccountToken.

* Removed Kusto.Data reference from ConnectionService and ScriptingListObjectsOperation. Moved creation of KustoConnectionStringBuilder to DataSourceFactory

* Added accountToken validation to DataSourceFactory Create.

* Renamed KustoConnectionDetails to DataSourceConnectionDetails. Renamed AzureToken to AuthToken.

* Refactored SchemaState and intellisense out of KustoClient to KustoIntellisenseClient. Added IIntellisenseClient. Added unit tests for KustoIntellisenseClient and KustoClient.

* Removed unused property dataSourceFactory from LanguageService > InitializeService. Moved KustoIntellisense functions from KustoIntellisenseHelper to KustoIntellisenseClient and made SchemaState private. Added IIntellisenseClient to IDataSource.

* Renamed directory from DataSourceIntellisense to Intellisense and updated namespace. Fixed namespace in ScriptDocumentInfo.

src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceBase.cs

@@ -7,9 +7,12 @@ using System.Collections.Generic;
 using System.Threading;
 using System.Data;
 using System.Threading.Tasks;
-using Kusto.Language;
+using Microsoft.Kusto.ServiceLayer.DataSource.Intellisense;
 using Microsoft.Kusto.ServiceLayer.Utility;
 using Microsoft.Kusto.ServiceLayer.DataSource.Metadata;
+using Microsoft.Kusto.ServiceLayer.LanguageServices;
+using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
+using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
 
 namespace Microsoft.Kusto.ServiceLayer.DataSource
 {
@@ -88,6 +91,14 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         public abstract string GenerateAlterFunctionScript(string functionName);
 
         public abstract string GenerateExecuteFunctionScript(string functionName);
+        public abstract ScriptFileMarker[] GetSemanticMarkers(ScriptParseInfo parseInfo, ScriptFile scriptFile, string queryText);
+
+        public abstract DefinitionResult GetDefinition(string queryText, int index, int startLine, int startColumn, bool throwOnError = false);
+
+        public abstract Hover GetHoverHelp(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, bool throwOnError = false);
+
+        public abstract CompletionItem[] GetAutoCompleteSuggestions(ScriptDocumentInfo scriptDocumentInfo, Position textPosition,
+            bool throwOnError = false);
 
         /// <inheritdoc/>
         public DataSourceType DataSourceType { get; protected set; }
@@ -96,7 +107,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         public abstract string ClusterName { get; }
 
         public abstract string DatabaseName { get; }
-        public abstract GlobalState SchemaState { get; }
 
         #endregion
     }

src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceFactory.cs

@@ -5,10 +5,10 @@ using Kusto.Data;
 using Microsoft.Kusto.ServiceLayer.Connection.Contracts;
 using Microsoft.Kusto.ServiceLayer.DataSource.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
-using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
+using Microsoft.Kusto.ServiceLayer.DataSource.Intellisense;
+using Microsoft.Kusto.ServiceLayer.LanguageServices;
 using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
 using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
-using Microsoft.Kusto.ServiceLayer.LanguageServices.Completion;
 using Microsoft.Kusto.ServiceLayer.Utility;
 
 namespace Microsoft.Kusto.ServiceLayer.DataSource
@@ -26,7 +26,8 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
                 {
                     var kustoConnectionDetails = MapKustoConnectionDetails(connectionDetails);
                     var kustoClient = new KustoClient(kustoConnectionDetails, ownerUri);
-                    return new KustoDataSource(kustoClient);
+                    var intellisenseClient = new KustoIntellisenseClient(kustoClient);
+                    return new KustoDataSource(kustoClient, intellisenseClient);
                 }
 
                 default:

src/Microsoft.Kusto.ServiceLayer/DataSource/IDataSource.cs

@@ -3,8 +3,11 @@ using System.Collections.Generic;
 using System.Data;
 using System.Threading;
 using System.Threading.Tasks;
-using Kusto.Language;
+using Microsoft.Kusto.ServiceLayer.DataSource.Intellisense;
 using Microsoft.Kusto.ServiceLayer.DataSource.Metadata;
+using Microsoft.Kusto.ServiceLayer.LanguageServices;
+using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
+using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
 
 namespace Microsoft.Kusto.ServiceLayer.DataSource
 {
@@ -27,8 +30,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// The current database name, if there is one.
         /// </summary>
         string DatabaseName { get; }
-        
-        GlobalState SchemaState { get; }
 
         /// <summary>
         /// Executes a query.
@@ -111,5 +112,10 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// <param name="functionName"></param>
         /// <returns></returns>
         string GenerateExecuteFunctionScript(string functionName);
+        
+        ScriptFileMarker[] GetSemanticMarkers(ScriptParseInfo parseInfo, ScriptFile scriptFile, string queryText);
+        DefinitionResult GetDefinition(string queryText, int index, int startLine, int startColumn, bool throwOnError = false);
+        Hover GetHoverHelp(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, bool throwOnError = false);
+        CompletionItem[] GetAutoCompleteSuggestions(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, bool throwOnError = false);
     }
 }

src/Microsoft.Kusto.ServiceLayer/DataSource/IKustoClient.cs

@@ -2,17 +2,11 @@ using System.Collections.Generic;
 using System.Data;
 using System.Threading;
 using System.Threading.Tasks;
-using Kusto.Language;
 
 namespace Microsoft.Kusto.ServiceLayer.DataSource
 {
     public interface IKustoClient
     {
-        /// <summary>
-        /// SchemaState used for getting intellisense info.
-        /// </summary>
-        GlobalState SchemaState { get; }
-
         string ClusterName { get; }
         
         string DatabaseName { get; }

src/Microsoft.Kusto.ServiceLayer/DataSource/Intellisense/IIntellisenseClient.cs

@@ -0,0 +1,15 @@
+using Microsoft.Kusto.ServiceLayer.LanguageServices;
+using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
+using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
+
+namespace Microsoft.Kusto.ServiceLayer.DataSource.Intellisense
+{
+    public interface IIntellisenseClient
+    {
+        void UpdateDatabase(string databaseName);
+        ScriptFileMarker[] GetSemanticMarkers(ScriptParseInfo parseInfo, ScriptFile scriptFile, string queryText);
+        DefinitionResult GetDefinition(string queryText, int index, int startLine, int startColumn, bool throwOnError = false);
+        Hover GetHoverHelp(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, bool throwOnError = false);
+        CompletionItem[] GetAutoCompleteSuggestions(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, bool throwOnError = false);
+    }
+}

src/Microsoft.Kusto.ServiceLayer/DataSource/Intellisense/KustoIntellisenseClient.cs

@@ -1,32 +1,141 @@
-//
-// 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.Collections.Generic;
 using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
 using Kusto.Language;
-using KustoDiagnostic = Kusto.Language.Diagnostic;
 using Kusto.Language.Editor;
-using Kusto.Language.Syntax;
 using Kusto.Language.Symbols;
+using Kusto.Language.Syntax;
 using Microsoft.Kusto.ServiceLayer.LanguageServices;
 using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
-using Microsoft.Kusto.ServiceLayer.LanguageServices.Completion;
 using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
+using Diagnostic = Kusto.Language.Diagnostic;
 
-namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
+namespace Microsoft.Kusto.ServiceLayer.DataSource.Intellisense
 {
-    /// <summary>
+    public class KustoIntellisenseClient : IIntellisenseClient
-    /// Kusto specific class for intellisense helper functions.
-    /// </summary>
-    public class KustoIntellisenseHelper
     {
+        private readonly IKustoClient _kustoClient;
+
+        /// <summary>
+        /// SchemaState used for getting intellisense info.
+        /// </summary>
+        private GlobalState _schemaState;
+
+        public KustoIntellisenseClient(IKustoClient kustoClient)
+        {
+            _kustoClient = kustoClient;
+            _schemaState = LoadSchemaState(kustoClient.DatabaseName, kustoClient.ClusterName);
+        }
+        
+        public void UpdateDatabase(string databaseName)
+        {
+            _schemaState = LoadSchemaState(databaseName, _kustoClient.ClusterName);
+        }
+        
+        private GlobalState LoadSchemaState(string databaseName, string clusterName)
+        {
+            IEnumerable<ShowDatabaseSchemaResult> tableSchemas = Enumerable.Empty<ShowDatabaseSchemaResult>();
+            IEnumerable<ShowFunctionsResult> functionSchemas = Enumerable.Empty<ShowFunctionsResult>();
+
+            if (!string.IsNullOrWhiteSpace(databaseName))
+            {
+                var source = new CancellationTokenSource();
+                Parallel.Invoke(() =>
+                    {
+                        tableSchemas =
+                            _kustoClient.ExecuteQueryAsync<ShowDatabaseSchemaResult>($".show database {databaseName} schema", source.Token, databaseName)
+                                .Result;
+                    },
+                    () =>
+                    {
+                        functionSchemas = _kustoClient.ExecuteQueryAsync<ShowFunctionsResult>(".show functions", source.Token, databaseName).Result;
+                    });
+            }
+
+            return AddOrUpdateDatabase(tableSchemas, functionSchemas, GlobalState.Default, databaseName,
+                clusterName);
+        }
+        
+        /// <summary>
+        /// Loads the schema for the specified database and returns a new <see cref="GlobalState"/> with the database added or updated.
+        /// </summary>
+        private GlobalState AddOrUpdateDatabase(IEnumerable<ShowDatabaseSchemaResult> tableSchemas,
+            IEnumerable<ShowFunctionsResult> functionSchemas, GlobalState globals,
+            string databaseName, string clusterName)
+        {
+            // try and show error from here.
+            DatabaseSymbol databaseSymbol = null;
+
+            if (databaseName != null)
+            {
+                databaseSymbol = LoadDatabase(tableSchemas, functionSchemas, databaseName);
+            }
+
+            if (databaseSymbol == null)
+            {
+                return globals;
+            }
+
+            var cluster = globals.GetCluster(clusterName);
+            if (cluster == null)
+            {
+                cluster = new ClusterSymbol(clusterName, new[] {databaseSymbol}, isOpen: true);
+                globals = globals.AddOrUpdateCluster(cluster);
+            }
+            else
+            {
+                cluster = cluster.AddOrUpdateDatabase(databaseSymbol);
+                globals = globals.AddOrUpdateCluster(cluster);
+            }
+
+            return globals.WithCluster(cluster).WithDatabase(databaseSymbol);
+        }
+        
+        /// <summary>
+        /// Loads the schema for the specified database into a <see cref="DatabaseSymbol"/>.
+        /// </summary>
+        private DatabaseSymbol LoadDatabase(IEnumerable<ShowDatabaseSchemaResult> tableSchemas,
+            IEnumerable<ShowFunctionsResult> functionSchemas,
+            string databaseName)
+        {
+            if (tableSchemas == null)
+            {
+                return null;
+            }
+
+            tableSchemas = tableSchemas
+                .Where(r => !string.IsNullOrEmpty(r.TableName) && !string.IsNullOrEmpty(r.ColumnName))
+                .ToArray();
+
+            var members = new List<Symbol>();
+            foreach (var table in tableSchemas.GroupBy(s => s.TableName))
+            {
+                var columns = table.Select(s => new ColumnSymbol(s.ColumnName, GetKustoType(s.ColumnType))).ToList();
+                var tableSymbol = new TableSymbol(table.Key, columns);
+                members.Add(tableSymbol);
+            }
+
+            if (functionSchemas == null)
+            {
+                return null;
+            }
+
+            foreach (var fun in functionSchemas)
+            {
+                var parameters = TranslateParameters(fun.Parameters);
+                var functionSymbol = new FunctionSymbol(fun.Name, fun.Body, parameters);
+                members.Add(functionSymbol);
+            }
+
+            return new DatabaseSymbol(databaseName, members);
+        }
+        
         /// <summary>
         /// Convert CLR type name into a Kusto scalar type.
         /// </summary>
-        private static ScalarSymbol GetKustoType(string clrTypeName)
+        private ScalarSymbol GetKustoType(string clrTypeName)
         {
             switch (clrTypeName)
             {
@@ -96,220 +205,121 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
                     throw new InvalidOperationException($"Unhandled clr type: {clrTypeName}");
             }
         }
-
+        
-        private static IReadOnlyList<Parameter> NoParameters = new Parameter[0];
-
         /// <summary>
         /// Translate Kusto parameter list declaration into into list of <see cref="Parameter"/> instances.
         /// </summary>
-        private static IReadOnlyList<Parameter> TranslateParameters(string parameters)
+        private IReadOnlyList<Parameter> TranslateParameters(string parameters)
         {
             parameters = parameters.Trim();
 
             if (string.IsNullOrEmpty(parameters) || parameters == "()")
-                return NoParameters;
+            {
+                return new Parameter[0];
+            }
 
             if (parameters[0] != '(')
+            {
                 parameters = "(" + parameters;
+            }
+
             if (parameters[parameters.Length - 1] != ')')
+            {
                 parameters = parameters + ")";
+            }
 
             var query = "let fn = " + parameters + " { };";
             var code = KustoCode.ParseAndAnalyze(query);
             var let = code.Syntax.GetFirstDescendant<LetStatement>();
             
-             FunctionSymbol function = let.Name.ReferencedSymbol is VariableSymbol variable
+            FunctionSymbol function = let.Name.ReferencedSymbol is VariableSymbol variable
                 ? variable.Type as FunctionSymbol
                 : let.Name.ReferencedSymbol as FunctionSymbol;
 
             return function.Signatures[0].Parameters;
         }
 
-        /// <summary>
+        public ScriptFileMarker[] GetSemanticMarkers(ScriptParseInfo parseInfo, ScriptFile scriptFile, string queryText)
-        /// Loads the schema for the specified databasea into a a <see cref="DatabaseSymbol"/>.
-        /// </summary>
-        private static DatabaseSymbol LoadDatabaseAsync(IEnumerable<ShowDatabaseSchemaResult> tableSchemas,
-            IEnumerable<ShowFunctionsResult> functionSchemas,
-            string databaseName)
         {
-            if (tableSchemas == null)
+            var kustoCodeService = new KustoCodeService(queryText, _schemaState);
+            var script = CodeScript.From(queryText, _schemaState);
+            var parseResult = new List<Diagnostic>();
+
+            foreach (var codeBlock in script.Blocks)
             {
-                return null;
+                parseResult.AddRange(codeBlock.Service.GetDiagnostics());
             }
 
-            tableSchemas = tableSchemas
-                .Where(r => !string.IsNullOrEmpty(r.TableName) && !string.IsNullOrEmpty(r.ColumnName))
-                .ToArray();
-
-            var members = new List<Symbol>();
-            foreach (var table in tableSchemas.GroupBy(s => s.TableName))
-            {
-                var columns = table.Select(s => new ColumnSymbol(s.ColumnName, GetKustoType(s.ColumnType))).ToList();
-                var tableSymbol = new TableSymbol(table.Key, columns);
-                members.Add(tableSymbol);
-            }
-
-            if (functionSchemas == null)
-            {
-                return null;
-            }
-
-            foreach (var fun in functionSchemas)
-            {
-                var parameters = TranslateParameters(fun.Parameters);
-                var functionSymbol = new FunctionSymbol(fun.Name, fun.Body, parameters);
-                members.Add(functionSymbol);
-            }
-
-            return new DatabaseSymbol(databaseName, members);
-        }
-
-        public static CompletionItemKind CreateCompletionItemKind(CompletionKind kustoKind)
-        {
-            CompletionItemKind kind = CompletionItemKind.Variable;
-            switch (kustoKind)
-            {
-                case CompletionKind.Syntax:
-                    kind = CompletionItemKind.Module;
-                    break;
-                case CompletionKind.Column:
-                    kind = CompletionItemKind.Field;
-                    break;
-                case CompletionKind.Variable:
-                    kind = CompletionItemKind.Variable;
-                    break;
-                case CompletionKind.Table:
-                    kind = CompletionItemKind.File;
-                    break;
-                case CompletionKind.Database:
-                    kind = CompletionItemKind.Method;
-                    break;
-                case CompletionKind.LocalFunction:
-                case CompletionKind.DatabaseFunction:
-                case CompletionKind.BuiltInFunction:
-                case CompletionKind.AggregateFunction:
-                    kind = CompletionItemKind.Function;
-                    break;
-                default:
-                    kind = CompletionItemKind.Keyword;
-                    break;
-            }
-
-            return kind;
-        }
-
-        /// <summary>
-        /// Gets default keyword when user if not connected to any Kusto cluster.
-        /// </summary>
-        public static LanguageServices.Contracts.CompletionItem[] GetDefaultKeywords(
-            ScriptDocumentInfo scriptDocumentInfo, Position textDocumentPosition)
-        {
-            var kustoCodeService = new KustoCodeService(scriptDocumentInfo.Contents, GlobalState.Default);
-            var script = CodeScript.From(scriptDocumentInfo.Contents, GlobalState.Default);
-            script.TryGetTextPosition(textDocumentPosition.Line + 1, textDocumentPosition.Character,
-                out int position); // Gets the actual offset based on line and local offset      
-            var completion = kustoCodeService.GetCompletionItems(position);
-
-            List<LanguageServices.Contracts.CompletionItem> completions =
-                new List<LanguageServices.Contracts.CompletionItem>();
-            foreach (var autoCompleteItem in completion.Items)
-            {
-                var label = autoCompleteItem.DisplayText;
-                // convert the completion item candidates into vscode format CompletionItems
-                completions.Add(AutoCompleteHelper.CreateCompletionItem(label, label + " keyword", label,
-                    CompletionItemKind.Keyword, scriptDocumentInfo.StartLine, scriptDocumentInfo.StartColumn,
-                    textDocumentPosition.Character));
-            }
-
-            return completions.ToArray();
-        }
-
-        /// <summary>
-        /// Gets default diagnostics when user if not connected to any Kusto cluster.
-        /// </summary>
-        public static ScriptFileMarker[] GetDefaultDiagnostics(ScriptParseInfo parseInfo, ScriptFile scriptFile,
-            string queryText)
-        {
-            var kustoCodeService = new KustoCodeService(queryText, GlobalState.Default);
-            var script = CodeScript.From(queryText, GlobalState.Default);
-            var parseResult = kustoCodeService.GetDiagnostics();
-
             parseInfo.ParseResult = parseResult;
-
+            
-            // build a list of Kusto script file markers from the errors.
+            if (!parseResult.Any())
-            List<ScriptFileMarker> markers = new List<ScriptFileMarker>();
-            if (parseResult != null && parseResult.Count() > 0)
             {
-                foreach (var error in parseResult)
+                return Array.Empty<ScriptFileMarker>();
-                {
+            }
-                    script.TryGetLineAndOffset(error.Start, out var startLine, out var startOffset);
+            
-                    script.TryGetLineAndOffset(error.End, out var endLine, out var endOffset);
+            // build a list of Kusto script file markers from the errors.
+            var markers = new List<ScriptFileMarker>();
+            
+            foreach (var error in parseResult)
+            {
+                script.TryGetLineAndOffset(error.Start, out var startLine, out var startOffset);
+                script.TryGetLineAndOffset(error.End, out var endLine, out var endOffset);
 
-                    // vscode specific format for error markers.
+                // vscode specific format for error markers.
-                    markers.Add(new ScriptFileMarker()
+                markers.Add(new ScriptFileMarker
+                {
+                    Message = error.Message,
+                    Level = ScriptFileMarkerLevel.Error,
+                    ScriptRegion = new ScriptRegion
                     {
-                        Message = error.Message,
+                        File = scriptFile.FilePath,
-                        Level = ScriptFileMarkerLevel.Error,
+                        StartLineNumber = startLine,
-                        ScriptRegion = new ScriptRegion()
+                        StartColumnNumber = startOffset,
-                        {
+                        StartOffset = 0,
-                            File = scriptFile.FilePath,
+                        EndLineNumber = endLine,
-                            StartLineNumber = startLine,
+                        EndColumnNumber = endOffset,
-                            StartColumnNumber = startOffset,
+                        EndOffset = 0
-                            StartOffset = 0,
+                    }
-                            EndLineNumber = endLine,
+                });
-                            EndColumnNumber = endOffset,
-                            EndOffset = 0
-                        }
-                    });
-                }
             }
 
             return markers.ToArray();
         }
-
+        
-        /// <summary>
+        public DefinitionResult GetDefinition(string queryText, int index, int startLine, int startColumn, bool throwOnError = false)
-        /// Loads the schema for the specified database and returns a new <see cref="GlobalState"/> with the database added or updated.
-        /// </summary>
-        public static GlobalState AddOrUpdateDatabase(IEnumerable<ShowDatabaseSchemaResult> tableSchemas,
-            IEnumerable<ShowFunctionsResult> functionSchemas, GlobalState globals,
-            string databaseName, string clusterName)
         {
-            // try and show error from here.
+            //TODOKusto: API wasnt working properly, need to check that part.
-            DatabaseSymbol databaseSymbol = null;
+            var abc = KustoCode.ParseAndAnalyze(queryText, _schemaState); 
+            var kustoCodeService = new KustoCodeService(abc);
+            //var kustoCodeService = new KustoCodeService(queryText, globals);
+            var relatedInfo = kustoCodeService.GetRelatedElements(index);
 
-            if (databaseName != null)
+            if (relatedInfo != null && relatedInfo.Elements.Count > 1)
             {
-                databaseSymbol = LoadDatabaseAsync(tableSchemas, functionSchemas, databaseName);
             }
 
-            if (databaseSymbol == null)
+            return null;
-            {
-                return globals;
-            }
-
-            var cluster = globals.GetCluster(clusterName);
-            if (cluster == null)
-            {
-                cluster = new ClusterSymbol(clusterName, new[] {databaseSymbol}, isOpen: true);
-                globals = globals.AddOrUpdateCluster(cluster);
-            }
-            else
-            {
-                cluster = cluster.AddOrUpdateDatabase(databaseSymbol);
-                globals = globals.AddOrUpdateCluster(cluster);
-            }
-
-            globals = globals.WithCluster(cluster).WithDatabase(databaseSymbol);
-
-            return globals;
         }
-
+        
-        /// <inheritdoc/>
+        public Hover GetHoverHelp(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, bool throwOnError = false)
-        public static LanguageServices.Contracts.CompletionItem[] GetAutoCompleteSuggestions(
-            ScriptDocumentInfo scriptDocumentInfo, Position textPosition, GlobalState schemaState,
-            bool throwOnError = false)
         {
-            var script = CodeScript.From(scriptDocumentInfo.Contents, schemaState);
+            var script = CodeScript.From(scriptDocumentInfo.Contents, _schemaState);
+            script.TryGetTextPosition(textPosition.Line + 1, textPosition.Character + 1, out int position);
+
+            var codeBlock = script.GetBlockAtPosition(position);
+            var quickInfo = codeBlock.Service.GetQuickInfo(position);
+
+            return AutoCompleteHelper.ConvertQuickInfoToHover(
+                quickInfo.Text,
+                "kusto",
+                scriptDocumentInfo.StartLine,
+                scriptDocumentInfo.StartColumn,
+                textPosition.Character);
+        }
+        
+        public LanguageServices.Contracts.CompletionItem[] GetAutoCompleteSuggestions(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, bool throwOnError = false)
+        {
+            var script = CodeScript.From(scriptDocumentInfo.Contents, _schemaState);
             script.TryGetTextPosition(textPosition.Line + 1, textPosition.Character + 1,
                 out int position); // Gets the actual offset based on line and local offset
 
@@ -334,86 +344,29 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
 
             return completions.ToArray();
         }
-
+        
-        /// <inheritdoc/>
+        private CompletionItemKind CreateCompletionItemKind(CompletionKind kustoKind)
-        public static Hover GetHoverHelp(ScriptDocumentInfo scriptDocumentInfo, Position textPosition,
-            GlobalState schemaState, bool throwOnError = false)
         {
-            var script = CodeScript.From(scriptDocumentInfo.Contents, schemaState);
+            switch (kustoKind)
-            script.TryGetTextPosition(textPosition.Line + 1, textPosition.Character + 1, out int position);
-
-            var codeBlock = script.GetBlockAtPosition(position);
-            var quickInfo = codeBlock.Service.GetQuickInfo(position);
-
-            return AutoCompleteHelper.ConvertQuickInfoToHover(
-                quickInfo.Text,
-                "kusto",
-                scriptDocumentInfo.StartLine,
-                scriptDocumentInfo.StartColumn,
-                textPosition.Character);
-        }
-
-        /// <inheritdoc/>
-        public static DefinitionResult GetDefinition(string queryText, int index, int startLine, int startColumn,
-            GlobalState schemaState, bool throwOnError = false)
-        {
-            var abc = KustoCode.ParseAndAnalyze(queryText,
-                schemaState); //TODOKusto: API wasnt working properly, need to check that part.
-            var kustoCodeService = new KustoCodeService(abc);
-            //var kustoCodeService = new KustoCodeService(queryText, globals);
-            var relatedInfo = kustoCodeService.GetRelatedElements(index);
-
-            if (relatedInfo != null && relatedInfo.Elements.Count > 1)
             {
+                case CompletionKind.Syntax:
+                    return CompletionItemKind.Module;
+                case CompletionKind.Column:
+                    return CompletionItemKind.Field;
+                case CompletionKind.Variable:
+                    return CompletionItemKind.Variable;
+                case CompletionKind.Table:
+                    return CompletionItemKind.File;
+                case CompletionKind.Database:
+                    return CompletionItemKind.Method;
+                case CompletionKind.LocalFunction:
+                case CompletionKind.DatabaseFunction:
+                case CompletionKind.BuiltInFunction:
+                case CompletionKind.AggregateFunction:
+                    return CompletionItemKind.Function;
+                default:
+                    return CompletionItemKind.Keyword;
             }
-
-            return null;
-        }
-
-        /// <inheritdoc/>
-        public static ScriptFileMarker[] GetSemanticMarkers(ScriptParseInfo parseInfo, ScriptFile scriptFile,
-            string queryText, GlobalState schemaState)
-        {
-            var kustoCodeService = new KustoCodeService(queryText, schemaState);
-            var script = CodeScript.From(queryText, schemaState);
-            var parseResult = new List<KustoDiagnostic>();
-
-            foreach (var codeBlock in script.Blocks)
-            {
-                parseResult.AddRange(codeBlock.Service.GetDiagnostics());
-            }
-
-            parseInfo.ParseResult = parseResult;
-
-            // build a list of Kusto script file markers from the errors.
-            List<ScriptFileMarker> markers = new List<ScriptFileMarker>();
-            if (parseResult != null && parseResult.Any())
-            {
-                foreach (var error in parseResult)
-                {
-                    script.TryGetLineAndOffset(error.Start, out var startLine, out var startOffset);
-                    script.TryGetLineAndOffset(error.End, out var endLine, out var endOffset);
-
-                    // vscode specific format for error markers.
-                    markers.Add(new ScriptFileMarker()
-                    {
-                        Message = error.Message,
-                        Level = ScriptFileMarkerLevel.Error,
-                        ScriptRegion = new ScriptRegion()
-                        {
-                            File = scriptFile.FilePath,
-                            StartLineNumber = startLine,
-                            StartColumnNumber = startOffset,
-                            StartOffset = 0,
-                            EndLineNumber = endLine,
-                            EndColumnNumber = endOffset,
-                            EndOffset = 0
-                        }
-                    });
-                }
-            }
-
-            return markers.ToArray();
         }
     }
-}
+}

src/Microsoft.Kusto.ServiceLayer/DataSource/Intellisense/KustoIntellisenseHelper.cs

@@ -0,0 +1,90 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System.Collections.Generic;
+using System.Linq;
+using Kusto.Language;
+using Kusto.Language.Editor;
+using Microsoft.Kusto.ServiceLayer.LanguageServices;
+using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
+using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
+
+namespace Microsoft.Kusto.ServiceLayer.DataSource.Intellisense
+{
+    /// <summary>
+    /// Kusto specific class for intellisense helper functions.
+    /// </summary>
+    public class KustoIntellisenseHelper
+    {
+        /// <summary>
+        /// Gets default keyword when user if not connected to any Kusto cluster.
+        /// </summary>
+        public static LanguageServices.Contracts.CompletionItem[] GetDefaultKeywords(
+            ScriptDocumentInfo scriptDocumentInfo, Position textDocumentPosition)
+        {
+            var kustoCodeService = new KustoCodeService(scriptDocumentInfo.Contents, GlobalState.Default);
+            var script = CodeScript.From(scriptDocumentInfo.Contents, GlobalState.Default);
+            script.TryGetTextPosition(textDocumentPosition.Line + 1, textDocumentPosition.Character,
+                out int position); // Gets the actual offset based on line and local offset      
+            var completion = kustoCodeService.GetCompletionItems(position);
+
+            List<LanguageServices.Contracts.CompletionItem> completions =
+                new List<LanguageServices.Contracts.CompletionItem>();
+            foreach (var autoCompleteItem in completion.Items)
+            {
+                var label = autoCompleteItem.DisplayText;
+                // convert the completion item candidates into vscode format CompletionItems
+                completions.Add(AutoCompleteHelper.CreateCompletionItem(label, label + " keyword", label,
+                    CompletionItemKind.Keyword, scriptDocumentInfo.StartLine, scriptDocumentInfo.StartColumn,
+                    textDocumentPosition.Character));
+            }
+
+            return completions.ToArray();
+        }
+
+        /// <summary>
+        /// Gets default diagnostics when user if not connected to any Kusto cluster.
+        /// </summary>
+        public static ScriptFileMarker[] GetDefaultDiagnostics(ScriptParseInfo parseInfo, ScriptFile scriptFile,
+            string queryText)
+        {
+            var kustoCodeService = new KustoCodeService(queryText, GlobalState.Default);
+            var script = CodeScript.From(queryText, GlobalState.Default);
+            var parseResult = kustoCodeService.GetDiagnostics();
+
+            parseInfo.ParseResult = parseResult;
+
+            // build a list of Kusto script file markers from the errors.
+            List<ScriptFileMarker> markers = new List<ScriptFileMarker>();
+            if (parseResult != null && parseResult.Count() > 0)
+            {
+                foreach (var error in parseResult)
+                {
+                    script.TryGetLineAndOffset(error.Start, out var startLine, out var startOffset);
+                    script.TryGetLineAndOffset(error.End, out var endLine, out var endOffset);
+
+                    // vscode specific format for error markers.
+                    markers.Add(new ScriptFileMarker()
+                    {
+                        Message = error.Message,
+                        Level = ScriptFileMarkerLevel.Error,
+                        ScriptRegion = new ScriptRegion
+                        {
+                            File = scriptFile.FilePath,
+                            StartLineNumber = startLine,
+                            StartColumnNumber = startOffset,
+                            StartOffset = 0,
+                            EndLineNumber = endLine,
+                            EndColumnNumber = endOffset,
+                            EndOffset = 0
+                        }
+                    });
+                }
+            }
+
+            return markers.ToArray();
+        }
+    }
+}

src/Microsoft.Kusto.ServiceLayer/DataSource/Intellisense/ScriptParseInfo.cs

@@ -7,7 +7,7 @@ using System.Collections.Generic;
 using Kusto.Language;
 using Kusto.Language.Editor;
 
-namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
+namespace Microsoft.Kusto.ServiceLayer.DataSource.Intellisense
 {
     /// <summary>
     /// Data Source specific class for storing cached metadata regarding a parsed KQL file.

src/Microsoft.Kusto.ServiceLayer/DataSource/Intellisense/ShowDatabaseSchemaResult.cs

@@ -1,4 +1,4 @@
-namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
+namespace Microsoft.Kusto.ServiceLayer.DataSource.Intellisense
 {
     public class ShowDatabaseSchemaResult
     {

src/Microsoft.Kusto.ServiceLayer/DataSource/Intellisense/ShowDatabasesResult.cs

@@ -1,4 +1,4 @@
-namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
+namespace Microsoft.Kusto.ServiceLayer.DataSource.Intellisense
 {
     public class ShowDatabasesResult
     {

src/Microsoft.Kusto.ServiceLayer/DataSource/Intellisense/ShowFunctionsResult.cs

@@ -1,4 +1,4 @@
-namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
+namespace Microsoft.Kusto.ServiceLayer.DataSource.Intellisense
 {
     public class ShowFunctionsResult
     {

src/Microsoft.Kusto.ServiceLayer/DataSource/KustoClient.cs

@@ -16,7 +16,6 @@ using Kusto.Language;
 using Kusto.Language.Editor;
 using Microsoft.Kusto.ServiceLayer.Connection;
 using Microsoft.Kusto.ServiceLayer.DataSource.Contracts;
-using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
 using Microsoft.Kusto.ServiceLayer.Utility;
 
 namespace Microsoft.Kusto.ServiceLayer.DataSource
@@ -30,12 +29,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
 
         [DebuggerBrowsable(DebuggerBrowsableState.Never)]
         private ICslQueryProvider _kustoQueryProvider;
-
+        
-        /// <summary>
-        /// SchemaState used for getting intellisense info.
-        /// </summary>
-        public GlobalState SchemaState { get; private set; }
-
         public string ClusterName { get; private set; }
         public string DatabaseName { get; private set; }
 
@@ -43,7 +37,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         {
             _ownerUri = ownerUri;
             Initialize(connectionDetails);
-            SchemaState = LoadSchemaState();
         }
 
         private string ParseDatabaseName(string databaseName)
@@ -55,31 +48,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
                 : databaseName;
         }
 
-        private GlobalState LoadSchemaState()
-        {
-            IEnumerable<ShowDatabaseSchemaResult> tableSchemas = Enumerable.Empty<ShowDatabaseSchemaResult>();
-            IEnumerable<ShowFunctionsResult> functionSchemas = Enumerable.Empty<ShowFunctionsResult>();
-
-            if (!string.IsNullOrWhiteSpace(DatabaseName))
-            {
-                var source = new CancellationTokenSource();
-                Parallel.Invoke(() =>
-                    {
-                        tableSchemas =
-                            ExecuteQueryAsync<ShowDatabaseSchemaResult>($".show database {DatabaseName} schema", source.Token, DatabaseName)
-                                .Result;
-                    },
-                    () =>
-                    {
-                        functionSchemas = ExecuteQueryAsync<ShowFunctionsResult>(".show functions", source.Token, DatabaseName).Result;
-                    });
-            }
-
-            return KustoIntellisenseHelper.AddOrUpdateDatabase(tableSchemas, functionSchemas,
-                GlobalState.Default,
-                DatabaseName, ClusterName);
-        }
-
         private void Initialize(DataSourceConnectionDetails connectionDetails)
         {
             var stringBuilder = GetKustoConnectionStringBuilder(connectionDetails);
@@ -275,7 +243,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         public void UpdateDatabase(string databaseName)
         {
             DatabaseName = ParseDatabaseName(databaseName);
-            SchemaState = LoadSchemaState();
         }
 
         public void Dispose()

src/Microsoft.Kusto.ServiceLayer/DataSource/KustoDataSource.cs

@@ -14,10 +14,13 @@ using System.Threading.Tasks;
 using Kusto.Cloud.Platform.Data;
 using Kusto.Data;
 using Kusto.Data.Data;
-using Kusto.Language;
+using Microsoft.Kusto.ServiceLayer.DataSource.Intellisense;
 using Microsoft.Kusto.ServiceLayer.DataSource.Metadata;
 using Microsoft.Kusto.ServiceLayer.DataSource.Models;
+using Microsoft.Kusto.ServiceLayer.LanguageServices;
+using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
 using Microsoft.Kusto.ServiceLayer.Utility;
+using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
 
 namespace Microsoft.Kusto.ServiceLayer.DataSource
 {
@@ -26,7 +29,8 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
     /// </summary>
     public class KustoDataSource : DataSourceBase
     {
-        private IKustoClient _kustoClient;
+        private readonly IKustoClient _kustoClient;
+        private readonly IIntellisenseClient _intellisenseClient;
 
         /// <summary>
         /// List of databases.
@@ -56,8 +60,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         public override string DatabaseName => _kustoClient.DatabaseName;
 
         public override string ClusterName => _kustoClient.ClusterName;
-        
-        public override GlobalState SchemaState => _kustoClient.SchemaState;
 
         // Some clusters have this signature. Queries might slightly differ for Aria
         private const string AriaProxyURL = "kusto.aria.microsoft.com"; 
@@ -77,9 +79,10 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// <summary>
         /// Prevents a default instance of the <see cref="IDataSource"/> class from being created.
         /// </summary>
-        public KustoDataSource(IKustoClient kustoClient)
+        public KustoDataSource(IKustoClient kustoClient, IIntellisenseClient intellisenseClient)
         {
             _kustoClient = kustoClient;
+            _intellisenseClient = intellisenseClient;
             // Check if a connection can be made
             ValidationUtils.IsTrue<ArgumentException>(Exists().Result,
                 $"Unable to connect. ClusterName = {ClusterName}, DatabaseName = {DatabaseName}");
@@ -246,6 +249,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         public override void UpdateDatabase(string databaseName)
         {
             _kustoClient.UpdateDatabase(databaseName);
+            _intellisenseClient.UpdateDatabase(databaseName);
         }
         
         /// <summary>
@@ -803,12 +807,32 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
                 ? string.Empty 
                 : $"{functionInfo.Name}{functionInfo.Parameters}";
         }
-
+        
         private string GenerateMetadataKey(string databaseName, string objectName)
         {
             return string.IsNullOrWhiteSpace(objectName) ? databaseName : $"{databaseName}.{objectName}";
         }
 
+        public override ScriptFileMarker[] GetSemanticMarkers(ScriptParseInfo parseInfo, ScriptFile scriptFile, string queryText)
+        {
+            return _intellisenseClient.GetSemanticMarkers(parseInfo, scriptFile, queryText);
+        }
+
+        public override DefinitionResult GetDefinition(string queryText, int index, int startLine, int startColumn, bool throwOnError = false)
+        {
+            return _intellisenseClient.GetDefinition(queryText, index, startLine, startColumn, throwOnError);
+        }
+
+        public override Hover GetHoverHelp(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, bool throwOnError = false)
+        {
+            return _intellisenseClient.GetHoverHelp(scriptDocumentInfo, textPosition, throwOnError);
+        }
+
+        public override CompletionItem[] GetAutoCompleteSuggestions(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, bool throwOnError = false)
+        {
+            return _intellisenseClient.GetAutoCompleteSuggestions(scriptDocumentInfo, textPosition, throwOnError);
+        }
+
         #endregion
     }
 }

src/Microsoft.Kusto.ServiceLayer/HostLoader.cs

@@ -78,7 +78,7 @@ namespace Microsoft.Kusto.ServiceLayer
             WorkspaceService<SqlToolsSettings>.Instance.InitializeService(serviceHost);
             serviceProvider.RegisterSingleService(WorkspaceService<SqlToolsSettings>.Instance);
 
-            LanguageService.Instance.InitializeService(serviceHost, connectedBindingQueue, dataSourceFactory);
+            LanguageService.Instance.InitializeService(serviceHost, connectedBindingQueue);
             serviceProvider.RegisterSingleService(LanguageService.Instance);
 
             ConnectionService.Instance.InitializeService(serviceHost, dataSourceConnectionFactory, connectedBindingQueue, dataSourceFactory);

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

@@ -15,10 +15,9 @@ using Microsoft.SqlServer.Management.SqlParser.Parser;
 using Microsoft.SqlServer.Management.SqlParser.SqlCodeDom;
 using Microsoft.SqlTools.Hosting.Protocol;
 using Microsoft.Kusto.ServiceLayer.DataSource;
-using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
 using Microsoft.Kusto.ServiceLayer.Connection;
 using Microsoft.Kusto.ServiceLayer.Connection.Contracts;
-using Microsoft.Kusto.ServiceLayer.LanguageServices.Completion;
+using Microsoft.Kusto.ServiceLayer.DataSource.Intellisense;
 using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
 using Microsoft.Kusto.ServiceLayer.Utility;
 using Microsoft.Kusto.ServiceLayer.Workspace;
@@ -67,9 +66,6 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
 
         internal const int PeekDefinitionTimeout = 10 * OneSecond;
 
-        // For testability only
-        internal Task DelayedDiagnosticsTask = null;
-
         private ConnectionService connectionService = null;
 
         private WorkspaceService<SqlToolsSettings> workspaceServiceInstance;
@@ -209,7 +205,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
         /// <param name="context"></param>
         /// <param name="dataSourceFactory"></param>
         /// <param name="connectedBindingQueue"></param>
-        public void InitializeService(ServiceHost serviceHost, IConnectedBindingQueue connectedBindingQueue, IDataSourceFactory dataSourceFactory)
+        public void InitializeService(ServiceHost serviceHost, IConnectedBindingQueue connectedBindingQueue)
         {
             _bindingQueue = connectedBindingQueue;
             // Register the requests that this service will handle
@@ -766,18 +762,16 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
                 connInfo.TryGetConnection("Default", out connection);
                 IDataSource dataSource = connection.GetUnderlyingConnection();
                 
-                return KustoIntellisenseHelper.GetDefinition(scriptFile.Contents, textDocumentPosition.Position.Character, 1, 1, dataSource.SchemaState);
+                return dataSource.GetDefinition(scriptFile.Contents, textDocumentPosition.Position.Character, 1, 1);
             }
-            else
+
+            // User is not connected.
+            return new DefinitionResult
             {
-                // User is not connected.
+                IsErrorResult = true,
-                return new DefinitionResult
+                Message = SR.PeekDefinitionNotConnectedError,
-                {
+                Locations = null
-                    IsErrorResult = true,
+            };
-                    Message = SR.PeekDefinitionNotConnectedError,
-                    Locations = null
-                };
-            }
         }
 
         /// <summary>
@@ -809,9 +803,8 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
                                 
                                 ReliableDataSourceConnection connection;
                                 connInfo.TryGetConnection("Default", out connection);
-                                IDataSource dataSource = connection.GetUnderlyingConnection();               
+                                IDataSource dataSource = connection.GetUnderlyingConnection();
-                                
+                                return dataSource.GetHoverHelp(scriptDocumentInfo, textDocumentPosition.Position);
-                                return KustoIntellisenseHelper.GetHoverHelp(scriptDocumentInfo, textDocumentPosition.Position, dataSource.SchemaState);
                             });
 
                         queueItem.ItemProcessed.WaitOne();
@@ -855,14 +848,17 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
 
             ScriptDocumentInfo scriptDocumentInfo = new ScriptDocumentInfo(textDocumentPosition, scriptFile, scriptParseInfo);
 
-            if(connInfo != null){
+            if (connInfo != null)
+            {
                 ReliableDataSourceConnection connection;
                 connInfo.TryGetConnection("Default", out connection);
                 IDataSource dataSource = connection.GetUnderlyingConnection();
-			    
+
-                resultCompletionItems = KustoIntellisenseHelper.GetAutoCompleteSuggestions(scriptDocumentInfo, textDocumentPosition.Position, dataSource.SchemaState);
+                resultCompletionItems =
+                    dataSource.GetAutoCompleteSuggestions(scriptDocumentInfo, textDocumentPosition.Position);
             }
-            else{
+            else
+            {
                 resultCompletionItems = DataSourceFactory.GetDefaultAutoComplete(DataSourceType.Kusto, scriptDocumentInfo, textDocumentPosition.Position);
             }
 
@@ -932,7 +928,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
             existingRequestCancellation = new CancellationTokenSource();
             Task.Factory.StartNew(
                 () =>
-                    this.DelayedDiagnosticsTask = DelayThenInvokeDiagnostics(
+                    DelayThenInvokeDiagnostics(
                         LanguageService.DiagnosticParseDelay,
                         filesToAnalyze,
                         eventContext,
@@ -1000,17 +996,19 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
                 ConnectionServiceInstance.TryFindConnection(
                     scriptFile.ClientUri,
                     out connInfo);
-                
+
-                if(connInfo != null){
+                if (connInfo != null)
+                {
                     connInfo.TryGetConnection("Default", out var connection);
                     IDataSource dataSource = connection.GetUnderlyingConnection();
-                    
+
-                    semanticMarkers = KustoIntellisenseHelper.GetSemanticMarkers(parseInfo, scriptFile, scriptFile.Contents, dataSource.SchemaState);
+                    semanticMarkers = dataSource.GetSemanticMarkers(parseInfo, scriptFile, scriptFile.Contents);
-			    }
+                }
-                else{
+                else
+                {
                     semanticMarkers = DataSourceFactory.GetDefaultSemanticMarkers(DataSourceType.Kusto, parseInfo, scriptFile, scriptFile.Contents);
                 }
-                
+
                 Logger.Write(TraceEventType.Verbose, "Analysis complete.");
 
                 await DiagnosticsHelper.PublishScriptDiagnostics(scriptFile, semanticMarkers, eventContext);

src/Microsoft.Kusto.ServiceLayer/LanguageServices/ScriptDocumentInfo.cs

@@ -3,12 +3,12 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using Microsoft.Kusto.ServiceLayer.DataSource.Intellisense;
+using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
 using Microsoft.SqlServer.Management.SqlParser.Parser;
 using Microsoft.SqlTools.Utility;
-using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
-using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
 
-namespace Microsoft.Kusto.ServiceLayer.LanguageServices.Completion
+namespace Microsoft.Kusto.ServiceLayer.LanguageServices
 {
     /// <summary>
     /// A class to calculate the numbers used by SQL parser using the text positions and content

test/Microsoft.Kusto.ServiceLayer.UnitTests/DataSource/DataSourceFactoryTests.cs

@@ -2,8 +2,8 @@ using System;
 using System.Collections.Generic;
 using Microsoft.Kusto.ServiceLayer.Connection.Contracts;
 using Microsoft.Kusto.ServiceLayer.DataSource;
-using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
+using Microsoft.Kusto.ServiceLayer.DataSource.Intellisense;
-using Microsoft.Kusto.ServiceLayer.LanguageServices.Completion;
+using Microsoft.Kusto.ServiceLayer.LanguageServices;
 using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
 using NUnit.Framework;
 

test/Microsoft.Kusto.ServiceLayer.UnitTests/DataSource/DataSourceIntellisense/KustoIntellisenseClientTests.cs

@@ -0,0 +1,174 @@
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Kusto.ServiceLayer.DataSource;
+using Microsoft.Kusto.ServiceLayer.DataSource.Intellisense;
+using Microsoft.Kusto.ServiceLayer.LanguageServices;
+using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
+using Moq;
+using NUnit.Framework;
+
+namespace Microsoft.Kusto.ServiceLayer.UnitTests.DataSource.DataSourceIntellisense
+{
+    public class KustoIntellisenseClientTests
+    {
+        private Mock<IKustoClient> GetMockKustoClient()
+        {
+            var kustoClientMock = new Mock<IKustoClient>();
+            kustoClientMock.Setup(x => x.ClusterName).Returns("https://fake.url.com");
+            kustoClientMock.Setup(x => x.DatabaseName).Returns("FakeDatabaseName");
+
+            var databaseSchema = new ShowDatabaseSchemaResult
+            {
+                DatabaseName = "FakeDatabaseName",
+                TableName = "FakeTableName",
+                ColumnName = "FakeColumnName",
+                ColumnType = "bool",
+                IsDefaultTable = false,
+                IsDefaultColumn = false,
+                PrettyName = "Fake Table Name",
+                Version = "",
+                Folder = "FakeTableFolder",
+                DocName = ""
+            };
+
+            var databaseSchemaResults = new List<ShowDatabaseSchemaResult> {databaseSchema} as IEnumerable<ShowDatabaseSchemaResult>;
+            kustoClientMock.Setup(x => x.ExecuteQueryAsync<ShowDatabaseSchemaResult>(It.IsAny<string>(), It.IsAny<CancellationToken>(), "FakeDatabaseName"))
+                .Returns(Task.FromResult(databaseSchemaResults));
+
+            var functionSchema = new ShowFunctionsResult
+            {
+                Name = "FakeFunctionName",
+                Parameters = "a:real, b:real",
+                Body = "a+b",
+                Folder = "FakeFunctionFolder",
+                DocString = ""
+            };
+
+            var functionSchemaResults = new List<ShowFunctionsResult> {functionSchema} as IEnumerable<ShowFunctionsResult>;
+            kustoClientMock.Setup(x => x.ExecuteQueryAsync<ShowFunctionsResult>(It.IsAny<string>(), It.IsAny<CancellationToken>(), "FakeDatabaseName"))
+                .Returns(Task.FromResult(functionSchemaResults));
+
+            return kustoClientMock;
+        }
+        
+        [Test]
+        public void GetSemanticMarkers_Returns_Error_For_InvalidText()
+        {
+            var kustoClientMock = GetMockKustoClient();
+            
+            var client = new KustoIntellisenseClient(kustoClientMock.Object);
+            var semanticMarkers = client.GetSemanticMarkers(new ScriptParseInfo(), new ScriptFile("", "", ""), "InvalidText");
+            
+            var semanticMarker = semanticMarkers.Single();
+            Assert.AreEqual(ScriptFileMarkerLevel.Error, semanticMarker.Level);
+            Assert.AreEqual("The name 'InvalidText' does not refer to any known column, table, variable or function.", semanticMarker.Message);
+            Assert.IsNotNull(semanticMarker.ScriptRegion);
+        }
+
+        [Test]
+        public void GetSemanticMarkers_Returns_Zero_SemanticMarkers_For_ValidQueryText()
+        {
+            var kustoClientMock = GetMockKustoClient();
+            
+            var client = new KustoIntellisenseClient(kustoClientMock.Object);
+            var queryText = @".show commands";
+            var semanticMarkers = client.GetSemanticMarkers(new ScriptParseInfo(), new ScriptFile("", "", ""), queryText);
+            
+            Assert.AreEqual(0, semanticMarkers.Length);
+        }
+
+        [Test]
+        public void GetDefinition_Returns_Null()
+        {
+            var kustoClientMock = GetMockKustoClient();
+            var client = new KustoIntellisenseClient(kustoClientMock.Object);
+            var definition = client.GetDefinition("queryText", 0, 1, 1);
+            
+            // finish these assertions once the function is implemented
+            Assert.IsNull(definition);
+        }
+
+        [Test]
+        public void GetHoverHelp_Returns_Hover()
+        {
+            var kustoClientMock = GetMockKustoClient();
+            var client = new KustoIntellisenseClient(kustoClientMock.Object);
+            var textDocumentPosition = new TextDocumentPosition
+            {
+                Position = new Position()
+            };
+            var scriptFile = new ScriptFile("", "", "");
+            var scriptParseInfo = new ScriptParseInfo();
+            var documentInfo = new ScriptDocumentInfo(textDocumentPosition, scriptFile, scriptParseInfo);
+            
+            var hover = client.GetHoverHelp(documentInfo, new Position());
+
+            Assert.IsNotNull(hover);
+        }
+
+        [Test]
+        public void GetAutoCompleteSuggestions_Returns_CompletionItems()
+        {
+            var kustoClientMock = GetMockKustoClient();
+            var client = new KustoIntellisenseClient(kustoClientMock.Object);
+
+            var position = new Position();
+            var textDocumentPosition = new TextDocumentPosition
+            {
+                Position = position
+            };
+            var scriptFile = new ScriptFile("", "", "");
+            var scriptParseInfo = new ScriptParseInfo();
+            var documentInfo = new ScriptDocumentInfo(textDocumentPosition, scriptFile, scriptParseInfo);
+            var items = client.GetAutoCompleteSuggestions(documentInfo, position);
+            
+            Assert.AreEqual(20, items.Length);
+        }
+        [Test]
+        public void UpdateDatabase_Updates_SchemaState()
+        {
+            var kustoClientMock = GetMockKustoClient();
+            
+            var databaseSchema = new ShowDatabaseSchemaResult
+            {
+                DatabaseName = "NewDatabaseName",
+                TableName = "NewTableName",
+                ColumnName = "NewColumnName",
+                ColumnType = "bool",
+                IsDefaultTable = false,
+                IsDefaultColumn = false,
+                PrettyName = "New Table Name",
+                Version = "",
+                Folder = "NewTableFolder",
+                DocName = ""
+            };
+
+            var databaseSchemaResults = new List<ShowDatabaseSchemaResult> {databaseSchema} as IEnumerable<ShowDatabaseSchemaResult>;
+            kustoClientMock.Setup(x => x.ExecuteQueryAsync<ShowDatabaseSchemaResult>(It.IsAny<string>(), It.IsAny<CancellationToken>(), "NewDatabaseName"))
+                .Returns(Task.FromResult(databaseSchemaResults));
+            
+            var client = new KustoIntellisenseClient(kustoClientMock.Object);
+            client.UpdateDatabase("NewDatabaseName");
+
+            var position = new Position();
+            var textDocumentPosition = new TextDocumentPosition
+            {
+                Position = position
+            };
+            var scriptFile = new ScriptFile("", "", "");
+            var scriptParseInfo = new ScriptParseInfo();
+            var documentInfo = new ScriptDocumentInfo(textDocumentPosition, scriptFile, scriptParseInfo);
+            var items = client.GetAutoCompleteSuggestions(documentInfo, position);
+            
+            Assert.AreEqual(19, items.Length);
+            var tableItem = items.FirstOrDefault(x => x.Detail == "Table");
+
+            // assert new table is being returned to show database has changed
+            Assert.IsNotNull(tableItem);
+            Assert.AreEqual("NewTableName", tableItem.InsertText);
+            Assert.AreEqual("NewTableName", tableItem.Label);
+        }
+    }
+}

test/Microsoft.Kusto.ServiceLayer.UnitTests/DataSource/DataSourceIntellisense/KustoIntellisenseHelperTests.cs

@@ -1,8 +1,5 @@
-using System.Linq;
+using Microsoft.Kusto.ServiceLayer.DataSource.Intellisense;
-using Kusto.Language.Editor;
+using Microsoft.Kusto.ServiceLayer.LanguageServices;
-using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
-using Microsoft.Kusto.ServiceLayer.LanguageServices.Completion;
-using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
 using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
 using NUnit.Framework;
 
@@ -10,37 +7,6 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.DataSource.DataSourceIntellisen
 {
     public class KustoIntellisenseHelperTests
     {
-        [TestCase(CompletionKind.Syntax, CompletionItemKind.Module)]
-        [TestCase(CompletionKind.Column, CompletionItemKind.Field)]
-        [TestCase(CompletionKind.Variable, CompletionItemKind.Variable)]
-        [TestCase(CompletionKind.Table, CompletionItemKind.File)]
-        [TestCase(CompletionKind.Database, CompletionItemKind.Method)]
-        [TestCase(CompletionKind.LocalFunction, CompletionItemKind.Function)]
-        [TestCase(CompletionKind.DatabaseFunction, CompletionItemKind.Function)]
-        [TestCase(CompletionKind.BuiltInFunction, CompletionItemKind.Function)]
-        [TestCase(CompletionKind.AggregateFunction, CompletionItemKind.Function)]
-        [TestCase(CompletionKind.Unknown, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.Keyword, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.Punctuation, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.Identifier, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.Example, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.ScalarPrefix, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.TabularPrefix, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.TabularSuffix, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.QueryPrefix, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.CommandPrefix, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.ScalarInfix, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.RenderChart, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.Parameter, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.Cluster, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.MaterialiedView, CompletionItemKind.Keyword)]
-        [TestCase(CompletionKind.ScalarType, CompletionItemKind.Keyword)]
-        public void CreateCompletionItemKind_Returns_Kind(CompletionKind completionKind, CompletionItemKind expected)
-        {
-            var result = KustoIntellisenseHelper.CreateCompletionItemKind(completionKind);
-            Assert.AreEqual(expected, result);
-        }
-
         [Test]
         public void GetDefaultKeywords_Returns_Keywords()
         {

test/Microsoft.Kusto.ServiceLayer.UnitTests/DataSource/KustoClientTests.cs

@@ -0,0 +1,58 @@
+using System;
+using Microsoft.Kusto.ServiceLayer.DataSource;
+using Microsoft.Kusto.ServiceLayer.DataSource.Contracts;
+using NUnit.Framework;
+
+namespace Microsoft.Kusto.ServiceLayer.UnitTests.DataSource
+{
+    public class KustoClientTests
+    {
+        [Test]
+        public void Constructor_Throws_ArgumentException_For_MissingToken()
+        {
+            var connectionDetails = new DataSourceConnectionDetails
+            {
+                UserToken = ""
+            };
+            
+            Assert.Throws<ArgumentException>(() => new KustoClient(connectionDetails, "ownerUri"));
+        }
+
+        [Test]
+        public void Constructor_Sets_ClusterName_With_DefaultDatabaseName()
+        {
+            string clusterName = "https://fake.url.com";
+            var connectionDetails = new DataSourceConnectionDetails
+            {
+                UserToken = "UserToken",
+                ServerName = clusterName,
+                DatabaseName = "",
+                AuthenticationType = "AzureMFA"
+            };
+
+            var client = new KustoClient(connectionDetails, "ownerUri");
+
+            Assert.AreEqual(clusterName, client.ClusterName);
+            Assert.AreEqual("NetDefaultDB", client.DatabaseName);
+        }
+
+        [TestCase("dstsAuth")]
+        [TestCase("AzureMFA")]
+        public void Constructor_Creates_Client_With_Valid_AuthenticationType(string authenticationType)
+        {
+            string clusterName = "https://fake.url.com";
+            var connectionDetails = new DataSourceConnectionDetails
+            {
+                UserToken = "UserToken",
+                ServerName = clusterName,
+                DatabaseName = "FakeDatabaseName",
+                AuthenticationType = authenticationType
+            };
+
+            var client = new KustoClient(connectionDetails, "ownerUri");
+
+            Assert.AreEqual(clusterName, client.ClusterName);
+            Assert.AreEqual("FakeDatabaseName", client.DatabaseName);
+        }
+    }
+}