ckaczor/sqltoolsservice
1
0
Fork 0
mirror of https://github.com/ckaczor/sqltoolsservice.git synced 2026-01-14 01:25:40 -05:00
Code Issues Packages Projects Releases Wiki Activity

3644 Kusto Azure Token Refresh (#1081)

Browse Source 
* 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
This commit is contained in:
Justin M
2020-09-28 13:24:23 -07:00
committed by GitHub
parent d6ff73d510
commit d061e781f4
23 changed files with 673 additions and 487 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionInfo.cs
View File

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

22
src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs
View File

@@ -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

1
src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/ConnectionDetails.cs
View File

@@ -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

22
src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/RequestSecurityTokenParams.cs Normal file
View File

@@ -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; }
}
}

15
src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/RequestSecurityTokenResponse.cs Normal file
View File

@@ -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; }
}
}

15
src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/SecurityTokenRequest.cs Normal file
View File

@@ -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");
}
}

44
src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceBase.cs
View File

@@ -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 string DatabaseName {
get
{
return _database;
}
set
{
lock(dataSourceLock)
{
_database = value;
}
}
}
public abstract string DatabaseName { get; }
public abstract GlobalState SchemaState { get; }
#endregion
}

3
src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceFactory.cs
View File

@@ -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:

224
src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/KustoIntellisenseHelper.cs
View File

@@ -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 databaseSymbol;
return new DatabaseSymbol(databaseName, members);
}
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){
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<LanguageServices.Contracts.CompletionItem> completions = new List<LanguageServices.Contracts.CompletionItem>();
foreach (var autoCompleteItem in completion.Items)
{
var label = autoCompleteItem.DisplayText;
// convert the completion item candidates into vscode format CompletionItems
completions.Add(AutoCompleteHelper.CreateCompletionItem(label, label + " keyword", label, CompletionItemKind.Keyword, scriptDocumentInfo.StartLine, scriptDocumentInfo.StartColumn, textDocumentPosition.Character));
}
List<LanguageServices.Contracts.CompletionItem> completions =
new List<LanguageServices.Contracts.CompletionItem>();
foreach (var autoCompleteItem in completion.Items)
{
var label = autoCompleteItem.DisplayText;
// convert the completion item candidates into vscode format CompletionItems
completions.Add(AutoCompleteHelper.CreateCompletionItem(label, label + " keyword", label,
CompletionItemKind.Keyword, scriptDocumentInfo.StartLine, scriptDocumentInfo.StartColumn,
textDocumentPosition.Character));
}
return completions.ToArray();
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)
{ // try and show error from here.
public static GlobalState AddOrUpdateDatabase(IEnumerable<ShowDatabaseSchemaResult> tableSchemas,
IEnumerable<ShowFunctionsResult> functionSchemas, GlobalState globals,
string databaseName, string clusterName)
{
// try and show error from here.
DatabaseSymbol databaseSymbol = null;
if(databaseName != null){
databaseSymbol = 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;
}
/// <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();
}
}
}

16
src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/ShowDatabaseSchemaResult.cs Normal file
View File

@@ -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;
}
}

14
src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/ShowDatabasesResult.cs Normal file
View File

@@ -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;
}
}

11
src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/ShowFunctionsResult.cs Normal file
View File

@@ -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;
}
}

11
src/Microsoft.Kusto.ServiceLayer/DataSource/Exceptions/DataSourceUnauthorizedException.cs Normal file
View File

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

39
src/Microsoft.Kusto.ServiceLayer/DataSource/IDataSource.cs
View File

@@ -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);
}
}

46
src/Microsoft.Kusto.ServiceLayer/DataSource/IKustoClient.cs Normal file
View File

@@ -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();
}
}

260
src/Microsoft.Kusto.ServiceLayer/DataSource/KustoClient.cs Normal file
View File

@@ -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();
}
}
}

355
src/Microsoft.Kusto.ServiceLayer/DataSource/KustoDataSource.cs
View File

@@ -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 ICslAdminProvider _kustoAdminProvider;
private IKustoClient _kustoClient;
/// <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);
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<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();
_kustoQueryProvider = null;
_kustoAdminProvider?.Dispose();
_kustoAdminProvider = null;
_kustoClient.Dispose();
}
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){
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)
public override void UpdateDatabase(string databaseName)
{
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.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);
}
/// <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
}
}

1
src/Microsoft.Kusto.ServiceLayer/DataSource/ReliableDataSourceConnection.cs
View File

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

3
src/Microsoft.Kusto.ServiceLayer/LanguageServices/ConnectedBindingContext.cs
View File

@@ -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

8
src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs
View File

@@ -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);

17
src/Microsoft.Kusto.ServiceLayer/QueryExecution/Batch.cs
View File

@@ -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;

24
src/Microsoft.Kusto.ServiceLayer/QueryExecution/Query.cs
View File

@@ -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;
}
}

2
src/Microsoft.Kusto.ServiceLayer/QueryExecution/ResultSet.cs
View File

@@ -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