Skip to content

Commit f61e38f

Browse files
jankucajankuca
committed
add SQL block variable info and input for changing
1 parent 55d8454 commit f61e38f

File tree

2 files changed

+192
-31
lines changed

2 files changed

+192
-31
lines changed

‎src/notebooks/deepnote/sqlCellStatusBarProvider.ts‎

Lines changed: 99 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,14 @@ import {
44
NotebookCell,
55
NotebookCellStatusBarItem,
66
NotebookCellStatusBarItemProvider,
7-
NotebookDocument,
7+
NotebookEdit,
88
ProviderResult,
9+
WorkspaceEdit,
10+
commands,
911
l10n,
10-
notebooks
12+
notebooks,
13+
window,
14+
workspace
1115
} from 'vscode';
1216
import { inject, injectable } from 'inversify';
1317

@@ -18,7 +22,7 @@ import { IIntegrationStorage } from './integrations/types';
1822
import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes';
1923

2024
/**
21-
* Provides status bar items for SQL cells showing the integration name
25+
* Provides status bar items for SQL cells showing the integration name and variable name
2226
*/
2327
@injectable()
2428
export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvider, IExtensionSyncActivationService {
@@ -42,6 +46,13 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid
4246
})
4347
);
4448

49+
// Register command to update SQL variable name
50+
this.disposables.push(
51+
commands.registerCommand('deepnote.updateSqlVariableName', async (cell: NotebookCell) => {
52+
await this.updateVariableName(cell);
53+
})
54+
);
55+
4556
// Dispose our emitter with the extension
4657
this.disposables.push(this._onDidChangeCellStatusBarItems);
4758
}
@@ -59,18 +70,7 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid
5970
return undefined;
6071
}
6172

62-
// Get the integration ID from cell metadata
63-
const integrationId = this.getIntegrationId(cell);
64-
if (!integrationId) {
65-
return undefined;
66-
}
67-
68-
// Don't show status bar for the internal DuckDB integration
69-
if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) {
70-
return undefined;
71-
}
72-
73-
return this.createStatusBarItem(cell.notebook, integrationId);
73+
return this.createStatusBarItems(cell);
7474
}
7575

7676
private getIntegrationId(cell: NotebookCell): string | undefined {
@@ -86,11 +86,29 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid
8686
return undefined;
8787
}
8888

89-
private async createStatusBarItem(
90-
notebook: NotebookDocument,
89+
private async createStatusBarItems(cell: NotebookCell): Promise<NotebookCellStatusBarItem[]> {
90+
const items: NotebookCellStatusBarItem[] = [];
91+
92+
// Add integration status bar item if integration ID is present
93+
const integrationId = this.getIntegrationId(cell);
94+
if (integrationId && integrationId !== DATAFRAME_SQL_INTEGRATION_ID) {
95+
const integrationItem = await this.createIntegrationStatusBarItem(cell, integrationId);
96+
if (integrationItem) {
97+
items.push(integrationItem);
98+
}
99+
}
100+
101+
// Always add variable status bar item for SQL cells
102+
items.push(this.createVariableStatusBarItem(cell));
103+
104+
return items;
105+
}
106+
107+
private async createIntegrationStatusBarItem(
108+
cell: NotebookCell,
91109
integrationId: string
92110
): Promise<NotebookCellStatusBarItem | undefined> {
93-
const projectId = notebook.metadata?.deepnoteProjectId;
111+
const projectId = cell.notebook.metadata?.deepnoteProjectId;
94112
if (!projectId) {
95113
return undefined;
96114
}
@@ -111,4 +129,67 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid
111129
}
112130
};
113131
}
132+
133+
private createVariableStatusBarItem(cell: NotebookCell): NotebookCellStatusBarItem {
134+
const variableName = this.getVariableName(cell);
135+
136+
return {
137+
text: `Variable: ${variableName}`,
138+
alignment: 1, // NotebookCellStatusBarAlignment.Left
139+
tooltip: l10n.t('Variable name for SQL query result\nClick to change'),
140+
command: {
141+
title: l10n.t('Change Variable Name'),
142+
command: 'deepnote.updateSqlVariableName',
143+
arguments: [cell]
144+
}
145+
};
146+
}
147+
148+
private getVariableName(cell: NotebookCell): string {
149+
const metadata = cell.metadata;
150+
if (metadata && typeof metadata === 'object') {
151+
const variableName = (metadata as Record<string, unknown>).deepnote_variable_name;
152+
if (typeof variableName === 'string' && variableName) {
153+
return variableName;
154+
}
155+
}
156+
157+
return 'df';
158+
}
159+
160+
private async updateVariableName(cell: NotebookCell): Promise<void> {
161+
const currentVariableName = this.getVariableName(cell);
162+
163+
const newVariableName = await window.showInputBox({
164+
prompt: l10n.t('Enter variable name for SQL query result'),
165+
value: currentVariableName,
166+
validateInput: (value) => {
167+
if (!value) {
168+
return l10n.t('Variable name cannot be empty');
169+
}
170+
if (!/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(value)) {
171+
return l10n.t('Variable name must be a valid Python identifier');
172+
}
173+
return undefined;
174+
}
175+
});
176+
177+
if (newVariableName === undefined || newVariableName === currentVariableName) {
178+
return;
179+
}
180+
181+
// Update cell metadata
182+
const edit = new WorkspaceEdit();
183+
const updatedMetadata = {
184+
...cell.metadata,
185+
deepnote_variable_name: newVariableName
186+
};
187+
188+
edit.set(cell.notebook.uri, [NotebookEdit.updateCellMetadata(cell.index, updatedMetadata)]);
189+
190+
await workspace.applyEdit(edit);
191+
192+
// Trigger status bar update
193+
this._onDidChangeCellStatusBarItems.fire();
194+
}
114195
}

