Kusto - ADS Core changes (#11750)

* ADS Core changes

* Generic fix to get language mode based on providerId

* Addressed comments

* Fix failing tests

* Fix for "New Notebook" from file menu option

* Remove line

* Fix for merge conflict

* Add removed comment back

* Addressed comments

Co-authored-by: Monica Gupta <mogupt@microsoft.com>

src/sql/workbench/browser/scriptingUtils.ts

@@ -45,7 +45,7 @@ export async function scriptSelect(connectionProfile: IConnectionProfile, metada
 	let paramDetails: azdata.ScriptingParamDetails = getScriptingParamDetails(connectionService, connectionResult, metadata);
 	const result = await scriptingService.script(connectionResult, metadata, ScriptOperation.Select, paramDetails);
 	if (result && result.script) {
-		const owner = await queryEditorService.newSqlEditor({ initalContent: result.script });
+		const owner = await queryEditorService.newSqlEditor({ initalContent: result.script }, connectionProfile?.providerName);
 		// Connect our editor to the input connection
 		let options: IConnectionCompletionOptions = {
 			params: { connectionType: ConnectionType.editor, runQueryOnCompletion: RunQueryOnConnectionMode.executeQuery, input: owner },

src/sql/workbench/contrib/query/browser/queryActions.ts

@@ -121,7 +121,7 @@ export function openNewQuery(accessor: ServicesAccessor, profile?: IConnectionPr
 	if (!profile) {
 		profile = getCurrentGlobalConnection(objectExplorerService, connectionManagementService, editorService);
 	}
-	return queryEditorService.newSqlEditor({ initalContent }).then((owner: IConnectableInput) => {
+	return queryEditorService.newSqlEditor({ initalContent }, profile?.providerName).then((owner: IConnectableInput) => {
 		// Connect our editor to the input connection
 		let options: IConnectionCompletionOptions = {
 			params: { connectionType: ConnectionType.editor, runQueryOnCompletion: onConnection, input: owner },

src/sql/workbench/contrib/query/browser/queryEditor.ts

@@ -29,7 +29,7 @@ import { BaseTextEditor } from 'vs/workbench/browser/parts/editor/textEditor';
 import { FileEditorInput } from 'vs/workbench/contrib/files/common/editors/fileEditorInput';
 import { URI } from 'vs/base/common/uri';
 import { IFileService, FileChangesEvent } from 'vs/platform/files/common/files';
-
+import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { QueryEditorInput, IQueryEditorStateChange } from 'sql/workbench/common/editor/query/queryEditorInput';
 import { QueryResultsEditor } from 'sql/workbench/contrib/query/browser/queryResultsEditor';
 import * as queryContext from 'sql/workbench/contrib/query/common/queryContext';
@@ -95,6 +95,7 @@ export class QueryEditor extends BaseEditor {
 		@IContextKeyService contextKeyService: IContextKeyService,
 		@IEditorGroupsService editorGroupService: IEditorGroupsService,
 		@IFileService fileService: IFileService,
+		@IConnectionManagementService private readonly connectionManagementService: IConnectionManagementService,
 		@IEditorService private readonly editorService: IEditorService,
 		@IInstantiationService private readonly instantiationService: IInstantiationService,
 		@IConfigurationService private readonly configurationService: IConfigurationService
@@ -185,14 +186,6 @@ export class QueryEditor extends BaseEditor {
 		this._actualQueryPlanAction = this.instantiationService.createInstance(actions.ActualQueryPlanAction, this);
 		this._toggleSqlcmdMode = this.instantiationService.createInstance(actions.ToggleSqlCmdModeAction, this, false);
 		this._exportAsNotebookAction = this.instantiationService.createInstance(actions.ExportAsNotebookAction, this);
-
-		this.setTaskbarContent();
-
-		this._register(this.configurationService.onDidChangeConfiguration(e => {
-			if (e.affectsConfiguration('workbench.enablePreviewFeatures')) {
-				this.setTaskbarContent();
-			}
-		}));
 	}
 
 	/**
@@ -204,6 +197,12 @@ export class QueryEditor extends BaseEditor {
 			this._changeConnectionAction.enabled = this.input.state.connected;
 			if (this.input.state.connected) {
 				this.listDatabasesActionItem.onConnected();
+				this.setTaskbarContent();
+				this._register(this.configurationService.onDidChangeConfiguration(e => {
+					if (e.affectsConfiguration('workbench.enablePreviewFeatures')) {
+						this.setTaskbarContent();
+					}
+				}));
 			} else {
 				this.listDatabasesActionItem.onDisconnect();
 			}
@@ -259,24 +258,10 @@ export class QueryEditor extends BaseEditor {
 		const separator = Taskbar.createTaskbarSeparator();
 		let content: ITaskbarContent[];
 		const previewFeaturesEnabled = this.configurationService.getValue('workbench')['enablePreviewFeatures'];
-		const notebookConvertActionsEnabled = this.configurationService.getValue('notebook')['showNotebookConvertActions'];
-		if (previewFeaturesEnabled) {
-			content = [
-				{ action: this._runQueryAction },
-				{ action: this._cancelQueryAction },
-				{ element: separator },
-				{ action: this._toggleConnectDatabaseAction },
-				{ action: this._changeConnectionAction },
-				{ action: this._listDatabasesAction },
-				{ element: separator },
-				{ action: this._estimatedQueryPlanAction }, // Preview
-				{ action: this._toggleSqlcmdMode }, // Preview
-			];
+		let connectionProfile = this.connectionManagementService.getConnectionProfile(this.input.uri);
 
-			if (notebookConvertActionsEnabled) {
-				content.push({ action: this._exportAsNotebookAction });
-			}
-		} else {
+		// TODO: Make it more generic, some way for extensions to register the commands it supports
+		if (connectionProfile?.providerName === 'KUSTO') {
 			content = [
 				{ action: this._runQueryAction },
 				{ action: this._cancelQueryAction },
@@ -285,11 +270,40 @@ export class QueryEditor extends BaseEditor {
 				{ action: this._changeConnectionAction },
 				{ action: this._listDatabasesAction }
 			];
-			const notebookConvertActionsEnabled = this.configurationService.getValue('notebook')['notebook.showNotebookConvertActions'];
-			if (notebookConvertActionsEnabled) {
-				content.push(
+		}
+		else {
+			const notebookConvertActionsEnabled = this.configurationService.getValue('notebook')['showNotebookConvertActions'];
+			if (previewFeaturesEnabled) {
+				content = [
+					{ action: this._runQueryAction },
+					{ action: this._cancelQueryAction },
 					{ element: separator },
-					{ action: this._exportAsNotebookAction });
+					{ action: this._toggleConnectDatabaseAction },
+					{ action: this._changeConnectionAction },
+					{ action: this._listDatabasesAction },
+					{ element: separator },
+					{ action: this._estimatedQueryPlanAction }, // Preview
+					{ action: this._toggleSqlcmdMode }, // Preview
+				];
+
+				if (notebookConvertActionsEnabled) {
+					content.push({ action: this._exportAsNotebookAction });
+				}
+			} else {
+				content = [
+					{ action: this._runQueryAction },
+					{ action: this._cancelQueryAction },
+					{ element: separator },
+					{ action: this._toggleConnectDatabaseAction },
+					{ action: this._changeConnectionAction },
+					{ action: this._listDatabasesAction }
+				];
+				const notebookConvertActionsEnabled = this.configurationService.getValue('notebook')['notebook.showNotebookConvertActions'];
+				if (notebookConvertActionsEnabled) {
+					content.push(
+						{ element: separator },
+						{ action: this._exportAsNotebookAction });
+				}
 			}
 		}
 

src/sql/workbench/contrib/query/browser/queryInputFactory.ts

@@ -28,7 +28,8 @@ const editorInputFactoryRegistry = Registry.as<IEditorInputFactoryRegistry>(Edit
 
 export class QueryEditorLanguageAssociation implements ILanguageAssociation {
 	static readonly isDefault = true;
-	static readonly languages = ['sql'];
+	static readonly languages = ['sql', 'kusto'];	//TODO Add language id here for new languages supported in query editor. Make it easier to contribute new extension's languageID
+
 
 	constructor(@IInstantiationService private readonly instantiationService: IInstantiationService,
 		@IObjectExplorerService private readonly objectExplorerService: IObjectExplorerService,

src/sql/workbench/services/connection/browser/connectionManagementService.ts

@@ -1431,6 +1431,13 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 		return connectionProvider && connectionProvider.properties;
 	}
 
+	/**
+	 * Gets languageMode property of provider if it exists. Defaults to 'sql'
+	*/
+	public getProviderLanguageMode(providerName?: string): string {
+		return this._providers.get(providerName)?.properties?.['languageMode'] || 'sql';
+	}
+
 	/**
 	 * Get known connection profiles including active connections, recent connections and saved connections.
 	 * @param activeConnectionsOnly Indicates whether only get the active connections, default value is false.

src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts

@@ -247,7 +247,7 @@ class SqlKernel extends Disposable implements nb.IKernel {
 			implementation: '',
 			implementation_version: '',
 			language_info: {
-				name: 'sql',
+				name: this._connectionManagementService.getProviderLanguageMode(this._currentConnection?.providerName),
 				version: '',
 			},
 			banner: '',

src/sql/workbench/services/queryEditor/browser/queryEditorService.ts

@@ -5,7 +5,7 @@
 
 import { QueryResultsInput } from 'sql/workbench/common/editor/query/queryResultsInput';
 import { EditDataInput } from 'sql/workbench/browser/editData/editDataInput';
-import { IConnectableInput } from 'sql/platform/connection/common/connectionManagement';
+import { IConnectableInput, IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { IQueryEditorService, INewSqlEditorOptions } from 'sql/workbench/services/queryEditor/common/queryEditorService';
 import { UntitledQueryEditorInput } from 'sql/workbench/common/editor/query/untitledQueryEditorInput';
 
@@ -36,22 +36,23 @@ export class QueryEditorService implements IQueryEditorService {
 		@IUntitledTextEditorService private _untitledEditorService: IUntitledTextEditorService,
 		@IInstantiationService private _instantiationService: IInstantiationService,
 		@IEditorService private _editorService: IEditorService,
-		@IConfigurationService private _configurationService: IConfigurationService
+		@IConfigurationService private _configurationService: IConfigurationService,
+		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService
 	) {
 	}
 
 	////// Public functions
 
 	/**
-	 * Creates new untitled document for SQL query and opens in new editor tab
+	 * Creates new untitled document for SQL/Kusto query and opens in new editor tab
 	 */
-	public async newSqlEditor(options: INewSqlEditorOptions = {}): Promise<IConnectableInput> {
+	public async newSqlEditor(options: INewSqlEditorOptions = {}, connectionProviderName?: string): Promise<IConnectableInput> {
 		options = mixin(options, defaults, false);
 		// Create file path and file URI
-		let docUri: URI = options.resource ?? URI.from({ scheme: Schemas.untitled, path: await this.createUntitledSqlFilePath() });
+		let docUri: URI = options.resource ?? URI.from({ scheme: Schemas.untitled, path: await this.createUntitledSqlFilePath(connectionProviderName) });
 
 		// Create a sql document pane with accoutrements
-		const fileInput = this._editorService.createEditorInput({ forceUntitled: true, resource: docUri, mode: 'sql' }) as UntitledTextEditorInput;
+		const fileInput = this._editorService.createEditorInput({ forceUntitled: true, resource: docUri, mode: this._connectionManagementService.getProviderLanguageMode(connectionProviderName) }) as UntitledTextEditorInput;
 		let untitledEditorModel = await fileInput.resolve() as UntitledTextEditorModel;
 		if (options.initalContent) {
 			untitledEditorModel.textEditorModel.setValue(options.initalContent);
@@ -99,8 +100,11 @@ export class QueryEditorService implements IQueryEditorService {
 	}
 
 	////// Private functions
+	private createUntitledSqlFilePath(providerName?: string): Promise<string> {
+		if (providerName === 'KUSTO') {
+			return this.createPrefixedSqlFilePath(providerName + 'Query');
+		}
 
-	private createUntitledSqlFilePath(): Promise<string> {
 		return this.createPrefixedSqlFilePath('SQLQuery');
 	}
 

src/sql/workbench/services/queryEditor/common/queryEditorService.ts

@@ -38,8 +38,8 @@ export interface IQueryEditorService {
 
 	_serviceBrand: undefined;
 
-	// Creates new untitled document for SQL queries and opens it in a new editor tab
-	newSqlEditor(options?: INewSqlEditorOptions): Promise<IConnectableInput>;
+	// Creates new untitled document for SQL/KUSTO queries and opens it in a new editor tab
+	newSqlEditor(options?: INewSqlEditorOptions, connectionProviderName?: string): Promise<IConnectableInput>;
 
 	// Creates new edit data session
 	newEditDataEditor(schemaName: string, tableName: string, queryString: string): Promise<IConnectableInput>;