Skip to content

Commit

Permalink
Remove get method from the UriStorage class (#14029)
Browse files Browse the repository at this point in the history
  • Loading branch information
DonJayamanne authored Aug 1, 2023
1 parent 45a9f56 commit 3ca8776
Show file tree
Hide file tree
Showing 8 changed files with 127 additions and 135 deletions.
6 changes: 0 additions & 6 deletions src/kernels/jupyter/connection/serverUriStorage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -105,12 +105,6 @@ export class JupyterServerUriStorage extends Disposables implements IJupyterServ
await this.newStorage.migrateMRU();
await Promise.all([this.oldStorage.clear(), this.newStorage.clear()]);
}
public async get(server: JupyterServerProviderHandle): Promise<IJupyterServerUriEntry | undefined> {
this.hookupStorageEvents();
await this.newStorage.migrateMRU();
const savedList = await this.getAll();
return savedList.find((item) => item.provider.id === server.id && item.provider.handle === server.handle);
}
public async add(
jupyterHandle: JupyterServerProviderHandle,
options?: { time: number; displayName: string }
Expand Down
79 changes: 48 additions & 31 deletions src/kernels/jupyter/connection/serverUriStorage.unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -769,7 +769,11 @@ suite('Server Uri Storage', async () => {
const timeOfNewHandle2BeforeUpdate = beforeUpdate.find((item) => item.provider.handle === 'NewHandle2')!;
assert.ok(timeOfNewHandle2BeforeUpdate);
await sleep(10);
await serverUriStorage.update({ id: 'NewId2', handle: 'NewHandle2', extensionId: JVSC_EXTENSION_ID });
await serverUriStorage.update({
id: 'NewId2',
handle: 'NewHandle2',
extensionId: JVSC_EXTENSION_ID
});
const afterUpdate = await serverUriStorage.getAll();
const timeOfNewHandle2AfterUpdate = afterUpdate.find((item) => item.provider.handle === 'NewHandle2')!;
assert.ok(timeOfNewHandle2BeforeUpdate);
Expand Down Expand Up @@ -841,7 +845,7 @@ suite('Server Uri Storage', async () => {
verify(fs.writeFile(anything(), anything())).atLeast(1);
assert.strictEqual(all.length, 0);
});
test('Add 10 new entries & add 11th, and add mroe and remove', async function () {
test('Add 10 new entries & add 11th, and add more and remove', async function () {
generateDummyData(8, true);
when(fs.exists(anything())).thenResolve(true);
when(fs.exists(uriEquals(globalStorageUri))).thenResolve(true);
Expand Down Expand Up @@ -889,7 +893,11 @@ suite('Server Uri Storage', async () => {
onDidAddEvent.reset();
onDidRemoveEvent.reset();
await serverUriStorage.add({ handle: 'NewHandle11', id: 'NewId11', extensionId: JVSC_EXTENSION_ID });
await serverUriStorage.update({ id: 'NewId11', handle: 'NewHandle11', extensionId: JVSC_EXTENSION_ID });
await serverUriStorage.update({
id: 'NewId11',
handle: 'NewHandle11',
extensionId: JVSC_EXTENSION_ID
});

all = await serverUriStorage.getAll();
assert.strictEqual(all.length, 10);
Expand Down Expand Up @@ -932,14 +940,16 @@ suite('Server Uri Storage', async () => {
}

// Should exist.
const server1 = await serverUriStorage.get({
id: UserJupyterServerPickerProviderId,
handle: 'handle1',
extensionId: JVSC_EXTENSION_ID
});
const time = await (
await serverUriStorage.getAll()
).find(
(item) =>
item.provider.id === UserJupyterServerPickerProviderId &&
item.provider.extensionId === JVSC_EXTENSION_ID &&
item.provider.handle === 'handle1'
);

assert.strictEqual(server1?.provider.id, UserJupyterServerPickerProviderId);
assert.strictEqual(server1?.provider.handle, 'handle1');
assert.ok(time);

// Remove this.
await serverUriStorage.remove({
Expand All @@ -949,33 +959,40 @@ suite('Server Uri Storage', async () => {
});

assert.isUndefined(
await serverUriStorage.get({
id: UserJupyterServerPickerProviderId,
handle: 'handle1',
extensionId: JVSC_EXTENSION_ID
})
await (
await serverUriStorage.getAll()
).find(
(item) =>
item.provider.id === UserJupyterServerPickerProviderId &&
item.provider.extensionId === JVSC_EXTENSION_ID &&
item.provider.handle === 'handle1'
)
);

// Bogus
const serverBogus = await serverUriStorage.get({
id: 'Bogus',
handle: 'handle1',
extensionId: JVSC_EXTENSION_ID
});

assert.isUndefined(serverBogus);
const noTime = await (
await serverUriStorage.getAll()
).find(
(item) =>
item.provider.id === 'Bogus' &&
item.provider.extensionId === JVSC_EXTENSION_ID &&
item.provider.handle === 'handle1'
);
assert.isUndefined(noTime);

// Add and it should exist.
await serverUriStorage.add({ handle: 'NewHandle11', id: 'NewId11', extensionId: JVSC_EXTENSION_ID });

const newServer = await serverUriStorage.get({
id: 'NewId11',
handle: 'NewHandle11',
extensionId: JVSC_EXTENSION_ID
});
const hastTime = await (
await serverUriStorage.getAll()
).find(
(item) =>
item.provider.id === 'NewId11' &&
item.provider.extensionId === JVSC_EXTENSION_ID &&
item.provider.handle === 'NewHandle11'
);

assert.strictEqual(newServer?.provider.id, 'NewId11');
assert.strictEqual(newServer?.provider.handle, 'NewHandle11');
assert.isOk(hastTime);
});

function generateDummyData(numberOfEntries: number = 2, generateNewDataAsWell: boolean = false) {
Expand All @@ -992,7 +1009,7 @@ suite('Server Uri Storage', async () => {
uris.push(`${uri}${Settings.JupyterServerRemoteLaunchNameSeparator}${displayName}`);
data.push({
index,
time: index
time: Date.now() - 1000 + index
});
itemsInNewStorage.push({
displayName,
Expand All @@ -1001,7 +1018,7 @@ suite('Server Uri Storage', async () => {
handle: `handle${index + 1}`,
extensionId: JVSC_EXTENSION_ID
},
time: index
time: Date.now() - 1000 + index
});
}
when(memento.get(Settings.JupyterServerUriList)).thenReturn(data);
Expand Down
1 change: 0 additions & 1 deletion src/kernels/jupyter/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,6 @@ export interface IJupyterServerUriStorage {
getAll(): Promise<IJupyterServerUriEntry[]>;
remove(serverProviderHandle: JupyterServerProviderHandle): Promise<void>;
clear(): Promise<void>;
get(serverProviderHandle: JupyterServerProviderHandle): Promise<IJupyterServerUriEntry | undefined>;
add(
serverProviderHandle: JupyterServerProviderHandle,
options?: { time: number; displayName: string }
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ import {
import { DisplayOptions } from '../../../kernels/displayOptions';
import { isPythonKernelConnection, isUserRegisteredKernelSpecConnection } from '../../../kernels/helpers';
import { ContributedKernelFinderKind } from '../../../kernels/internalTypes';
import { IJupyterUriProviderRegistration } from '../../../kernels/jupyter/types';
import { IInternalJupyterUriProvider, IJupyterUriProviderRegistration } from '../../../kernels/jupyter/types';
import { initializeInteractiveOrNotebookTelemetryBasedOnUserAction } from '../../../kernels/telemetry/helper';
import { sendKernelTelemetryEvent } from '../../../kernels/telemetry/sendKernelTelemetryEvent';
import {
Expand Down Expand Up @@ -50,7 +50,10 @@ import {
@injectable()
export class KernelSourceCommandHandler implements IExtensionSyncActivationService {
private localDisposables: IDisposable[] = [];
private readonly providerMappings = new Map<string, IDisposable[]>();
private readonly providerMappings = new Map<
string,
{ disposables: IDisposable[]; provider: IInternalJupyterUriProvider }
>();
private kernelSpecsSourceRegistered = false;
constructor(
@inject(IControllerRegistration) private readonly controllerRegistration: IControllerRegistration,
Expand Down Expand Up @@ -165,11 +168,12 @@ export class KernelSourceCommandHandler implements IExtensionSyncActivationServi
);
const existingItems = new Set<string>();
uriRegistration.providers.map((provider) => {
const id = `${provider.extensionId}:${provider.id}`;
if (provider.id === TestingKernelPickerProviderId) {
return;
}
existingItems.add(provider.id);
if (this.providerMappings.has(provider.id)) {
existingItems.add(id);
if (this.providerMappings.has(id)) {
return;
}
const providerItemNb = notebooks.registerKernelSourceActionProvider(JupyterNotebookView, {
Expand Down Expand Up @@ -212,12 +216,12 @@ export class KernelSourceCommandHandler implements IExtensionSyncActivationServi
});
this.localDisposables.push(providerItemNb);
this.localDisposables.push(providerItemIW);
this.providerMappings.set(provider.id, [providerItemNb, providerItemIW]);
this.providerMappings.set(id, { disposables: [providerItemNb, providerItemIW], provider });
});
this.providerMappings.forEach((disposables, providerId) => {
if (!existingItems.has(providerId)) {
this.providerMappings.forEach(({ disposables }, id) => {
if (!existingItems.has(id)) {
disposeAllDisposables(disposables);
this.providerMappings.delete(providerId);
this.providerMappings.delete(id);
}
});
}
Expand Down Expand Up @@ -252,10 +256,15 @@ export class KernelSourceCommandHandler implements IExtensionSyncActivationServi
if (!notebook) {
return;
}
const id = `${extensionId}:${providerId}`;
const provider = this.providerMappings.get(id)?.provider;
if (!provider) {
return;
}
const selector = ServiceContainer.instance.get<IRemoteNotebookKernelSourceSelector>(
IRemoteNotebookKernelSourceSelector
);
const kernel = await selector.selectRemoteKernel(notebook, extensionId, providerId);
const kernel = await selector.selectRemoteKernel(notebook, provider);
return this.getSelectedController(notebook, kernel);
}
private async getSelectedController(notebook: NotebookDocument, kernel?: KernelConnectionMetadata) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ import { JupyterServerSelector } from '../../../kernels/jupyter/connection/serve
import {
IJupyterServerUriStorage,
IInternalJupyterUriProvider,
IJupyterUriProviderRegistration,
IRemoteKernelFinder
} from '../../../kernels/jupyter/types';
import { IKernelFinder, KernelConnectionMetadata, RemoteKernelConnectionMetadata } from '../../../kernels/types';
Expand All @@ -42,6 +41,7 @@ import { QuickPickKernelItemProvider } from './quickPickKernelItemProvider';
import { ConnectionQuickPickItem, IQuickPickKernelItemProvider, MultiStepResult } from './types';
import { JupyterConnection } from '../../../kernels/jupyter/connection/jupyterConnection';
import { CreateAndSelectItemFromQuickPick } from './baseKernelSelector';
import { generateIdFromRemoteProvider } from '../../../kernels/jupyter/jupyterUtils';

enum KernelFinderEntityQuickPickType {
KernelFinder = 'finder',
Expand Down Expand Up @@ -73,25 +73,18 @@ export class RemoteNotebookKernelSourceSelector implements IRemoteNotebookKernel
private cancellationTokenSource: CancellationTokenSource | undefined;
constructor(
@inject(IKernelFinder) private readonly kernelFinder: IKernelFinder,
@inject(IJupyterUriProviderRegistration)
private readonly uriProviderRegistration: IJupyterUriProviderRegistration,
@inject(IJupyterServerUriStorage) private readonly serverUriStorage: IJupyterServerUriStorage,
@inject(JupyterServerSelector) private readonly serverSelector: JupyterServerSelector,
@inject(JupyterConnection) private readonly jupyterConnection: JupyterConnection
) {}
public async selectRemoteKernel(
notebook: NotebookDocument,
extensionId: string,
providerId: string
provider: IInternalJupyterUriProvider
): Promise<RemoteKernelConnectionMetadata | undefined> {
// Reject if it's not our type
if (notebook.notebookType !== JupyterNotebookView && notebook.notebookType !== InteractiveWindowView) {
return;
}
const provider = await this.uriProviderRegistration.getProvider(extensionId, providerId);
if (!provider) {
throw new Error(`Remote Provider Id ${providerId} not found`);
}
this.localDisposables.forEach((d) => d.dispose());
this.localDisposables = [];
this.cancellationTokenSource?.cancel();
Expand Down Expand Up @@ -132,34 +125,38 @@ export class RemoteNotebookKernelSourceSelector implements IRemoteNotebookKernel
const servers = this.kernelFinder.registered.filter((info) => info.kind === 'remote') as IRemoteKernelFinder[];
const items: (ContributedKernelFinderQuickPickItem | KernelProviderItemsQuickPickItem | QuickPickItem)[] = [];

for (const server of servers) {
// remote server
const savedURI = await this.serverUriStorage.get(server.serverUri.provider);
if (token.isCancellationRequested) {
return;
}

const idAndHandle = savedURI?.provider;
if (idAndHandle && idAndHandle.id === provider.id) {
// local server
const uriDate = new Date(savedURI.time);
items.push({
type: KernelFinderEntityQuickPickType.KernelFinder,
kernelFinderInfo: server,
idAndHandle,
label: server.displayName,
detail: DataScience.jupyterSelectURIMRUDetail(uriDate),
buttons: provider.removeHandle
? [
{
iconPath: new ThemeIcon('trash'),
tooltip: DataScience.removeRemoteJupyterServerEntryInQuickPick
}
]
: []
});
}
}
await Promise.all(
servers
.filter((s) => s.serverUri.provider.id === provider.id)
.map(async (server) => {
// remote server
const lastUsedTime = await (
await this.serverUriStorage.getAll()
).find(
(item) =>
generateIdFromRemoteProvider(item.provider) ===
generateIdFromRemoteProvider(server.serverUri.provider)
);
if (token.isCancellationRequested || !lastUsedTime) {
return;
}
items.push({
type: KernelFinderEntityQuickPickType.KernelFinder,
kernelFinderInfo: server,
idAndHandle: server.serverUri.provider,
label: server.displayName,
detail: DataScience.jupyterSelectURIMRUDetail(new Date(lastUsedTime.time)),
buttons: provider.removeHandle
? [
{
iconPath: new ThemeIcon('trash'),
tooltip: DataScience.removeRemoteJupyterServerEntryInQuickPick
}
]
: []
});
})
);

if (provider.getQuickPickEntryItems && provider.handleQuickPick) {
if (items.length > 0) {
Expand Down
Loading

0 comments on commit 3ca8776

Please sign in to comment.