‎src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts‎

Lines changed: 93 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -38,25 +38,45 @@ suite('SqlCellStatusBarProvider', () => {
3838
assert.isUndefined(result);
3939
});
4040

41-
test('returns undefined for SQL cells without integration ID', async () => {
41+
test('returns only variable status bar item for SQL cells without integration ID', async () => {
4242
const cell = createMockCell('sql', {});
4343

4444
const result = await provider.provideCellStatusBarItems(cell, cancellationToken);
4545

46-
assert.isUndefined(result);
46+
assert.isDefined(result);
47+
assert.isArray(result);
48+
const items = result as any[];
49+
assert.strictEqual(items.length, 1);
50+
51+
// Check variable status bar item
52+
const variableItem = items[0];
53+
assert.strictEqual(variableItem.text, 'Variable: df');
54+
assert.strictEqual(variableItem.alignment, 1);
55+
assert.isDefined(variableItem.command);
56+
assert.strictEqual(variableItem.command.command, 'deepnote.updateSqlVariableName');
4757
});
4858

49-
test('returns undefined for SQL cells with dataframe integration ID', async () => {
59+
test('returns only variable status bar item for SQL cells with dataframe integration ID', async () => {
5060
const cell = createMockCell('sql', {
5161
sql_integration_id: DATAFRAME_SQL_INTEGRATION_ID
5262
});
5363

5464
const result = await provider.provideCellStatusBarItems(cell, cancellationToken);
5565

56-
assert.isUndefined(result);
66+
assert.isDefined(result);
67+
assert.isArray(result);
68+
const items = result as any[];
69+
assert.strictEqual(items.length, 1);
70+
71+
// Check variable status bar item
72+
const variableItem = items[0];
73+
assert.strictEqual(variableItem.text, 'Variable: df');
74+
assert.strictEqual(variableItem.alignment, 1);
75+
assert.isDefined(variableItem.command);
76+
assert.strictEqual(variableItem.command.command, 'deepnote.updateSqlVariableName');
5777
});
5878

59-
test('returns status bar item for SQL cell with integration ID', async () => {
79+
test('returns status bar items for SQL cell with integration ID', async () => {
6080
const integrationId = 'postgres-123';
6181
const cell = createMockCell(
6282
'sql',
@@ -82,11 +102,24 @@ suite('SqlCellStatusBarProvider', () => {
82102
const result = await provider.provideCellStatusBarItems(cell, cancellationToken);
83103

84104
assert.isDefined(result);
85-
assert.strictEqual((result as any).text, '$(database) My Postgres DB');
86-
assert.strictEqual((result as any).alignment, 1); // NotebookCellStatusBarAlignment.Left
87-
assert.isDefined((result as any).command);
88-
assert.strictEqual((result as any).command.command, 'deepnote.manageIntegrations');
89-
assert.deepStrictEqual((result as any).command.arguments, [integrationId]);
105+
assert.isArray(result);
106+
const items = result as any[];
107+
assert.strictEqual(items.length, 2);
108+
109+
// Check integration status bar item
110+
const integrationItem = items[0];
111+
assert.strictEqual(integrationItem.text, '$(database) My Postgres DB');
112+
assert.strictEqual(integrationItem.alignment, 1);
113+
assert.isDefined(integrationItem.command);
114+
assert.strictEqual(integrationItem.command.command, 'deepnote.manageIntegrations');
115+
assert.deepStrictEqual(integrationItem.command.arguments, [integrationId]);
116+
117+
// Check variable status bar item
118+
const variableItem = items[1];
119+
assert.strictEqual(variableItem.text, 'Variable: df');
120+
assert.strictEqual(variableItem.alignment, 1);
121+
assert.isDefined(variableItem.command);
122+
assert.strictEqual(variableItem.command.command, 'deepnote.updateSqlVariableName');
90123
});
91124

92125
test('shows "Unknown integration (configure)" when config not found', async () => {
@@ -106,18 +139,65 @@ suite('SqlCellStatusBarProvider', () => {
106139
const result = await provider.provideCellStatusBarItems(cell, cancellationToken);
107140

108141
assert.isDefined(result);
109-
assert.strictEqual((result as any).text, '$(database) Unknown integration (configure)');
142+
assert.isArray(result);
143+
const items = result as any[];
144+
assert.strictEqual(items.length, 2);
145+
assert.strictEqual(items[0].text, '$(database) Unknown integration (configure)');
146+
assert.strictEqual(items[1].text, 'Variable: df');
110147
});
111148

112-
test('returns undefined when notebook has no project ID', async () => {
149+
test('returns only variable item when notebook has no project ID', async () => {
113150
const integrationId = 'postgres-123';
114151
const cell = createMockCell('sql', {
115152
sql_integration_id: integrationId
116153
});
117154

118155
const result = await provider.provideCellStatusBarItems(cell, cancellationToken);
119156

120-
assert.isUndefined(result);
157+
assert.isDefined(result);
158+
assert.isArray(result);
159+
const items = result as any[];
160+
assert.strictEqual(items.length, 1);
161+
162+
// Check variable status bar item is still shown
163+
const variableItem = items[0];
164+
assert.strictEqual(variableItem.text, 'Variable: df');
165+
});
166+
167+
test('shows custom variable name when set in metadata', async () => {
168+
const integrationId = 'postgres-123';
169+
const cell = createMockCell(
170+
'sql',
171+
{
172+
sql_integration_id: integrationId,
173+
deepnote_variable_name: 'my_results'
174+
},
175+
{
176+
deepnoteProjectId: 'project-1'
177+
}
178+
);
179+
180+
when(integrationStorage.getProjectIntegrationConfig(anything(), anything())).thenResolve({
181+
id: integrationId,
182+
name: 'My Postgres DB',
183+
type: IntegrationType.Postgres,
184+
host: 'localhost',
185+
port: 5432,
186+
database: 'test',
187+
username: 'user',
188+
password: 'pass'
189+
});
190+
191+
const result = await provider.provideCellStatusBarItems(cell, cancellationToken);
192+
193+
assert.isDefined(result);
194+
assert.isArray(result);
195+
const items = result as any[];
196+
assert.strictEqual(items.length, 2);
197+
198+
// Check variable status bar item shows custom name
199+
const variableItem = items[1];
200+
assert.strictEqual(variableItem.text, 'Variable: my_results');
121201
});
122202

123203
function createMockCell(

0 commit comments

Comments
 (0)