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>
-    /// Kusto specific class for intellisense helper functions.
-    /// </summary>
-    public class KustoIntellisenseHelper
+    public class KustoIntellisenseClient : IIntellisenseClient
     {
+        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>
-        /// 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)
+        public ScriptFileMarker[] GetSemanticMarkers(ScriptParseInfo parseInfo, ScriptFile scriptFile, string queryText)
         {
-            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.
-            List<ScriptFileMarker> markers = new List<ScriptFileMarker>();
-            if (parseResult != null && parseResult.Count() > 0)
+            
+            if (!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);
+                return Array.Empty<ScriptFileMarker>();
+            }
+            
+            // 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.
-                    markers.Add(new ScriptFileMarker()
+                // vscode specific format for error markers.
+                markers.Add(new ScriptFileMarker
+                {
+                    Message = error.Message,
+                    Level = ScriptFileMarkerLevel.Error,
+                    ScriptRegion = new ScriptRegion
                     {
-                        Message = error.Message,
-                        Level = ScriptFileMarkerLevel.Error,
-                        ScriptRegion = new ScriptRegion()
-                        {
-                            File = scriptFile.FilePath,
-                            StartLineNumber = startLine,
-                            StartColumnNumber = startOffset,
-                            StartOffset = 0,
-                            EndLineNumber = endLine,
-                            EndColumnNumber = endOffset,
-                            EndOffset = 0
-                        }
-                    });
-                }
+                        File = scriptFile.FilePath,
+                        StartLineNumber = startLine,
+                        StartColumnNumber = startOffset,
+                        StartOffset = 0,
+                        EndLineNumber = endLine,
+                        EndColumnNumber = endOffset,
+                        EndOffset = 0
+                    }
+                });
             }
 
             return markers.ToArray();
         }
-
-        /// <summary>
-        /// 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)
+        
+        public DefinitionResult GetDefinition(string queryText, int index, int startLine, int startColumn, bool throwOnError = false)
         {
-            // try and show error from here.
-            DatabaseSymbol databaseSymbol = null;
+            //TODOKusto: API wasnt working properly, need to check that part.
+            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 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;
+            return null;
         }
-
-        /// <inheritdoc/>
-        public static LanguageServices.Contracts.CompletionItem[] GetAutoCompleteSuggestions(
-            ScriptDocumentInfo scriptDocumentInfo, Position textPosition, GlobalState schemaState,
-            bool throwOnError = false)
+        
+        public Hover GetHoverHelp(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, 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/>
-        public static Hover GetHoverHelp(ScriptDocumentInfo scriptDocumentInfo, Position textPosition,
-            GlobalState schemaState, bool throwOnError = false)
+        
+        private CompletionItemKind CreateCompletionItemKind(CompletionKind kustoKind)
         {
-            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);
-        }
-
-        /// <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)
+            switch (kustoKind)
             {
+                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
     }
 }