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

Initial implementation for VSCode Notebook support (#17885)

Browse Source
This commit is contained in:
Cory Rivera
2022-01-03 15:59:37 -08:00
committed by GitHub
parent af5575a852
commit 2ecc3d35ca
45 changed files with 1533 additions and 260 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/sql/workbench/services/notebook/browser/models/cell.ts
View File

@@ -611,7 +611,8 @@ export class CellModel extends Disposable implements ICellModel {
if (tryMatchCellMagic(this.source[0]) !== ads_execute_command || !this._isCommandExecutionSettingEnabled) {
const future = kernel.requestExecute({
code: content,
stop_on_error: true
stop_on_error: true,
notebookUri: this.notebookModel.notebookUri
}, false);
this.setFuture(future as FutureInternal);
this.fireExecutionStateChanged();

3
src/sql/workbench/services/notebook/browser/models/clientSession.ts
View File

@@ -77,6 +77,9 @@ export class ClientSession implements IClientSession {
}
private async startServer(kernelSpec: nb.IKernelSpec): Promise<void> {
if (!this._executeManager) {
throw new Error(localize('NoExecuteManager', "Server could not start because a provider was not defined for this notebook file type."));
}
let serverManager = this._executeManager.serverManager;
if (serverManager) {
await serverManager.startServer(kernelSpec);

21
src/sql/workbench/services/notebook/browser/models/notebookModel.ts
View File

@@ -179,13 +179,6 @@ export class NotebookModel extends Disposable implements INotebookModel {
return manager;
}
public getExecuteManager(providerId: string): IExecuteManager | undefined {
if (providerId) {
return this.executeManagers.find(manager => manager.providerId === providerId);
}
return undefined;
}
public get notebookOptions(): INotebookModelOptions {
return this._notebookOptions;
}
@@ -519,7 +512,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
public async requestModelLoad(): Promise<void> {
try {
this.setDefaultKernelAndProviderId();
await this.setDefaultKernelAndProviderId();
this.trySetLanguageFromLangInfo();
} catch (error) {
this._inErrorState = true;
@@ -975,7 +968,15 @@ export class NotebookModel extends Disposable implements INotebookModel {
this._activeClientSession = clientSession;
}
public setDefaultKernelAndProviderId() {
public async setDefaultKernelAndProviderId(): Promise<void> {
if (!this._defaultKernel) {
await this.executeManager.sessionManager.ready;
if (this.executeManager.sessionManager.specs) {
let defaultKernelName = this.executeManager.sessionManager.specs.defaultKernel;
this._defaultKernel = this.executeManager.sessionManager.specs.kernels.find(kernel => kernel.name === defaultKernelName);
}
}
if (this._capabilitiesService?.providers) {
let providers = this._capabilitiesService.providers;
for (const server in providers) {
@@ -1416,7 +1417,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
this._onProviderIdChanged.fire(this._providerId);
await this.shutdownActiveSession();
let manager = this.getExecuteManager(providerId);
let manager = this.executeManager;
if (manager) {
await this.startSession(manager, displayName, false, kernelAlias);
} else {

6
src/sql/workbench/services/notebook/browser/models/notebookUtils.ts
View File

@@ -35,14 +35,14 @@ export function getProvidersForFileName(fileName: string, notebookService: INote
return providers;
}
export function getStandardKernelsForProvider(providerId: string, notebookService: INotebookService): IStandardKernelWithProvider[] {
export async function getStandardKernelsForProvider(providerId: string, notebookService: INotebookService): Promise<IStandardKernelWithProvider[]> {
if (!providerId || !notebookService) {
return [];
}
let standardKernels = notebookService.getStandardKernelsForProvider(providerId);
let standardKernels = await notebookService.getStandardKernelsForProvider(providerId);
if (!standardKernels || standardKernels.length === 0) {
// Fall back to using SQL provider instead
standardKernels = notebookService.getStandardKernelsForProvider(SQL_NOTEBOOK_PROVIDER) ?? [];
standardKernels = await notebookService.getStandardKernelsForProvider(SQL_NOTEBOOK_PROVIDER) ?? [];
}
standardKernels.forEach(kernel => {
Object.assign(<IStandardKernelWithProvider>kernel, {

2
src/sql/workbench/services/notebook/browser/notebookService.ts
View File

@@ -74,7 +74,7 @@ export interface INotebookService {
getProvidersForFileType(fileType: string): string[] | undefined;
getStandardKernelsForProvider(provider: string): azdata.nb.IStandardKernel[] | undefined;
getStandardKernelsForProvider(provider: string): Promise<azdata.nb.IStandardKernel[] | undefined>;
getOrCreateSerializationManager(providerId: string, uri: URI): Promise<ISerializationManager>;

101
src/sql/workbench/services/notebook/browser/notebookServiceImpl.ts
View File

@@ -50,7 +50,7 @@ import { IEditorGroupsService } from 'vs/workbench/services/editor/common/editor
import { IEditorInput, IEditorPane } from 'vs/workbench/common/editor';
import { isINotebookInput } from 'sql/workbench/services/notebook/browser/interface';
import { INotebookShowOptions } from 'sql/workbench/api/common/sqlExtHost.protocol';
import { NotebookLanguage } from 'sql/workbench/common/constants';
import { JUPYTER_PROVIDER_ID, NotebookLanguage } from 'sql/workbench/common/constants';
import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
import { SqlSerializationProvider } from 'sql/workbench/services/notebook/browser/sql/sqlSerializationProvider';
@@ -134,6 +134,31 @@ export class ExecuteProviderDescriptor {
}
}
export class StandardKernelsDescriptor {
private _instanceReady = new Deferred<nb.IStandardKernel[]>();
constructor(private readonly _providerId: string, private _instance?: nb.IStandardKernel[]) {
if (_instance) {
this._instanceReady.resolve(_instance);
}
}
public get providerId(): string {
return this._providerId;
}
public get instanceReady(): Promise<nb.IStandardKernel[]> {
return this._instanceReady.promise;
}
public get instance(): nb.IStandardKernel[] | undefined {
return this._instance;
}
public set instance(value: nb.IStandardKernel[]) {
this._instance = value;
this._instanceReady.resolve(value);
}
}
export const NotebookUriNotDefined = localize('notebookUriNotDefined', "No URI was passed when creating a notebook manager");
export const NotebookServiceNoProviderRegistered = localize('notebookServiceNoProvider', "Notebook provider does not exist");
export const FailToSaveTrustState = 'Failed to save trust state to cache';
@@ -154,7 +179,7 @@ export class NotebookService extends Disposable implements INotebookService {
private _onNotebookEditorRename = new Emitter<INotebookEditor>();
private _editors = new Map<string, INotebookEditor>();
private _fileToProviderDescriptions = new Map<string, ProviderDescriptionRegistration[]>();
private _providerToStandardKernels = new Map<string, nb.IStandardKernel[]>();
private _providerToStandardKernels = new Map<string, StandardKernelsDescriptor>();
private _registrationComplete = new Deferred<void>();
private _isRegistrationComplete = false;
private _trustedCacheQueue: URI[] = [];
@@ -291,13 +316,14 @@ export class NotebookService extends Disposable implements INotebookService {
let sqlNotebookKernels = this._providerToStandardKernels.get(notebookConstants.SQL);
if (sqlNotebookKernels) {
let sqlConnectionTypes = this._queryManagementService.getRegisteredProviders();
let kernel = sqlNotebookKernels.find(p => p.name === notebookConstants.SQL);
let kernel = sqlNotebookKernels.instance.find(p => p.name === notebookConstants.SQL);
if (kernel) {
this._providerToStandardKernels.set(notebookConstants.SQL, [{
let descriptor = new StandardKernelsDescriptor(notebookConstants.SQL, [{
name: notebookConstants.SQL,
displayName: notebookConstants.SQL,
connectionProviderIds: sqlConnectionTypes
}]);
this._providerToStandardKernels.set(notebookConstants.SQL, descriptor);
}
}
this._isRegistrationComplete = true;
@@ -325,6 +351,17 @@ export class NotebookService extends Disposable implements INotebookService {
this._executeProviders.set(p.id, new ExecuteProviderDescriptor(p.id));
}
this.addStandardKernels(registration);
} else {
// Standard kernels might get registered later for VSCode notebooks, so add a descriptor to wait on
let descriptor = new StandardKernelsDescriptor(p.id);
this._providerToStandardKernels.set(p.id.toUpperCase(), descriptor);
}
// Emit activation event if the provider is not one of the default options
if (p.id !== SQL_NOTEBOOK_PROVIDER && p.id !== JUPYTER_PROVIDER_ID) {
this._extensionService.whenInstalledExtensionsRegistered().then(() => {
this._extensionService.activateByEvent(`onNotebook:${p.id}`).catch(err => onUnexpectedError(err));
}).catch(err => onUnexpectedError(err));
}
}
@@ -392,18 +429,24 @@ export class NotebookService extends Disposable implements INotebookService {
// kernels to the dropdown
private addStandardKernels(provider: ProviderDescriptionRegistration) {
let providerUpperCase = provider.provider.toUpperCase();
let standardKernels = this._providerToStandardKernels.get(providerUpperCase);
let descriptor = this._providerToStandardKernels.get(providerUpperCase);
if (!descriptor) {
descriptor = new StandardKernelsDescriptor(provider.provider);
}
let standardKernels = descriptor.instance;
if (!standardKernels) {
standardKernels = [];
}
provider.standardKernels.forEach(kernel => {
standardKernels.push(kernel);
});
// Filter out unusable kernels when running on a SAW
if (this.productService.quality === 'saw') {
standardKernels = standardKernels.filter(kernel => !kernel.blockedOnSAW);
}
this._providerToStandardKernels.set(providerUpperCase, standardKernels);
descriptor.instance = standardKernels;
this._providerToStandardKernels.set(providerUpperCase, descriptor);
}
getSupportedFileExtensions(): string[] {
@@ -415,8 +458,12 @@ export class NotebookService extends Disposable implements INotebookService {
return providers?.map(provider => provider.provider);
}
getStandardKernelsForProvider(provider: string): nb.IStandardKernel[] | undefined {
return this._providerToStandardKernels.get(provider.toUpperCase());
public async getStandardKernelsForProvider(provider: string): Promise<nb.IStandardKernel[] | undefined> {
let descriptor = this._providerToStandardKernels.get(provider.toUpperCase());
if (descriptor) {
return this.waitOnStandardKernelsAvailability(descriptor);
}
return undefined;
}
private shutdown(): void {
@@ -588,11 +635,12 @@ export class NotebookService extends Disposable implements INotebookService {
private async getExecuteProviderInstance(providerId: string, timeout?: number): Promise<IExecuteProvider> {
let providerDescriptor = this._executeProviders.get(providerId);
let kernelDescriptor = this._providerToStandardKernels.get(providerId.toUpperCase());
let instance: IExecuteProvider;
// Try get from actual provider, waiting on its registration
if (providerDescriptor) {
if (!providerDescriptor.instance) {
if (providerDescriptor && kernelDescriptor) {
if (!providerDescriptor.instance || !kernelDescriptor.instance) {
// Await extension registration before awaiting provider registration
try {
await this._extensionService.whenInstalledExtensionsRegistered();
@@ -600,6 +648,12 @@ export class NotebookService extends Disposable implements INotebookService {
this._logService.error(error);
}
instance = await this.waitOnExecuteProviderAvailability(providerDescriptor, timeout);
if (instance) {
let kernels = await this.waitOnStandardKernelsAvailability(kernelDescriptor, timeout);
if (!kernels) {
instance = undefined;
}
}
} else {
instance = providerDescriptor.instance;
}
@@ -621,9 +675,12 @@ export class NotebookService extends Disposable implements INotebookService {
private waitOnSerializationProviderAvailability(providerDescriptor: SerializationProviderDescriptor, timeout?: number): Promise<ISerializationProvider | undefined> {
// Wait up to 30 seconds for the provider to be registered
timeout = timeout ?? 30000;
let promises: Promise<ISerializationProvider>[] = [
let promises: Promise<ISerializationProvider | undefined>[] = [
providerDescriptor.instanceReady,
new Promise<ISerializationProvider>((resolve, reject) => setTimeout(() => resolve(undefined), timeout))
new Promise<ISerializationProvider | undefined>((resolve, reject) => setTimeout(() => {
onUnexpectedError(localize('serializationProviderTimeout', 'Waiting for Serialization Provider availability timed out for notebook provider \'{0}\'', providerDescriptor.providerId));
resolve(undefined);
}, timeout))
];
return Promise.race(promises);
}
@@ -631,9 +688,25 @@ export class NotebookService extends Disposable implements INotebookService {
private waitOnExecuteProviderAvailability(providerDescriptor: ExecuteProviderDescriptor, timeout?: number): Promise<IExecuteProvider | undefined> {
// Wait up to 30 seconds for the provider to be registered
timeout = timeout ?? 30000;
let promises: Promise<IExecuteProvider>[] = [
let promises: Promise<IExecuteProvider | undefined>[] = [
providerDescriptor.instanceReady,
new Promise<IExecuteProvider>((resolve, reject) => setTimeout(() => resolve(undefined), timeout))
new Promise<IExecuteProvider | undefined>((resolve, reject) => setTimeout(() => {
onUnexpectedError(localize('executeProviderTimeout', 'Waiting for Execute Provider availability timed out for notebook provider \'{0}\'', providerDescriptor.providerId));
resolve(undefined);
}, timeout))
];
return Promise.race(promises);
}
private waitOnStandardKernelsAvailability(kernelsDescriptor: StandardKernelsDescriptor, timeout?: number): Promise<nb.IStandardKernel[] | undefined> {
// Wait up to 30 seconds for the kernels to be registered
timeout = timeout ?? 30000;
let promises: Promise<nb.IStandardKernel[] | undefined>[] = [
kernelsDescriptor.instanceReady,
new Promise<nb.IStandardKernel[] | undefined>((resolve, reject) => setTimeout(() => {
onUnexpectedError(localize('standardKernelsTimeout', 'Waiting for Standard Kernels availability timed out for notebook provider \'{0}\'', kernelsDescriptor.providerId));
resolve(undefined);
}, timeout))
];
return Promise.race(promises);
}