3644 Kusto Azure Token Refresh (#1081)

* 3644 Created KustoClient to handle interaction with Kusto servers. Refactored KustoDataSource and moved execution logic into KustoClient.

* 3644 Added RequestSecurityTokenParams, RequestSecurityTokenResponse, and SecurityTokenRequest to Kusto. Moved intellisense functions from KustoDataSource to KustoIntellisenseHelper. Added SchemaState to readonly property on KustoDataSource. Added catch block to Batch.cs to catch KustoRequestExceptions.

* 3644 Removed unused reference from ConnectionDetails and ConnectedBindingContext. Added UpdateAzureToken function to IKustoClient, KustoClient, IDataSource, KustoDataSource, and KustoDataSource. Added dataSource.Dispose to ReliableDataSourceConnection > Close. Added RefreshAzureToken to ConnectionService to refresh azure token.

* 3644 Removed unused properties from RequestSecurityTokenParams and RequestSecurityTokenResponse

* 3644 Added default to DbColumnWrapper to UnknownTypeName when null. Moved database query logic from KustoIntellisenseHelper to KustoClient. Moved KustoIntellisenseHelper data objects out of class. Changed SchemaState to load through separate tasks.

* 3644 Changed ReRunQuery logic in Kusto Batch to flip back to false if the query fails a second time so it can be rejected

* 3644 Updated GetAutoCompleteSuggestions and GetHoverHelp with changes from main

* 3644 Added AccountId to RequestSecurityTokenParams and set value in ConnectionService. Added throw to Batch.cs to ensure exceptions that are non-Unauthorized bubble up to ADS.

* 3644 Changed KustoUnauthorizedException to take original exception as inner exception. Changed catch block to only throw KustoUnauthorizedException when FailureCode is 401

* 3644 Renamed KustoUnauthorizedException to DataSourceUnauthorizedException. Moved logic to throw exception down into KustoClient. Changed retryLogic in Batch.cs to a decrementing count

* 3644 Changed logic in Batch.cs for throwing InvalidOperationException

src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionInfo.cs

@@ -156,6 +156,11 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
                 ReliableDataSourceConnection connection;
                 _connectionTypeToConnectionMap.TryRemove(type, out connection);
             }
-        } 
+        }
+
+        public void UpdateAzureToken(string token)
+        {
+            ConnectionDetails.AzureAccountToken = token;
+        }
     }
 }

src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs

@@ -63,9 +63,8 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
 
         /// <summary>
         /// Map from script URIs to ConnectionInfo objects
-        /// This is internal for testing access only
         /// </summary>
-        internal Dictionary<string, ConnectionInfo> OwnerToConnectionMap { get; } = new Dictionary<string, ConnectionInfo>();
+        private Dictionary<string, ConnectionInfo> OwnerToConnectionMap { get; } = new Dictionary<string, ConnectionInfo>();
 
         /// <summary>
         /// Database Lock manager instance
@@ -270,6 +269,25 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             return completeParams;
         }
 
