diff --git a/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionInfo.cs b/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionInfo.cs index 576c7b27..78fe4154 100644 --- a/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionInfo.cs +++ b/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; + } } } diff --git a/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs b/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs index 9a63e974..c66481bd 100644 --- a/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs +++ b/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs @@ -63,9 +63,8 @@ namespace Microsoft.Kusto.ServiceLayer.Connection /// /// Map from script URIs to ConnectionInfo objects - /// This is internal for testing access only /// - internal Dictionary OwnerToConnectionMap { get; } = new Dictionary(); + private Dictionary OwnerToConnectionMap { get; } = new Dictionary(); /// /// 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 diff --git a/src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/ConnectionDetails.cs b/src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/ConnectionDetails.cs index b722f12f..893e9348 100644 --- a/src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/ConnectionDetails.cs +++ b/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 diff --git a/src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/RequestSecurityTokenParams.cs b/src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/RequestSecurityTokenParams.cs new file mode 100644 index 00000000..6aae20f0 --- /dev/null +++ b/src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/RequestSecurityTokenParams.cs @@ -0,0 +1,22 @@ +namespace Microsoft.Kusto.ServiceLayer.Connection.Contracts +{ + public class RequestSecurityTokenParams + { + /// + /// Gets or sets the address of the authority to issue token. + /// + public string Authority { get; set; } + + /// + /// Gets or sets the provider that indicates the type of linked account to query. + /// + public string Provider { get; set; } + + /// + /// Gets or sets the identifier of the target resource that is the recipient of the requested token. + /// + public string Resource { get; set; } + + public string AccountId { get; set; } + } +} \ No newline at end of file diff --git a/src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/RequestSecurityTokenResponse.cs b/src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/RequestSecurityTokenResponse.cs new file mode 100644 index 00000000..c332beea --- /dev/null +++ b/src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/RequestSecurityTokenResponse.cs @@ -0,0 +1,15 @@ +namespace Microsoft.Kusto.ServiceLayer.Connection.Contracts +{ + public class RequestSecurityTokenResponse + { + /// + /// Gets or sets the key that uniquely identifies a particular linked account. + /// + public string AccountKey { get; set; } + + /// + /// Gets or sets the access token. + /// + public string Token { get; set; } + } +} \ No newline at end of file diff --git a/src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/SecurityTokenRequest.cs b/src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/SecurityTokenRequest.cs new file mode 100644 index 00000000..a7c99316 --- /dev/null +++ b/src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/SecurityTokenRequest.cs @@ -0,0 +1,15 @@ +using Microsoft.SqlTools.Hosting.Protocol.Contracts; + +namespace Microsoft.Kusto.ServiceLayer.Connection.Contracts +{ + /// + /// SecurityToken Request mapping entry + /// + public class SecurityTokenRequest + { + public static readonly + RequestType Type = + RequestType.Create( + "account/securityTokenRequest"); + } +} \ No newline at end of file diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceBase.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceBase.cs index 730f952a..6bd35bf7 100644 --- a/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceBase.cs +++ b/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 { /// public abstract class DataSourceBase : IDataSource { - protected Object dataSourceLock = new Object(); - - private string _database; - #region IDisposable /// @@ -87,17 +79,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource /// public abstract void UpdateDatabase(string databaseName); - /// - public abstract CompletionItem[] GetAutoCompleteSuggestions(ScriptDocumentInfo queryText, Position index, bool throwOnError = false); - /// - public abstract Hover GetHoverHelp(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, bool throwOnError = false); - - /// - public abstract DefinitionResult GetDefinition(string queryText, int index, int startLine, int startColumn, bool throwOnError = false); - - /// - public abstract ScriptFileMarker[] GetSemanticMarkers(ScriptParseInfo parseInfo, ScriptFile scriptFile, string queryText); - /// public abstract Task Exists(); @@ -108,27 +89,16 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource public abstract string GenerateExecuteFunctionScript(string functionName); + public abstract void UpdateAzureToken(string azureToken); + /// public DataSourceType DataSourceType { get; protected set; } - + /// - public string ClusterName { get; protected set; } + public abstract string ClusterName { get; } - /// - public string DatabaseName { - get - { - return _database; - } - - set - { - lock(dataSourceLock) - { - _database = value; - } - } - } + public abstract string DatabaseName { get; } + public abstract GlobalState SchemaState { get; } #endregion } diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceFactory.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceFactory.cs index 594db8c4..6ca18479 100644 --- a/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceFactory.cs +++ b/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: diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/KustoIntellisenseHelper.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/KustoIntellisenseHelper.cs index 75677e1f..42929fd2 100644 --- a/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/KustoIntellisenseHelper.cs +++ b/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 /// 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; - } - /// /// Convert CLR type name into a Kusto scalar type. /// @@ -161,20 +125,20 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense /// /// Loads the schema for the specified databasea into a a . /// - private static async Task LoadDatabaseAsync(IDataSource dataSource, string databaseName, bool throwOnError = false) + private static DatabaseSymbol LoadDatabaseAsync(IEnumerable tableSchemas, + IEnumerable functionSchemas, + string databaseName) { - var members = new List(); - CancellationTokenSource source = new CancellationTokenSource(); - CancellationToken cancellationToken = source.Token; - - var tableSchemas = await dataSource.ExecuteControlCommandAsync($".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(); 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(".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 databaseSymbol; + return new DatabaseSymbol(databaseName, members); } public static CompletionItemKind CreateCompletionItemKind(CompletionKind kustoKind) @@ -234,33 +198,41 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense /// /// Gets default keyword when user if not connected to any Kusto cluster. /// - 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); + 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 completions = new List(); - 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)); - } + List completions = + new List(); + 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(); + return completions.ToArray(); } /// /// Gets default diagnostics when user if not connected to any Kusto cluster. /// - 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 markers = new List(); if (parseResult != null && parseResult.Count() > 0) @@ -295,22 +267,27 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense /// /// Loads the schema for the specified database and returns a new with the database added or updated. /// - public static async Task AddOrUpdateDatabaseAsync(IDataSource dataSource, GlobalState globals, string databaseName, string clusterName, bool throwOnError) - { // try and show error from here. + public static GlobalState AddOrUpdateDatabase(IEnumerable tableSchemas, + IEnumerable functionSchemas, GlobalState globals, + string databaseName, string clusterName) + { + // try and show error from here. DatabaseSymbol databaseSymbol = null; - if(databaseName != null){ - databaseSymbol = await LoadDatabaseAsync(dataSource, databaseName, throwOnError).ConfigureAwait(false); + if (databaseName != null) + { + 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; } + /// + 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(); + 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(); + } + + /// + 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); + } + + /// + 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; + } + + /// + 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(); + + 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 markers = new List(); + 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(); + } } } diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/ShowDatabaseSchemaResult.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/ShowDatabaseSchemaResult.cs new file mode 100644 index 00000000..d194d72d --- /dev/null +++ b/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; + } +} \ No newline at end of file diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/ShowDatabasesResult.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/ShowDatabasesResult.cs new file mode 100644 index 00000000..0cbb930a --- /dev/null +++ b/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; + } +} \ No newline at end of file diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/ShowFunctionsResult.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/ShowFunctionsResult.cs new file mode 100644 index 00000000..812da976 --- /dev/null +++ b/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; + } +} \ No newline at end of file diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/Exceptions/DataSourceUnauthorizedException.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/Exceptions/DataSourceUnauthorizedException.cs new file mode 100644 index 00000000..51ef3fa2 --- /dev/null +++ b/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) + { + } + } +} \ No newline at end of file diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/IDataSource.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/IDataSource.cs index 7f8988f4..1979f430 100644 --- a/src/Microsoft.Kusto.ServiceLayer/DataSource/IDataSource.cs +++ b/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 /// /// The current database name, if there is one. /// - string DatabaseName { get; set; } + string DatabaseName { get; } + + GlobalState SchemaState { get; } /// /// Executes a query. @@ -88,29 +86,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource /// Object metadata. void UpdateDatabase(string databaseName); - /// - /// Gets autocomplete suggestions at given position. - /// - /// Object metadata. - CompletionItem[] GetAutoCompleteSuggestions(ScriptDocumentInfo queryText, Position index, bool throwOnError = false); - /// - /// Gets quick info hover tooltips for the current position. - /// - /// Object metadata. - Hover GetHoverHelp(ScriptDocumentInfo scriptDocumentInfo, Position textPosition, bool throwOnError = false); - - /// - /// Gets definition for a selected query text. - /// - /// Object metadata. - DefinitionResult GetDefinition(string queryText, int index, int startLine, int startColumn, bool throwOnError = false); - - /// - /// Gets a list of semantic diagnostic marks for the provided script file - /// - /// Object metadata. - ScriptFileMarker[] GetSemanticMarkers(ScriptParseInfo parseInfo, ScriptFile scriptFile, string queryText); - /// /// Tells whether the data source exists. /// @@ -136,5 +111,11 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource /// /// string GenerateExecuteFunctionScript(string functionName); + + /// + /// Updates Azure Token + /// + /// + void UpdateAzureToken(string azureToken); } } \ No newline at end of file diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/IKustoClient.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/IKustoClient.cs new file mode 100644 index 00000000..861f6cf2 --- /dev/null +++ b/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 + { + /// + /// SchemaState used for getting intellisense info. + /// + GlobalState SchemaState { get; } + + string ClusterName { get; } + + string DatabaseName { get; } + + void UpdateAzureToken(string azureAccountToken); + + IDataReader ExecuteQuery(string query, CancellationToken cancellationToken, string databaseName = null); + + /// + /// Executes a query or command against a kusto cluster and returns a sequence of result row instances. + /// + Task> ExecuteControlCommandAsync(string command, bool throwOnError, CancellationToken cancellationToken); + + /// + /// Executes a query. + /// + /// The query. + /// The results. + Task ExecuteQueryAsync(string query, CancellationToken cancellationToken, string databaseName = null); + + /// + /// Executes a Kusto control command. + /// + /// The command. + void ExecuteControlCommand(string command); + + void UpdateDatabase(string databaseName); + + void Dispose(); + } +} \ No newline at end of file diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoClient.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoClient.cs new file mode 100644 index 00000000..fc9539a8 --- /dev/null +++ b/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; + + /// + /// SchemaState used for getting intellisense info. + /// + 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 tableSchemas = Enumerable.Empty(); + IEnumerable functionSchemas = Enumerable.Empty(); + + Parallel.Invoke(() => + { + tableSchemas = ExecuteControlCommandAsync( + $".show database {DatabaseName} schema", + false, source.Token).Result; + }, () => + { + functionSchemas = ExecuteControlCommandAsync(".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); + } + + /// + /// 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" + /// + /// A connection string coming over the Data management protocol + 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( + !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; + } + + /// + /// Extracts the database name from the connectionString if it exists + /// otherwise it takes the first database name from the server + /// + /// + /// Database Name + 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); + } + } + + /// + /// Executes a query or command against a kusto cluster and returns a sequence of result row instances. + /// + public async Task> ExecuteControlCommandAsync(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(tableReader); + } + catch (DataSourceUnauthorizedException) + { + throw; + } + catch (Exception) when (!throwOnError) + { + return null; + } + } + + /// + /// Executes a query. + /// + /// The query. + /// The results. + public Task 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); + } + + /// + /// Executes a Kusto control command. + /// + /// The command. + 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(); + } + } +} \ No newline at end of file diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoDataSource.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoDataSource.cs index d16d1fda..a958121a 100644 --- a/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoDataSource.cs +++ b/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 /// public class KustoDataSource : DataSourceBase { - private ICslQueryProvider _kustoQueryProvider; - - private ICslAdminProvider _kustoAdminProvider; + private IKustoClient _kustoClient; /// /// List of databases. @@ -66,6 +53,12 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource /// private ConcurrentDictionary> _functionMetadata = new ConcurrentDictionary>(); + 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 /// /// Prevents a default instance of the class from being created. /// - public KustoDataSource(string connectionString, string azureAccountToken) + public KustoDataSource(IKustoClient kustoClient) { - ClusterName = GetClusterName(connectionString); - UserToken = azureAccountToken; - DatabaseName = GetDatabaseName(connectionString); - SchemaState = Task.Run(() => - KustoIntellisenseHelper.AddOrUpdateDatabaseAsync(this, GlobalState.Default, DatabaseName, ClusterName, - throwOnError: false)).Result; + _kustoClient = kustoClient; // Check if a connection can be made ValidationUtils.IsTrue(Exists().Result, $"Unable to connect. ClusterName = {ClusterName}, DatabaseName = {DatabaseName}"); } - - /// - /// 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" - /// - /// A connection string coming over the Data management protocol - 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"); - } - - /// - /// Extracts the database name from the connectionString if it exists - /// otherwise it takes the first database name from the server - /// - /// - /// Database Name - 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; - } - } - - /// - /// SchemaState used for getting intellisense info. - /// - public GlobalState SchemaState { get; private set; } - - /// - /// The AAD user token. - /// - public string UserToken { get; private set; } - - /// - /// The AAD application client id. - /// - public string ApplicationClientId { get; private set; } - - /// - /// The AAD application client key. - /// - 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; - } - } - + /// /// Disposes resources. /// @@ -209,11 +94,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource // Dispose managed resources. if (disposing) { - _kustoQueryProvider?.Dispose(); - _kustoQueryProvider = null; - - _kustoAdminProvider?.Dispose(); - _kustoAdminProvider = null; + _kustoClient.Dispose(); } base.Dispose(disposing); @@ -228,42 +109,10 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource /// The results. public override Task 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); - } - /// public override async Task Exists() { @@ -282,68 +131,17 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource } #endregion - - /// - /// Executes a Kusto control command. - /// - /// The command. - 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( - !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 } /// - public override void UpdateDatabase(string databaseName){ - DatabaseName = databaseName; - SchemaState = Task.Run(() => KustoIntellisenseHelper.AddOrUpdateDatabaseAsync(this, GlobalState.Default, DatabaseName, ClusterName, throwOnError: false)).Result; - } - - /// - 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 completions = new List(); - 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(); - } - - /// - 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); - - } - - /// - 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; - } - - /// - public override ScriptFileMarker[] GetSemanticMarkers(ScriptParseInfo parseInfo, ScriptFile scriptFile, string queryText) + public override void UpdateDatabase(string databaseName) { - var kustoCodeService = new KustoCodeService(queryText, SchemaState); - var script = CodeScript.From(queryText, SchemaState); - var parseResult = new List(); - - 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 markers = new List(); - 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(); + _kustoClient.UpdateDatabase(databaseName); } - + /// /// Clears everything /// @@ -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 } } diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/ReliableDataSourceConnection.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/ReliableDataSourceConnection.cs index 3e80d718..b0c30c0c 100644 --- a/src/Microsoft.Kusto.ServiceLayer/DataSource/ReliableDataSourceConnection.cs +++ b/src/Microsoft.Kusto.ServiceLayer/DataSource/ReliableDataSourceConnection.cs @@ -226,6 +226,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection /// public void Close() { + _dataSource?.Dispose(); } /// diff --git a/src/Microsoft.Kusto.ServiceLayer/LanguageServices/ConnectedBindingContext.cs b/src/Microsoft.Kusto.ServiceLayer/LanguageServices/ConnectedBindingContext.cs index bde3a381..7491a08f 100644 --- a/src/Microsoft.Kusto.ServiceLayer/LanguageServices/ConnectedBindingContext.cs +++ b/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 diff --git a/src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs b/src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs index 7fdcdfbf..4c6ac1f4 100644 --- a/src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs +++ b/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); diff --git a/src/Microsoft.Kusto.ServiceLayer/QueryExecution/Batch.cs b/src/Microsoft.Kusto.ServiceLayer/QueryExecution/Batch.cs index 60fe2d19..f3557cf8 100644 --- a/src/Microsoft.Kusto.ServiceLayer/QueryExecution/Batch.cs +++ b/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 /// 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; diff --git a/src/Microsoft.Kusto.ServiceLayer/QueryExecution/Query.cs b/src/Microsoft.Kusto.ServiceLayer/QueryExecution/Query.cs index 75708c23..aaa2eadc 100644 --- a/src/Microsoft.Kusto.ServiceLayer/QueryExecution/Query.cs +++ b/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 /// 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; } } diff --git a/src/Microsoft.Kusto.ServiceLayer/QueryExecution/ResultSet.cs b/src/Microsoft.Kusto.ServiceLayer/QueryExecution/ResultSet.cs index bc3b0422..5af04b95 100644 --- a/src/Microsoft.Kusto.ServiceLayer/QueryExecution/ResultSet.cs +++ b/src/Microsoft.Kusto.ServiceLayer/QueryExecution/ResultSet.cs @@ -74,7 +74,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution /// /// Row count to use in special scenarios where we want to override the number of rows. /// - private long? rowCountOverride=null; + private long? rowCountOverride = null; /// /// The special action which applied to this result set