Add Notebook <-> SQL convert (#11590)

* Add Notebook <-> SQL convert

* Update STS

extensions/mssql/src/constants.ts

@@ -31,6 +31,8 @@ export const hdfsRootPath = '/';
 export const clusterEndpointsProperty = 'clusterEndpoints';
 export const isBigDataClusterProperty = 'isBigDataCluster';
 
+export const ViewType = 'view';
+
 // SERVICE NAMES //////////////////////////////////////////////////////////
 export const ObjectExplorerService = 'objectexplorer';
 export const CmsService = 'cmsService';
@@ -38,8 +40,8 @@ export const DacFxService = 'dacfxService';
 export const SchemaCompareService = 'schemaCompareService';
 export const LanguageExtensionService = 'languageExtensionService';
 export const objectExplorerPrefix: string = 'objectexplorer://';
-export const ViewType = 'view';
 export const SqlAssessmentService = 'sqlAssessmentService';
+export const NotebookConvertService = 'notebookConvertService';
 
 export enum BuiltInCommands {
 	SetContext = 'setContext'

extensions/mssql/src/contracts.ts

@@ -747,3 +747,33 @@ export class CompletionExtensionParams {
 export namespace CompletionExtLoadRequest {
 	export const type = new RequestType<CompletionExtensionParams, boolean, void, void>('completion/extLoad');
 }
+
+// ------------------------------- < Load Completion Extension Request > ------------------------------------
+
+/// ------------------------------- <Convert Notebook> -----------------------------
+
+export interface ConvertNotebookToSqlParams {
+	content: string;
+}
+
+export namespace ConvertNotebookToSqlRequest {
+	export const type = new RequestType<ConvertNotebookToSqlParams, ConvertNotebookToSqlResult, void, void>('notebookconvert/convertnotebooktosql');
+}
+
+export interface ConvertNotebookToSqlResult extends azdata.ResultStatus {
+	content: string;
+}
+
+export interface ConvertSqlToNotebookParams {
+	clientUri: string;
+}
+
+export namespace ConvertSqlToNotebookRequest {
+	export const type = new RequestType<ConvertSqlToNotebookParams, ConvertSqlToNotebookResult, void, void>('notebookconvert/convertsqltonotebook');
+}
+
+export interface ConvertSqlToNotebookResult extends azdata.ResultStatus {
+	content: string;
+}
+
+// ------------------------------- <Convert Notebook> -----------------------------

extensions/mssql/src/main.ts

@@ -30,6 +30,7 @@ import { SqlToolsServer } from './sqlToolsServer';
 import { promises as fs } from 'fs';
 import { IconPathHelper } from './iconHelper';
 import * as nls from 'vscode-nls';
+import { INotebookConvertService } from './notebookConvert/notebookConvertService';
 
 const localize = nls.loadMessageBundle();
 const msgSampleCodeDataFrame = localize('msgSampleCodeDataFrame', "This sample code loads the file into a data frame and shows the first 10 results.");
@@ -79,6 +80,24 @@ export async function activate(context: vscode.ExtensionContext): Promise<IExten
 	context.subscriptions.push(server);
 	await server.start(appContext);
 
+	vscode.commands.registerCommand('mssql.exportSqlAsNotebook', async (uri: vscode.Uri) => {
+		const result = await appContext.getService<INotebookConvertService>(Constants.NotebookConvertService).convertSqlToNotebook(uri.toString());
+		const title = findNextUntitledEditorName();
+		const untitledUri = vscode.Uri.parse(`untitled:${title}`);
+		await azdata.nb.showNotebookDocument(untitledUri, { initialContent: result.content });
+	});
+
+	vscode.commands.registerCommand('mssql.exportNotebookToSql', async (uri: vscode.Uri) => {
+		// SqlToolsService doesn't currently store anything about Notebook documents so we have to pass the raw JSON to it directly
+		// We use vscode.workspace.textDocuments here because the azdata.nb.notebookDocuments don't actually contain their contents
+		// (they're left out for perf purposes)
+		const doc = vscode.workspace.textDocuments.find(doc => doc.uri.toString() === uri.toString());
+		const result = await appContext.getService<INotebookConvertService>(Constants.NotebookConvertService).convertNotebookToSql(doc.getText());
+
+		const sqlDoc = await vscode.workspace.openTextDocument({ language: 'sql', content: result.content });
+		await vscode.commands.executeCommand('vscode.open', sqlDoc.uri);
+	});
+
 	return createMssqlApi(appContext);
 }
 

extensions/mssql/src/notebookConvert/notebookConvertService.ts

@@ -0,0 +1,58 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { AppContext } from '../appContext';
+import { SqlOpsDataClient, ISqlOpsFeature } from 'dataprotocol-client';
+import { ClientCapabilities } from 'vscode-languageclient';
+import * as constants from '../constants';
+import * as contracts from '../contracts';
+
+export interface INotebookConvertService {
+	convertNotebookToSql(content: string): Promise<contracts.ConvertNotebookToSqlResult | undefined>;
+	convertSqlToNotebook(content: string): Promise<contracts.ConvertSqlToNotebookResult | undefined>;
+}
+
+export class NotebookConvertService implements INotebookConvertService {
+	public static asFeature(context: AppContext): ISqlOpsFeature {
+		return class extends NotebookConvertService {
+			constructor(client: SqlOpsDataClient) {
+				super(context, client);
+			}
+
+			fillClientCapabilities(capabilities: ClientCapabilities): void {
+			}
+
+			initialize(): void {
+			}
+		};
+	}
+
+	private constructor(context: AppContext, protected readonly client: SqlOpsDataClient) {
+		context.registerService(constants.NotebookConvertService, this);
+	}
+
+	async convertNotebookToSql(content: string): Promise<contracts.ConvertNotebookToSqlResult | undefined> {
+		let params: contracts.ConvertNotebookToSqlParams = { content: content };
+		try {
+			return this.client.sendRequest(contracts.ConvertNotebookToSqlRequest.type, params);
+		}
+		catch (e) {
+			this.client.logFailedRequest(contracts.ConvertNotebookToSqlRequest.type, e);
+		}
+
+		return undefined;
+	}
+	async convertSqlToNotebook(content: string): Promise<contracts.ConvertSqlToNotebookResult | undefined> {
+		let params: contracts.ConvertSqlToNotebookParams = { clientUri: content };
+		try {
+			return this.client.sendRequest(contracts.ConvertSqlToNotebookRequest.type, params);
+		}
+		catch (e) {
+			this.client.logFailedRequest(contracts.ConvertSqlToNotebookRequest.type, e);
+		}
+
+		return undefined;
+	}
+}

extensions/mssql/src/sqlToolsServer.ts

@@ -23,6 +23,7 @@ import { promises as fs } from 'fs';
 import * as nls from 'vscode-nls';
 import { LanguageExtensionService } from './languageExtension/languageExtensionService';
 import { SqlAssessmentService } from './sqlAssessment/sqlAssessmentService';
+import { NotebookConvertService } from './notebookConvert/notebookConvertService';
 
 const localize = nls.loadMessageBundle();
 const outputChannel = vscode.window.createOutputChannel(Constants.serviceName);
@@ -160,7 +161,8 @@ function getClientOptions(context: AppContext): ClientOptions {
 			LanguageExtensionService.asFeature(context),
 			DacFxService.asFeature(context),
 			CmsService.asFeature(context),
-			SqlAssessmentService.asFeature(context)
+			SqlAssessmentService.asFeature(context),
+			NotebookConvertService.asFeature(context)
 		],
 		outputChannel: new CustomOutputChannel()
 	};