+        internal void RefreshAzureToken(string ownerUri)
+        {
+            ConnectionInfo existingConnection = OwnerToConnectionMap[ownerUri];
+            
+            var requestMessage = new RequestSecurityTokenParams
+            {
+                AccountId = existingConnection.ConnectionDetails.GetOptionValue("azureAccount", string.Empty),
+                Authority = existingConnection.ConnectionDetails.GetOptionValue("azureTenantId", string.Empty),
+                Provider = "Azure",
+                Resource = "SQL"
+            };
+
+            var response = Instance.ServiceHost.SendRequest(SecurityTokenRequest.Type, requestMessage, true).Result;
+            existingConnection.UpdateAzureToken(response.Token);
+            
+            existingConnection.TryGetConnection(ConnectionType.Query, out var reliableDataSourceConnection);
+            reliableDataSourceConnection.GetUnderlyingConnection().UpdateAzureToken(response.Token);
+        }
+
         private void TryCloseConnectionTemporaryConnection(ConnectParams connectionParams, ConnectionInfo connectionInfo)
         {
             try

src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/ConnectionDetails.cs

@@ -3,7 +3,6 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
-using Microsoft.Kusto.ServiceLayer.Utility;
 using Microsoft.SqlTools.Utility;
 
 namespace Microsoft.Kusto.ServiceLayer.Connection.Contracts

src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/RequestSecurityTokenParams.cs

@@ -0,0 +1,22 @@
+namespace Microsoft.Kusto.ServiceLayer.Connection.Contracts
+{
+    public class RequestSecurityTokenParams
+    {
+        /// <summary>
+        /// Gets or sets the address of the authority to issue token.
+        /// </summary>
+        public string Authority { get; set; }
+
+        /// <summary>
+        /// Gets or sets the provider that indicates the type of linked account to query.
+        /// </summary>
+        public string Provider { get; set; }
+
+        /// <summary>
+        /// Gets or sets the identifier of the target resource that is the recipient of the requested token.
+        /// </summary>
+        public string Resource { get; set; }
+        
+        public string AccountId { get; set; }
+    }
+}

src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/RequestSecurityTokenResponse.cs

@@ -0,0 +1,15 @@
+namespace Microsoft.Kusto.ServiceLayer.Connection.Contracts
+{
+    public class RequestSecurityTokenResponse
+    {
+        /// <summary>
+        /// Gets or sets the key that uniquely identifies a particular linked account.
+        /// </summary>
+        public string AccountKey { get; set; }
+
+        /// <summary>
+        /// Gets or sets the access token.
+        /// </summary>
+        public string Token { get; set; }
+    }
+}

src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/SecurityTokenRequest.cs

@@ -0,0 +1,15 @@
+using Microsoft.SqlTools.Hosting.Protocol.Contracts;
+
+namespace Microsoft.Kusto.ServiceLayer.Connection.Contracts
+{
+    /// <summary>
+    /// SecurityToken Request mapping entry 
+    /// </summary>
+    public class SecurityTokenRequest
+    {
+        public static readonly
+            RequestType<RequestSecurityTokenParams, RequestSecurityTokenResponse> Type =
+                RequestType<RequestSecurityTokenParams, RequestSecurityTokenResponse>.Create(
+                    "account/securityTokenRequest");
+    }
+}

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

@@ -7,23 +7,15 @@ using System.Collections.Generic;
 using System.Threading;
 using System.Data;
 using System.Threading.Tasks;
+using Kusto.Language;
 using Microsoft.Kusto.ServiceLayer.Utility;
-using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
 using Microsoft.Kusto.ServiceLayer.DataSource.Metadata;
-using Microsoft.Kusto.ServiceLayer.LanguageServices;
-using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
-using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
-using Microsoft.Kusto.ServiceLayer.LanguageServices.Completion;
 
 namespace Microsoft.Kusto.ServiceLayer.DataSource
 {
     /// <inheritdoc cref="IDataSource"/>
     public abstract class DataSourceBase : IDataSource
     {
-        protected Object dataSourceLock = new Object();
-
-        private string _database;
-
         #region IDisposable
 
         /// <summary>
@@ -87,17 +79,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// <inheritdoc/>
         public abstract void UpdateDatabase(string databaseName);
 
-        /// <inheritdoc/>
-        public abstract CompletionItem[] GetAutoCompleteSuggestions(ScriptDocumentInfo queryText, Position index, bool throwOnError = false);
-        /// <inheritdoc/>
-        public abstract Hover GetHoverHelp(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, bool throwOnError = false);
-        
-        /// <inheritdoc/>
-        public abstract DefinitionResult GetDefinition(string queryText, int index, int startLine, int startColumn, bool throwOnError = false);
-
-        /// <inheritdoc/>
-        public abstract ScriptFileMarker[] GetSemanticMarkers(ScriptParseInfo parseInfo, ScriptFile scriptFile, string queryText);
-
         /// <inheritdoc/>
         public abstract Task<bool> Exists();
 
@@ -108,27 +89,16 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
 
         public abstract string GenerateExecuteFunctionScript(string functionName);
 
+        public abstract void UpdateAzureToken(string azureToken);
+
         /// <inheritdoc/>
         public DataSourceType DataSourceType { get; protected set; }
-
+        
         /// <inheritdoc/>
-        public string ClusterName { get; protected set; }
+        public abstract string ClusterName { get; }
 
-        /// <inheritdoc/>
+        public abstract string DatabaseName { get; }
-        public string DatabaseName { 
+        public abstract GlobalState SchemaState { get; }
-            get
-            {
-                return _database;
-            }
-            
-            set
-            {
-                lock(dataSourceLock)
-                {
-                    _database = value;
-                }
-            }
-        }
 
         #endregion
     }

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

@@ -22,7 +22,8 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
             {
                 case DataSourceType.Kusto:
                 {
-                    return new KustoDataSource(connectionString, azureAccountToken);
+                    var kustoClient = new KustoClient(connectionString, azureAccountToken);
+                    return new KustoDataSource(kustoClient);
                 }
 
                 default:

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

@@ -5,10 +5,9 @@
 
 using System;
 using System.Collections.Generic;
-using System.Threading;
-using System.Threading.Tasks;
 using System.Linq;
 using Kusto.Language;
+using KustoDiagnostic = Kusto.Language.Diagnostic;
 using Kusto.Language.Editor;
 using Kusto.Language.Syntax;
 using Kusto.Language.Symbols;
@@ -24,41 +23,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
     /// </summary>
     public class KustoIntellisenseHelper
     {
-        public class ShowDatabasesResult
-        {
-            public string DatabaseName;
-            public string PersistentStorage;
-            public string Version;
-            public bool IsCurrent;
-            public string DatabaseAccessMode;
-            public string PrettyName;
-            public bool CurrentUserIsUnrestrictedViewer;
-            public string DatabaseId;
-        }
-
-        public class ShowDatabaseSchemaResult
-        {
-            public string DatabaseName;
-            public string TableName;
-            public string ColumnName;
-            public string ColumnType;
-            public bool IsDefaultTable;
-            public bool IsDefaultColumn;
-            public string PrettyName;
-            public string Version;
-            public string Folder;
-            public string DocName;
-        }
-
-        public class ShowFunctionsResult
-        {
-            public string Name;
-            public string Parameters;
-            public string Body;
-            public string Folder;
-            public string DocString;
-        }
-
         /// <summary>
         /// Convert CLR type name into a Kusto scalar type.
         /// </summary>
@@ -161,20 +125,20 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
         /// <summary>
         /// Loads the schema for the specified databasea into a a <see cref="DatabaseSymbol"/>.
         /// </summary>
-        private static async Task<DatabaseSymbol> LoadDatabaseAsync(IDataSource dataSource, string databaseName, bool throwOnError = false)
+        private static DatabaseSymbol LoadDatabaseAsync(IEnumerable<ShowDatabaseSchemaResult> tableSchemas,
+            IEnumerable<ShowFunctionsResult> functionSchemas,
+            string databaseName)
         {
-            var members = new List<Symbol>();
-            CancellationTokenSource source = new CancellationTokenSource();
-            CancellationToken cancellationToken = source.Token;
-
-            var tableSchemas = await dataSource.ExecuteControlCommandAsync<ShowDatabaseSchemaResult>($".show database {databaseName} schema", throwOnError, cancellationToken).ConfigureAwait(false);
             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();
@@ -182,9 +146,10 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
                 members.Add(tableSymbol);
             }
 
-            var functionSchemas = await dataSource.ExecuteControlCommandAsync<ShowFunctionsResult>(".show functions", throwOnError, cancellationToken).ConfigureAwait(false);
             if (functionSchemas == null)
+            {
                 return null;
+            }
 
             foreach (var fun in functionSchemas)
             {
@@ -193,8 +158,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
                 members.Add(functionSymbol);
             }
 
-            var databaseSymbol = new DatabaseSymbol(databaseName, members);
+            return new DatabaseSymbol(databaseName, members);
-            return databaseSymbol;
         }
 
         public static CompletionItemKind CreateCompletionItemKind(CompletionKind kustoKind)
@@ -234,33 +198,41 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
         /// <summary>
         /// Gets default keyword when user if not connected to any Kusto cluster.
         /// </summary>
-        public static LanguageServices.Contracts.CompletionItem[] GetDefaultKeywords(ScriptDocumentInfo scriptDocumentInfo, Position textDocumentPosition){
+        public static LanguageServices.Contracts.CompletionItem[] GetDefaultKeywords(
-                var kustoCodeService = new KustoCodeService(scriptDocumentInfo.Contents, GlobalState.Default);
+            ScriptDocumentInfo scriptDocumentInfo, Position textDocumentPosition)
-                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 kustoCodeService = new KustoCodeService(scriptDocumentInfo.Contents, GlobalState.Default);
-                var completion = kustoCodeService.GetCompletionItems(position);
+            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>();
+            List<LanguageServices.Contracts.CompletionItem> completions =
-                foreach (var autoCompleteItem in completion.Items)
+                new List<LanguageServices.Contracts.CompletionItem>();
-                {
+            foreach (var autoCompleteItem in completion.Items)
-                    var label = autoCompleteItem.DisplayText;
+            {
-                    // convert the completion item candidates into vscode format CompletionItems
+                var label = autoCompleteItem.DisplayText;
-                    completions.Add(AutoCompleteHelper.CreateCompletionItem(label, label + " keyword", label, CompletionItemKind.Keyword, scriptDocumentInfo.StartLine, scriptDocumentInfo.StartColumn, textDocumentPosition.Character));
+                // 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();
+            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){
+        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)
@@ -295,22 +267,27 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
         /// <summary>
         /// Loads the schema for the specified database and returns a new <see cref="GlobalState"/> with the database added or updated.
         /// </summary>
-        public static async Task<GlobalState> AddOrUpdateDatabaseAsync(IDataSource dataSource, GlobalState globals, string databaseName, string clusterName, bool throwOnError)
+        public static GlobalState AddOrUpdateDatabase(IEnumerable<ShowDatabaseSchemaResult> tableSchemas,
-        {   // try and show error from here.
+            IEnumerable<ShowFunctionsResult> functionSchemas, GlobalState globals,
+            string databaseName, string clusterName)
+        {
+            // try and show error from here.
             DatabaseSymbol databaseSymbol = null;
 
-            if(databaseName != null){
+            if (databaseName != null)
-                databaseSymbol = await LoadDatabaseAsync(dataSource, databaseName, throwOnError).ConfigureAwait(false);
+            {
+                databaseSymbol = LoadDatabaseAsync(tableSchemas, functionSchemas, databaseName);
             }
 
-            if(databaseSymbol == null){
+            if (databaseSymbol == null)
+            {
                 return globals;
             }
 
             var cluster = globals.GetCluster(clusterName);
             if (cluster == null)
             {
-                cluster = new ClusterSymbol(clusterName, new[] { databaseSymbol }, isOpen: true);
+                cluster = new ClusterSymbol(clusterName, new[] {databaseSymbol}, isOpen: true);
                 globals = globals.AddOrUpdateCluster(cluster);
             }
             else
@@ -324,5 +301,116 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
             return globals;
         }
 
+        /// <inheritdoc/>
+        public static LanguageServices.Contracts.CompletionItem[] GetAutoCompleteSuggestions(
+            ScriptDocumentInfo scriptDocumentInfo, Position textPosition, GlobalState schemaState,
+            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
+
+            var codeBlock = script.GetBlockAtPosition(position);
+            var completion = codeBlock.Service.GetCompletionItems(position);
+            scriptDocumentInfo.ScriptParseInfo.CurrentSuggestions =
+                completion.Items; // this is declaration item so removed for now, but keep the info when api gets updated
+
+            var completions = new List<LanguageServices.Contracts.CompletionItem>();
+            foreach (var autoCompleteItem in completion.Items)
+            {
+                var label = autoCompleteItem.DisplayText;
+                var insertText = autoCompleteItem.Kind == CompletionKind.Table
+                    ? KustoQueryUtils.EscapeName(label)
+                    : label;
+                
+                var completionKind = CreateCompletionItemKind(autoCompleteItem.Kind);
+                completions.Add(AutoCompleteHelper.CreateCompletionItem(label, autoCompleteItem.Kind.ToString(),
+                    insertText, completionKind, scriptDocumentInfo.StartLine, scriptDocumentInfo.StartColumn,
+                    textPosition.Character));
+            }
+
+            return completions.ToArray();
+        }
+
+        /// <inheritdoc/>
+        public static Hover GetHoverHelp(ScriptDocumentInfo scriptDocumentInfo, Position textPosition,
+            GlobalState schemaState, bool throwOnError = false)
+        {
+            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)
+            {
+            }
+
+            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/DataSourceIntellisense/ShowDatabaseSchemaResult.cs

@@ -0,0 +1,16 @@
+namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
+{
+    public class ShowDatabaseSchemaResult
+    {
+        public string DatabaseName;
+        public string TableName;
+        public string ColumnName;
+        public string ColumnType;
+        public bool IsDefaultTable;
+        public bool IsDefaultColumn;
+        public string PrettyName;
+        public string Version;
+        public string Folder;
+        public string DocName;
+    }
+}

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

@@ -0,0 +1,14 @@
+namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
+{
+    public class ShowDatabasesResult
+    {
+        public string DatabaseName;
+        public string PersistentStorage;
+        public string Version;
+        public bool IsCurrent;
+        public string DatabaseAccessMode;
+        public string PrettyName;
+        public bool CurrentUserIsUnrestrictedViewer;
+        public string DatabaseId;
+    }
+}

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

@@ -0,0 +1,11 @@
+namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
+{
+    public class ShowFunctionsResult
+    {
+        public string Name;
+        public string Parameters;
+        public string Body;
+        public string Folder;
+        public string DocString;
+    }
+}

src/Microsoft.Kusto.ServiceLayer/DataSource/Exceptions/DataSourceUnauthorizedException.cs

@@ -0,0 +1,11 @@
+using System;
+
+namespace Microsoft.Kusto.ServiceLayer.DataSource.Exceptions
+{
+    public class DataSourceUnauthorizedException : Exception
+    {
+        public DataSourceUnauthorizedException(Exception ex) : base (ex.Message, ex)
+        {
+        }
+    }
+}

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

@@ -3,12 +3,8 @@ using System.Collections.Generic;
 using System.Data;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
+using Kusto.Language;
 using Microsoft.Kusto.ServiceLayer.DataSource.Metadata;
-using Microsoft.Kusto.ServiceLayer.LanguageServices;
-using Microsoft.Kusto.ServiceLayer.LanguageServices.Completion;
-using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
-using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
 
 namespace Microsoft.Kusto.ServiceLayer.DataSource
 {
@@ -30,7 +26,9 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// <summary>
         /// The current database name, if there is one.
         /// </summary>
-        string DatabaseName { get; set; }
+        string DatabaseName { get; }
+        
+        GlobalState SchemaState { get; }
 
         /// <summary>
         /// Executes a query.
@@ -88,29 +86,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// <param name="updateDatabase">Object metadata.</param>
         void UpdateDatabase(string databaseName);
 
-        /// <summary>
-        /// Gets autocomplete suggestions at given position.
-        /// </summary>
-        /// <param name="GetAutoCompleteSuggestions">Object metadata.</param>
-        CompletionItem[] GetAutoCompleteSuggestions(ScriptDocumentInfo queryText, Position index, bool throwOnError = false);
-        /// <summary>
-        /// Gets quick info hover tooltips for the current position.
-        /// </summary>
-        /// <param name="GetHoverHelp">Object metadata.</param>
-        Hover GetHoverHelp(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, bool throwOnError = false);
-
-        /// <summary>
-        /// Gets definition for a selected query text.
-        /// </summary>
-        /// <param name="GetDefinition">Object metadata.</param>
-        DefinitionResult GetDefinition(string queryText, int index, int startLine, int startColumn, bool throwOnError = false);
-        
-        /// <summary>
-        /// Gets a list of semantic diagnostic marks for the provided script file
-        /// </summary>
-        /// <param name="GetSemanticMarkers">Object metadata.</param>
-        ScriptFileMarker[] GetSemanticMarkers(ScriptParseInfo parseInfo, ScriptFile scriptFile, string queryText);
-
         /// <summary>
         /// Tells whether the data source exists.
         /// </summary>
@@ -136,5 +111,11 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// <param name="functionName"></param>
         /// <returns></returns>
         string GenerateExecuteFunctionScript(string functionName);
+
+        /// <summary>
+        /// Updates Azure Token
+        /// </summary>
+        /// <param name="azureToken"></param>
+        void UpdateAzureToken(string azureToken);
     }
 }

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

@@ -0,0 +1,46 @@
+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; }
+        
+        void UpdateAzureToken(string azureAccountToken);
+        
+        IDataReader ExecuteQuery(string query, CancellationToken cancellationToken, string databaseName = null);
+
+        /// <summary>
+        /// Executes a query or command against a kusto cluster and returns a sequence of result row instances.
+        /// </summary>
+        Task<IEnumerable<T>> ExecuteControlCommandAsync<T>(string command, bool throwOnError, CancellationToken cancellationToken);
+
+        /// <summary>
+        /// Executes a query.
+        /// </summary>
+        /// <param name="query">The query.</param>
+        /// <returns>The results.</returns>
+        Task<IDataReader> ExecuteQueryAsync(string query, CancellationToken cancellationToken, string databaseName = null);
+
+        /// <summary>
+        /// Executes a Kusto control command.
+        /// </summary>
+        /// <param name="command">The command.</param>
+        void ExecuteControlCommand(string command);
+
+        void UpdateDatabase(string databaseName);
+        
+        void Dispose();
+    }
+}

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

@@ -0,0 +1,260 @@
+using System;
+using System.Collections.Generic;
+using System.Data;
+using System.Diagnostics;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+using Kusto.Cloud.Platform.Data;
+using Kusto.Data;
+using Kusto.Data.Common;
+using Kusto.Data.Data;
+using Kusto.Data.Exceptions;
+using Kusto.Data.Net.Client;
+using Kusto.Language;
+using Kusto.Language.Editor;
+using Microsoft.Data.SqlClient;
+using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
+using Microsoft.Kusto.ServiceLayer.DataSource.Exceptions;
+using Microsoft.Kusto.ServiceLayer.Utility;
+
+namespace Microsoft.Kusto.ServiceLayer.DataSource
+{
+    public class KustoClient : IKustoClient
+    {
+        [DebuggerBrowsable(DebuggerBrowsableState.Never)]
+        private ICslAdminProvider _kustoAdminProvider;
+
+        [DebuggerBrowsable(DebuggerBrowsableState.Never)]
+        private ICslQueryProvider _kustoQueryProvider;
+
+        /// <summary>
+        /// SchemaState used for getting intellisense info.
+        /// </summary>
+        public GlobalState SchemaState { get; private set; }
+
+        public string ClusterName { get; }
+        public string DatabaseName { get; private set; }
+
+        public KustoClient(string connectionString, string azureAccountToken)
+        {
+            ClusterName = GetClusterName(connectionString);
+            var databaseName = new SqlConnectionStringBuilder(connectionString).InitialCatalog;
+            Initialize(ClusterName, databaseName, azureAccountToken);
+            DatabaseName = string.IsNullOrWhiteSpace(databaseName) ? GetFirstDatabaseName() : databaseName;
+            SchemaState = LoadSchemaState();
+        }
+
+        private GlobalState LoadSchemaState()
+        {
+            CancellationTokenSource source = new CancellationTokenSource();
+
+            IEnumerable<ShowDatabaseSchemaResult> tableSchemas = Enumerable.Empty<ShowDatabaseSchemaResult>();
+            IEnumerable<ShowFunctionsResult> functionSchemas = Enumerable.Empty<ShowFunctionsResult>();
+
+            Parallel.Invoke(() =>
+            {
+                tableSchemas = ExecuteControlCommandAsync<ShowDatabaseSchemaResult>(
+                    $".show database {DatabaseName} schema",
+                    false, source.Token).Result;
+            }, () =>
+            {
+                functionSchemas = ExecuteControlCommandAsync<ShowFunctionsResult>(".show functions", false,
+                    source.Token).Result;
+            });
+
+            return KustoIntellisenseHelper.AddOrUpdateDatabase(tableSchemas, functionSchemas,
+                GlobalState.Default,
+                DatabaseName, ClusterName);
+        }
+
+        private void Initialize(string clusterName, string databaseName, string azureAccountToken)
+        {
+            var stringBuilder = GetKustoConnectionStringBuilder(clusterName, databaseName, azureAccountToken, "", "");
+            _kustoQueryProvider = KustoClientFactory.CreateCslQueryProvider(stringBuilder);
+            _kustoAdminProvider = KustoClientFactory.CreateCslAdminProvider(stringBuilder);
+        }
+
+        public void UpdateAzureToken(string azureAccountToken)
+        {
+            _kustoQueryProvider.Dispose();
+            _kustoAdminProvider.Dispose();
+            Initialize(ClusterName, DatabaseName, azureAccountToken);
+        }
+
+        /// <summary>
+        /// Extracts the cluster name from the connectionstring. The string looks like the following:
+        /// "Data Source=clustername.kusto.windows.net;User ID=;Password=;Pooling=False;Application Name=azdata-GeneralConnection"
+        /// <summary>
+        /// <param name="connectionString">A connection string coming over the Data management protocol</param>
+        private string GetClusterName(string connectionString)
+        {
+            var csb = new SqlConnectionStringBuilder(connectionString);
+
+            // If there is no https:// prefix, add it
+            Uri uri;
+            if ((Uri.TryCreate(csb.DataSource, UriKind.Absolute, out uri) ||
+                 Uri.TryCreate("https://" + csb.DataSource, UriKind.Absolute, out uri)) &&
+                (uri.Scheme == Uri.UriSchemeHttp || uri.Scheme == Uri.UriSchemeHttps))
+            {
+                return uri.AbsoluteUri;
+            }
+
+            throw new ArgumentException("Expected a URL of the form clustername.kusto.windows.net");
+        }
+
+        private KustoConnectionStringBuilder GetKustoConnectionStringBuilder(string clusterName, string databaseName,
+            string userToken, string applicationClientId, string applicationKey)
+        {
+            ValidationUtils.IsNotNull(clusterName, nameof(clusterName));
+            ValidationUtils.IsTrue<ArgumentException>(
+                !string.IsNullOrWhiteSpace(userToken)
+                || (!string.IsNullOrWhiteSpace(applicationClientId) && !string.IsNullOrWhiteSpace(applicationKey)),
+                $"the Kusto authentication is not specified - either set {nameof(userToken)}, or set {nameof(applicationClientId)} and {nameof(applicationKey)}");
+
+            var kcsb = new KustoConnectionStringBuilder
+            {
+                DataSource = clusterName,
+
+                // Perform federated auth based on the AAD user token, or based on the AAD application client id and key.
+                FederatedSecurity = true
+            };
+
+            if (!string.IsNullOrWhiteSpace(databaseName))
+            {
+                kcsb.InitialCatalog = databaseName;
+            }
+
+            if (!string.IsNullOrWhiteSpace(userToken))
+            {
+                kcsb.UserToken = userToken;
+            }
+
+            if (!string.IsNullOrWhiteSpace(applicationClientId))
+            {
+                kcsb.ApplicationClientId = applicationClientId;
+            }
+
+            if (!string.IsNullOrWhiteSpace(applicationKey))
+            {
+                kcsb.ApplicationKey = applicationKey;
+            }
+
+            return kcsb;
+        }
+
+        /// <summary>
+        /// Extracts the database name from the connectionString if it exists
+        /// otherwise it takes the first database name from the server
+        /// </summary>
+        /// <param name="connectionString"></param>
+        /// <returns>Database Name</returns>
+        private string GetFirstDatabaseName()
+        {
+            var source = new CancellationTokenSource();
+            string query = ".show databases | project DatabaseName";
+
+            using (var reader = ExecuteQuery(query, source.Token))
+            {
+                var rows = reader.ToEnumerable();
+                var row = rows?.FirstOrDefault();
+                return row?[0].ToString() ?? string.Empty;
+            }
+        }
+
+        public IDataReader ExecuteQuery(string query, CancellationToken cancellationToken, string databaseName = null)
+        {
+            ValidationUtils.IsArgumentNotNullOrWhiteSpace(query, nameof(query));
+
+            var clientRequestProperties = new ClientRequestProperties
+            {
+                ClientRequestId = Guid.NewGuid().ToString()
+            };
+            clientRequestProperties.SetOption(ClientRequestProperties.OptionNoTruncation, true);
+            cancellationToken.Register(() => CancelQuery(clientRequestProperties.ClientRequestId));
+
+            var kustoCodeService = new KustoCodeService(query);
+            var minimalQuery = kustoCodeService.GetMinimalText(MinimalTextKind.RemoveLeadingWhitespaceAndComments);
+
+            try
+            {
+                IDataReader origReader = _kustoQueryProvider.ExecuteQuery(
+                    KustoQueryUtils.IsClusterLevelQuery(minimalQuery) ? "" : databaseName,
+                    minimalQuery,
+                    clientRequestProperties);
+                
+                return new KustoResultsReader(origReader);
+            }
+            catch (KustoRequestException exception) when (exception.FailureCode == 401) // Unauthorized
+            {
+                throw new DataSourceUnauthorizedException(exception);
+            }
+        }
+
+        /// <summary>
+        /// Executes a query or command against a kusto cluster and returns a sequence of result row instances.
+        /// </summary>
+        public async Task<IEnumerable<T>> ExecuteControlCommandAsync<T>(string command, bool throwOnError,
+            CancellationToken cancellationToken)
+        {
+            try
+            {
+                var resultReader = await ExecuteQueryAsync(command, cancellationToken, DatabaseName);
+                var results = KustoDataReaderParser.ParseV1(resultReader, null);
+                var tableReader = results[WellKnownDataSet.PrimaryResult].Single().TableData.CreateDataReader();
+                return new ObjectReader<T>(tableReader);
+            }
+            catch (DataSourceUnauthorizedException)
+            {
+                throw;
+            }
+            catch (Exception) when (!throwOnError)
+            {
+                return null;
+            }
+        }
+
+        /// <summary>
+        /// Executes a query.
+        /// </summary>
+        /// <param name="query">The query.</param>
+        /// <returns>The results.</returns>
+        public Task<IDataReader> ExecuteQueryAsync(string query, CancellationToken cancellationToken,
+            string databaseName = null)
+        {
+            var reader = ExecuteQuery(query, cancellationToken, databaseName);
+            return Task.FromResult(reader);
+        }
+
+        private void CancelQuery(string clientRequestId)
+        {
+            var query = $".cancel query \"{clientRequestId}\"";
+            ExecuteControlCommand(query);
+        }
+
+        /// <summary>
+        /// Executes a Kusto control command.
+        /// </summary>
+        /// <param name="command">The command.</param>
+        public void ExecuteControlCommand(string command)
+        {
+            ValidationUtils.IsArgumentNotNullOrWhiteSpace(command, nameof(command));
+
+            using (var adminOutput = _kustoAdminProvider.ExecuteControlCommand(command, null))
+            {
+            }
+        }
+
+        public void UpdateDatabase(string databaseName)
+        {
+            DatabaseName = databaseName;
+            SchemaState = LoadSchemaState();
+        }
+
+        public void Dispose()
+        {
+            _kustoQueryProvider.Dispose();
+            _kustoAdminProvider.Dispose();
+        }
+    }
+}

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

@@ -6,8 +6,6 @@ using System.Collections.Generic;
 using System.Threading;
 using System.Collections.Concurrent;
 using System.Data;
-using System.Data.SqlClient;
-using System.Diagnostics;
 using System.Globalization;
 using System.Linq;
 using System.Text;
@@ -15,20 +13,11 @@ using Newtonsoft.Json;
 using System.Threading.Tasks;
 using Kusto.Cloud.Platform.Data;
 using Kusto.Data;
-using Kusto.Data.Common;
 using Kusto.Data.Data;
-using Kusto.Data.Net.Client;
 using Kusto.Language;
-using KustoDiagnostic = Kusto.Language.Diagnostic;
-using Kusto.Language.Editor;
-using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
 using Microsoft.Kusto.ServiceLayer.DataSource.Metadata;
 using Microsoft.Kusto.ServiceLayer.DataSource.Models;
-using Microsoft.Kusto.ServiceLayer.LanguageServices;
-using Microsoft.Kusto.ServiceLayer.LanguageServices.Completion;
-using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
 using Microsoft.Kusto.ServiceLayer.Utility;
-using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
 
 namespace Microsoft.Kusto.ServiceLayer.DataSource
 {
@@ -37,9 +26,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
     /// </summary>
     public class KustoDataSource : DataSourceBase
     {
-        private ICslQueryProvider _kustoQueryProvider;
+        private IKustoClient _kustoClient;
-
-        private ICslAdminProvider _kustoAdminProvider;
 
         /// <summary>
         /// List of databases.
@@ -66,6 +53,12 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// </summary>
         private ConcurrentDictionary<string, IEnumerable<FunctionMetadata>> _functionMetadata = new ConcurrentDictionary<string, IEnumerable<FunctionMetadata>>();
 
+        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"; 
 
@@ -84,122 +77,14 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// <summary>
         /// Prevents a default instance of the <see cref="IDataSource"/> class from being created.
         /// </summary>
-        public KustoDataSource(string connectionString, string azureAccountToken)
+        public KustoDataSource(IKustoClient kustoClient)
         {
-            ClusterName = GetClusterName(connectionString);
+            _kustoClient = kustoClient;
-            UserToken = azureAccountToken;
-            DatabaseName = GetDatabaseName(connectionString);
-            SchemaState = Task.Run(() =>
-                KustoIntellisenseHelper.AddOrUpdateDatabaseAsync(this, GlobalState.Default, DatabaseName, ClusterName,
-                    throwOnError: false)).Result;
             // Check if a connection can be made
             ValidationUtils.IsTrue<ArgumentException>(Exists().Result,
                 $"Unable to connect. ClusterName = {ClusterName}, DatabaseName = {DatabaseName}");
         }
-
+        
-        /// <summary>
-        /// Extracts the cluster name from the connectionstring. The string looks like the following:
-        /// "Data Source=clustername.kusto.windows.net;User ID=;Password=;Pooling=False;Application Name=azdata-GeneralConnection"
-        /// <summary>
-        /// <param name="connectionString">A connection string coming over the Data management protocol</param>
-        private static string GetClusterName(string connectionString)
-        {
-            var csb = new SqlConnectionStringBuilder(connectionString);
-
-            // If there is no https:// prefix, add it
-            Uri uri;
-            if ((Uri.TryCreate(csb.DataSource, UriKind.Absolute, out uri) || Uri.TryCreate("https://" + csb.DataSource, UriKind.Absolute, out uri)) &&
-                (uri.Scheme == Uri.UriSchemeHttp || uri.Scheme == Uri.UriSchemeHttps))
-            {
-                return uri.AbsoluteUri;
-            }
-
-            throw new ArgumentException("Expected a URL of the form clustername.kusto.windows.net");
-        }
-
-        /// <summary>
-        /// Extracts the database name from the connectionString if it exists
-        /// otherwise it takes the first database name from the server
-        /// </summary>
-        /// <param name="connectionString"></param>
-        /// <returns>Database Name</returns>
-        private string GetDatabaseName(string connectionString)
-        {
-            var csb = new SqlConnectionStringBuilder(connectionString);
-
-            if (!string.IsNullOrWhiteSpace(csb.InitialCatalog))
-            {
-                return csb.InitialCatalog;    
-            }
-            
-            CancellationTokenSource source = new CancellationTokenSource();
-            CancellationToken token = source.Token;
-
-            string query = ".show databases | project DatabaseName";
-
-            using (var reader = ExecuteQuery(query, token))
-            {
-                var rows = reader.ToEnumerable();
-                var row = rows?.FirstOrDefault();
-                return row?[0].ToString() ?? string.Empty;
-            }
-        }
-
-        /// <summary>
-        /// SchemaState used for getting intellisense info.
-        /// </summary>
-        public GlobalState SchemaState { get; private set; }
-
-        /// <summary>
-        /// The AAD user token.
-        /// </summary>
-        public string UserToken { get; private set; }
-
-        /// <summary>
-        /// The AAD application client id.
-        /// </summary>
-        public string ApplicationClientId { get; private set; }
-
-        /// <summary>
-        /// The AAD application client key.
-        /// </summary>
-        public string ApplicationKey { get; private set; }
-
-        // The Kusto query provider.
-        [DebuggerBrowsable(DebuggerBrowsableState.Never)]
-        private ICslQueryProvider KustoQueryProvider
-        {
-            get
-            {
-                if (_kustoQueryProvider == null)
-                {
-                    var kcsb = GetKustoConnectionStringBuilder();
-                    _kustoQueryProvider = KustoClientFactory.CreateCslQueryProvider(kcsb);
-                }
-
-                return _kustoQueryProvider;
-            }
-        }
-
-        [DebuggerBrowsable(DebuggerBrowsableState.Never)]
-        private ICslAdminProvider KustoAdminProvider
-        {
-            get
-            {
-                if (_kustoAdminProvider == null)
-                {
-                    var kcsb = GetKustoConnectionStringBuilder();
-                    _kustoAdminProvider = KustoClientFactory.CreateCslAdminProvider(kcsb);
-                    if (!string.IsNullOrWhiteSpace(DatabaseName))
-                    {
-                        _kustoAdminProvider.DefaultDatabaseName = DatabaseName;
-                    }
-                }
-
-                return _kustoAdminProvider;
-            }
-        }
-
         /// <summary>
         /// Disposes resources.
         /// </summary>
@@ -209,11 +94,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
             // Dispose managed resources.
             if (disposing)
             {
-                _kustoQueryProvider?.Dispose();
+                _kustoClient.Dispose();
-                _kustoQueryProvider = null;
-
-                _kustoAdminProvider?.Dispose();
-                _kustoAdminProvider = null;
             }
 
             base.Dispose(disposing);
@@ -228,42 +109,10 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// <returns>The results.</returns>
         public override Task<IDataReader> ExecuteQueryAsync(string query, CancellationToken cancellationToken, string databaseName = null)
         {
-            var reader = ExecuteQuery(query, cancellationToken, databaseName);
+            var reader = _kustoClient.ExecuteQuery(query, cancellationToken, databaseName);
             return Task.FromResult(reader);
         }
 
-        private IDataReader ExecuteQuery(string query, CancellationToken cancellationToken, string databaseName = null)
-        {
-            ValidationUtils.IsArgumentNotNullOrWhiteSpace(query, nameof(query));
-
-            var clientRequestProperties = new ClientRequestProperties
-            {
-                ClientRequestId = Guid.NewGuid().ToString()
-            };
-            clientRequestProperties.SetOption(ClientRequestProperties.OptionNoTruncation, true);
-
-            if(cancellationToken != null)
-            {
-                cancellationToken.Register(() => CancelQuery(clientRequestProperties.ClientRequestId));
-            }
-            
-            var kustoCodeService = new KustoCodeService(query);
-            query = kustoCodeService.GetMinimalText(MinimalTextKind.RemoveLeadingWhitespaceAndComments);
-
-            IDataReader origReader = KustoQueryProvider.ExecuteQuery(
-                KustoQueryUtils.IsClusterLevelQuery(query) ? "" : databaseName, 
-                query, 
-                clientRequestProperties);
-
-            return new KustoResultsReader(origReader);
-        }
-
-        private void CancelQuery(string clientRequestId)
-        {
-            var query = $".cancel query \"{clientRequestId}\"";
-            ExecuteControlCommand(query);
-        }
-
         /// <inheritdoc/>
         public override async Task<bool> Exists()
         {
@@ -282,68 +131,17 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         }
 
         #endregion
-
+        
-        /// <summary>
-        /// Executes a Kusto control command.
-        /// </summary>
-        /// <param name="command">The command.</param>
-        public void ExecuteControlCommand(string command)
-        {
-            ValidationUtils.IsArgumentNotNullOrWhiteSpace(command, nameof(command));
-
-            using (var adminOutput = KustoAdminProvider.ExecuteControlCommand(command, null))
-            {
-            }
-        }
-
-        private KustoConnectionStringBuilder GetKustoConnectionStringBuilder()
-        {
-            ValidationUtils.IsNotNull(ClusterName, nameof(ClusterName));
-            ValidationUtils.IsTrue<ArgumentException>(
-                !string.IsNullOrWhiteSpace(UserToken)
-                || (!string.IsNullOrWhiteSpace(ApplicationClientId) && !string.IsNullOrWhiteSpace(ApplicationKey)),
-                $"the Kusto authentication is not specified - either set {nameof(UserToken)}, or set {nameof(ApplicationClientId)} and {nameof(ApplicationKey)}");
-
-            var kcsb = new KustoConnectionStringBuilder
-            {
-                DataSource = ClusterName,
-
-                // Perform federated auth based on the AAD user token, or based on the AAD application client id and key.
-                FederatedSecurity = true
-            };
-
-            if (!string.IsNullOrWhiteSpace(DatabaseName))
-            {
-                kcsb.InitialCatalog = DatabaseName;
-            }
-
-            if (!string.IsNullOrWhiteSpace(UserToken))
-            {
-                kcsb.UserToken = UserToken;
-            }
-
-            if (!string.IsNullOrWhiteSpace(ApplicationClientId))
-            {
-                kcsb.ApplicationClientId = ApplicationClientId;
-            }
-
-            if (!string.IsNullOrWhiteSpace(ApplicationKey))
-            {
-                kcsb.ApplicationKey = ApplicationKey;
-            }
-
-            return kcsb;
-        }
-
         #region IDataSource
 
-        protected DiagnosticsInfo GetClusterDiagnostics(){
+        private DiagnosticsInfo GetClusterDiagnostics()
+        {
             CancellationTokenSource source = new CancellationTokenSource();
             CancellationToken token = source.Token;
             DiagnosticsInfo clusterDiagnostics = new DiagnosticsInfo();
 
             var query =  ".show diagnostics | extend Passed= (IsHealthy) and not(IsScaleOutRequired) | extend Summary = strcat('Cluster is ', iif(Passed, '', 'NOT'), 'healthy.'),Details=pack('MachinesTotal', MachinesTotal, 'DiskCacheCapacity', round(ClusterDataCapacityFactor,1)) | project Action = 'Cluster Diagnostics', Category='Info', Summary, Details;";
-            using (var reader = ExecuteQuery(query, token))
+            using (var reader = _kustoClient.ExecuteQuery(query, token))
                 {
                     while(reader.Read()) 
                     {
@@ -382,7 +180,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
                 query =  ".show cluster extents | summarize sum(OriginalSize) by tostring(DatabaseName)";
             }
 
-            using (var reader = ExecuteQuery(query, token))
+            using (var reader = _kustoClient.ExecuteQuery(query, token))
             {
                 _databaseMetadata = reader.ToEnumerable()
                     .Where(row => !string.IsNullOrWhiteSpace(row["DatabaseName"].ToString()))
@@ -432,111 +230,11 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         }
 
         /// <inheritdoc/>
-        public override void UpdateDatabase(string databaseName){
+        public override void UpdateDatabase(string databaseName)
-            DatabaseName = databaseName;
-            SchemaState = Task.Run(() => KustoIntellisenseHelper.AddOrUpdateDatabaseAsync(this, GlobalState.Default, DatabaseName, ClusterName, throwOnError: false)).Result;
-        }
-
-        /// <inheritdoc/>
-        public override 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
-            
-            var codeBlock = script.GetBlockAtPosition(position);
-            var completion = codeBlock.Service.GetCompletionItems(position);
-            scriptDocumentInfo.ScriptParseInfo.CurrentSuggestions = completion.Items;         // this is declaration item so removed for now, but keep the info when api gets updated
-
-            List<LanguageServices.Contracts.CompletionItem> completions = new List<LanguageServices.Contracts.CompletionItem>();
-            foreach (var autoCompleteItem in completion.Items)
-            {
-                var label = autoCompleteItem.DisplayText;
-                var insertText = autoCompleteItem.Kind.ToString() == "Table" ? KustoQueryUtils.EscapeName(label) : label;
-                var completionKind = KustoIntellisenseHelper.CreateCompletionItemKind(autoCompleteItem.Kind);
-                completions.Add(AutoCompleteHelper.CreateCompletionItem(label, autoCompleteItem.Kind.ToString(),
-                    insertText, completionKind, scriptDocumentInfo.StartLine, scriptDocumentInfo.StartColumn,
-                    textPosition.Character));
-            }
-
-            return completions.ToArray();
-        }
-
-        /// <inheritdoc/>
-        public override Hover GetHoverHelp(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);
-
-            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 override DefinitionResult GetDefinition(string queryText, int index, int startLine, int startColumn, 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)
-            {
-
-            }
-
-            return null;
-        }
-
-        /// <inheritdoc/>
-        public override ScriptFileMarker[] GetSemanticMarkers(ScriptParseInfo parseInfo, ScriptFile scriptFile, string queryText)
         {
-            var kustoCodeService = new KustoCodeService(queryText, SchemaState);
+            _kustoClient.UpdateDatabase(databaseName);
-            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.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();
         }
-
+        
         /// <summary>
         /// Clears everything
         /// </summary>
@@ -728,7 +426,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
             query.Append($" | where TableName == '{tableName}' ");
             query.Append(" | project TableName, ColumnName, ColumnType, Folder");
 
-            using (var reader = ExecuteQuery(query.ToString(), token, databaseName))
+            using (var reader = _kustoClient.ExecuteQuery(query.ToString(), token, databaseName))
             {
                 var columns = reader.ToEnumerable()
                     .Select(row => new ColumnInfo
@@ -752,7 +450,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
 
             string query = $".show database {databaseName} cslschema";
 
-            using (var reader = ExecuteQuery(query, token, databaseName))
+            using (var reader = _kustoClient.ExecuteQuery(query, token, databaseName))
             {
                 return reader.ToEnumerable()
                     .Select(row => new TableInfo
@@ -1028,7 +726,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
 
             string query = ".show functions";
 
-            using (var reader = ExecuteQuery(query, token, databaseName))
+            using (var reader = _kustoClient.ExecuteQuery(query, token, databaseName))
             {
                 return reader.ToEnumerable()
                     .Select(row => new FunctionInfo
@@ -1050,7 +748,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
 
             string query = $".show function {functionName}";
 
-            using (var reader = ExecuteQuery(query, token, DatabaseName))
+            using (var reader = _kustoClient.ExecuteQuery(query, token, DatabaseName))
             {
                 return reader.ToEnumerable()
                     .Select(row => new FunctionInfo
@@ -1097,6 +795,11 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         {
             return string.IsNullOrWhiteSpace(objectName) ? databaseName : $"{databaseName}.{objectName}";
         }
+
+        public override void UpdateAzureToken(string azureToken)
+        {
+            _kustoClient.UpdateAzureToken(azureToken);
+        }
         #endregion
     }
 }

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

@@ -226,6 +226,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
         /// </summary>
         public void Close()
         {
+            _dataSource?.Dispose();
         }
         
         /// <summary>	

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

@@ -11,9 +11,6 @@ using Microsoft.SqlServer.Management.SqlParser.Binder;
 using Microsoft.SqlServer.Management.SqlParser.Common;
 using Microsoft.SqlServer.Management.SqlParser.MetadataProvider;
 using Microsoft.SqlServer.Management.SqlParser.Parser;
-using Kusto.Data.Net.Client;
-using Kusto.Data.Common;
-using Kusto.Data;
 using Microsoft.Kusto.ServiceLayer.DataSource;
 
 namespace Microsoft.Kusto.ServiceLayer.LanguageServices

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

@@ -771,7 +771,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
                 connInfo.TryGetConnection("Default", out connection);
                 IDataSource dataSource = connection.GetUnderlyingConnection();
                 
-                return dataSource.GetDefinition(scriptFile.Contents, textDocumentPosition.Position.Character, 1, 1);
+                return KustoIntellisenseHelper.GetDefinition(scriptFile.Contents, textDocumentPosition.Position.Character, 1, 1, dataSource.SchemaState);
             }
             else
             {
@@ -816,7 +816,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
                                 connInfo.TryGetConnection("Default", out connection);
                                 IDataSource dataSource = connection.GetUnderlyingConnection();               
                                 
-                                return dataSource.GetHoverHelp(scriptDocumentInfo, textDocumentPosition.Position);
+                                return KustoIntellisenseHelper.GetHoverHelp(scriptDocumentInfo, textDocumentPosition.Position, dataSource.SchemaState);
                             });
 
                         queueItem.ItemProcessed.WaitOne();
@@ -865,7 +865,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
                 connInfo.TryGetConnection("Default", out connection);
                 IDataSource dataSource = connection.GetUnderlyingConnection();
 			    
-                resultCompletionItems = dataSource.GetAutoCompleteSuggestions(scriptDocumentInfo, textDocumentPosition.Position);
+                resultCompletionItems = KustoIntellisenseHelper.GetAutoCompleteSuggestions(scriptDocumentInfo, textDocumentPosition.Position, dataSource.SchemaState);
             }
             else{
                 resultCompletionItems = DataSourceFactory.GetDefaultAutoComplete(DataSourceType.Kusto, scriptDocumentInfo, textDocumentPosition.Position);
@@ -1010,7 +1010,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
                     connInfo.TryGetConnection("Default", out var connection);
                     IDataSource dataSource = connection.GetUnderlyingConnection();
                     
-                    semanticMarkers = dataSource.GetSemanticMarkers(parseInfo, scriptFile, scriptFile.Contents);
+                    semanticMarkers = KustoIntellisenseHelper.GetSemanticMarkers(parseInfo, scriptFile, scriptFile.Contents, dataSource.SchemaState);
 			    }
                 else{
                     semanticMarkers = DataSourceFactory.GetDefaultSemanticMarkers(DataSourceType.Kusto, parseInfo, scriptFile, scriptFile.Contents);

src/Microsoft.Kusto.ServiceLayer/QueryExecution/Batch.cs

@@ -6,7 +6,6 @@ using System;
 using System.Collections.Generic;
 using System.Data;
 using System.Data.Common;
-using System.Diagnostics;
 using System.Data.SqlClient;
 using System.Linq;
 using System.Threading;
@@ -16,7 +15,7 @@ using Microsoft.Kusto.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.Kusto.ServiceLayer.QueryExecution.DataStorage;
 using Microsoft.SqlTools.Utility;
 using System.Globalization;
-using System.Collections.ObjectModel;
+using Microsoft.Kusto.ServiceLayer.DataSource.Exceptions;
 
 namespace Microsoft.Kusto.ServiceLayer.QueryExecution
 {
@@ -67,6 +66,8 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
         /// </summary>
         private readonly bool getFullColumnSchema;
 
+        private int _retryCount;
+
         #endregion
 
         internal Batch(string batchText, SelectionData selection, int ordinalId,
@@ -87,7 +88,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
             this.outputFileFactory = outputFileFactory;
             specialAction = new SpecialAction();
             BatchExecutionCount = executionCount > 0 ? executionCount : 1;
-
+            _retryCount = 1;
             this.getFullColumnSchema = getFullColumnSchema;
         }
 
@@ -251,7 +252,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
         public async Task Execute(ReliableDataSourceConnection conn, CancellationToken cancellationToken)
         {
             // Sanity check to make sure we haven't already run this batch
-            if (HasExecuted)
+            if (HasExecuted && _retryCount < 0)
             {
                 throw new InvalidOperationException("Batch has already executed.");
             }
@@ -266,6 +267,12 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
             {
                 await DoExecute(conn, cancellationToken);
             }
+            catch (DataSourceUnauthorizedException)
+            {
+                // Rerun the query once if unauthorized
+                _retryCount--;
+                throw;
+            }
             catch (TaskCanceledException)
             {
                 // Cancellation isn't considered an error condition
@@ -290,7 +297,6 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
                     await BatchCompletion(this);
                 }
             }
-
         }
 
         private async Task DoExecute(ReliableDataSourceConnection conn, CancellationToken cancellationToken)
@@ -308,6 +314,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
                 {
                     await ExecuteOnce(conn, cancellationToken);
                 }
+                
                 catch (DbException dbe)
                 {
                     HasError = true;

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

@@ -4,8 +4,6 @@
 //
 
 using System;
-using System.Data.Common;
-using System.Data.SqlClient;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -18,8 +16,8 @@ using Microsoft.SqlTools.Utility;
 using Microsoft.SqlTools.ServiceLayer.BatchParser.ExecutionEngineCode;
 using System.Collections.Generic;
 using System.Diagnostics;
+using Microsoft.Kusto.ServiceLayer.DataSource.Exceptions;
 using Microsoft.Kusto.ServiceLayer.Utility;
-using System.Text;
 
 namespace Microsoft.Kusto.ServiceLayer.QueryExecution
 {
@@ -333,7 +331,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
         /// </summary>
         private async Task ExecuteInternal()
         {
-            ReliableDataSourceConnection sqlConn = null;
+            ReliableDataSourceConnection queryConnection = null;
             try
             {
                 // check for cancellation token before actually making connection
@@ -341,7 +339,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
 
                 // Mark that we've internally executed
                 hasExecuteBeenCalled = true;
-                
+
                 // Don't actually execute if there aren't any batches to execute
                 if (Batches.Length == 0)
                 {
@@ -349,16 +347,19 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
                     {
                         await BatchMessageSent(new ResultMessage(SR.QueryServiceCompletedSuccessfully, false, null));
                     }
+
                     if (QueryCompleted != null)
                     {
                         await QueryCompleted(this);
                     }
+
                     return;
                 }
-                
+
                 // Locate and setup the connection
-                ReliableDataSourceConnection queryConnection = await ConnectionService.Instance.GetOrOpenConnection(editorConnection.OwnerUri, ConnectionType.Query);
+                queryConnection = await ConnectionService.Instance.GetOrOpenConnection(editorConnection.OwnerUri,
-                
+                        ConnectionType.Query);
+
                 // Execute beforeBatches synchronously, before the user defined batches 
                 foreach (Batch b in BeforeBatches)
                 {
@@ -390,6 +391,11 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
                     await QueryCompleted(this);
                 }
             }
+            catch (DataSourceUnauthorizedException)
+            {
+                ConnectionService.Instance.RefreshAzureToken(editorConnection.OwnerUri);
+                await ExecuteInternal();
+            }
             catch (Exception e)
             {
                 HasErrored = true;
@@ -409,7 +415,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
                 {
                     if (b.HasError)
                     {
-                        ConnectionService.EnsureConnectionIsOpen(sqlConn);
+                        ConnectionService.EnsureConnectionIsOpen(queryConnection);
                         break;
                     }
                 }

src/Microsoft.Kusto.ServiceLayer/QueryExecution/ResultSet.cs

@@ -74,7 +74,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
         /// <summary>
         /// Row count to use in special scenarios where we want to override the number of rows.
         /// </summary>
-        private long? rowCountOverride=null;
+        private long? rowCountOverride = null;
 
         /// <summary>
         /// The special action which applied to this result set