Skip to content

Commit 4678416

Browse files
Merge branch 'main' into themis/diff7
2 parents 7c0478c + c49bed2 commit 4678416

7 files changed

Lines changed: 185 additions & 21 deletions

File tree

sqlmesh/lsp/custom.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ class RenderModelEntry(PydanticModel):
3636

3737
name: str
3838
fqn: str
39-
description: str
39+
description: t.Optional[str] = None
4040
rendered_query: str
4141

4242

vscode/extension/src/commands/renderModel.ts

Lines changed: 35 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,12 @@ import * as vscode from 'vscode'
22
import { LSPClient } from '../lsp/lsp'
33
import { isErr } from '@bus/result'
44
import { RenderModelEntry } from '../lsp/custom'
5+
import { RenderedModelProvider } from '../providers/renderedModelProvider'
56

6-
export function renderModel(lspClient?: LSPClient) {
7+
export function renderModel(
8+
lspClient?: LSPClient,
9+
renderedModelProvider?: RenderedModelProvider,
10+
) {
711
return async () => {
812
// Get the current active editor
913
const activeEditor = vscode.window.activeTextEditor
@@ -42,10 +46,10 @@ export function renderModel(lspClient?: LSPClient) {
4246
// If multiple models, let user choose
4347
let selectedModel: RenderModelEntry
4448
if (result.value.models.length > 1) {
45-
const items = result.value.models.map((model: RenderModelEntry) => ({
49+
const items = result.value.models.map(model => ({
4650
label: model.name,
4751
description: model.fqn,
48-
detail: model.description,
52+
detail: model.description ? model.description : undefined,
4953
model: model,
5054
}))
5155

@@ -62,29 +66,43 @@ export function renderModel(lspClient?: LSPClient) {
6266
selectedModel = result.value.models[0]
6367
}
6468

65-
// Create a new untitled document with the rendered SQL
66-
const document = await vscode.workspace.openTextDocument({
67-
language: 'sql',
68-
content: selectedModel.rendered_query,
69-
})
69+
if (!renderedModelProvider) {
70+
vscode.window.showErrorMessage('Rendered model provider not available')
71+
return
72+
}
73+
74+
// Store the rendered content and get a virtual URI
75+
const uri = renderedModelProvider.storeRenderedModel(
76+
selectedModel.name,
77+
selectedModel.rendered_query,
78+
)
79+
80+
// Open the virtual document
81+
const document = await vscode.workspace.openTextDocument(uri)
7082

7183
// Determine the view column for side-by-side display
72-
let viewColumn: vscode.ViewColumn
73-
if (activeEditor) {
74-
// Open beside the current editor
75-
viewColumn = activeEditor.viewColumn
76-
? activeEditor.viewColumn + 1
77-
: vscode.ViewColumn.Two
78-
} else {
79-
// If no active editor, open in column two
80-
viewColumn = vscode.ViewColumn.Two
84+
// Find the rightmost column with an editor
85+
let maxColumn = vscode.ViewColumn.One
86+
for (const editor of vscode.window.visibleTextEditors) {
87+
if (editor.viewColumn && editor.viewColumn > maxColumn) {
88+
maxColumn = editor.viewColumn
89+
}
8190
}
8291

92+
// Open in the next column after the rightmost editor
93+
const viewColumn = maxColumn + 1
94+
8395
// Open the document in the editor as a preview (preview: true is default)
8496
await vscode.window.showTextDocument(document, {
8597
viewColumn: viewColumn,
8698
preview: true,
8799
preserveFocus: false,
88100
})
101+
102+
// Execute "Keep Open" command to convert preview tab to permanent tab
103+
await vscode.commands.executeCommand('workbench.action.keepEditor')
104+
105+
// Explicitly set the language mode to SQL for syntax highlighting
106+
await vscode.languages.setTextDocumentLanguage(document, 'sql')
89107
}
90108
}

vscode/extension/src/extension.ts

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import {
2222
} from './utilities/errors'
2323
import { selector, completionProvider } from './completion/completion'
2424
import { LineagePanel } from './webviews/lineagePanel'
25+
import { RenderedModelProvider } from './providers/renderedModelProvider'
2526

2627
let lspClient: LSPClient | undefined
2728

@@ -65,10 +66,20 @@ export async function activate(context: vscode.ExtensionContext) {
6566

6667
lspClient = new LSPClient()
6768

69+
// Create and register the rendered model provider
70+
const renderedModelProvider = new RenderedModelProvider()
71+
context.subscriptions.push(
72+
vscode.workspace.registerTextDocumentContentProvider(
73+
RenderedModelProvider.getScheme(),
74+
renderedModelProvider,
75+
),
76+
renderedModelProvider,
77+
)
78+
6879
context.subscriptions.push(
6980
vscode.commands.registerCommand(
7081
'sqlmesh.renderModel',
71-
renderModel(lspClient),
82+
renderModel(lspClient, renderedModelProvider),
7283
),
7384
)
7485

vscode/extension/src/lsp/custom.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ interface RenderModelResponse {
2121
export interface RenderModelEntry {
2222
name: string
2323
fqn: string
24-
description: string
24+
description: string | null | undefined
2525
rendered_query: string
2626
}
2727

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
import * as vscode from 'vscode'
2+
3+
/**
4+
* Content provider for read-only rendered SQL models
5+
*/
6+
export class RenderedModelProvider
7+
implements vscode.TextDocumentContentProvider
8+
{
9+
private static readonly scheme = 'sqlmesh-rendered'
10+
11+
private renderedModels = new Map<string, string>()
12+
13+
// Event emitter for content changes
14+
private _onDidChange = new vscode.EventEmitter<vscode.Uri>()
15+
readonly onDidChange = this._onDidChange.event
16+
17+
/**
18+
* Provide text content for a given URI
19+
*/
20+
provideTextDocumentContent(uri: vscode.Uri): string {
21+
const key = uri.toString()
22+
return this.renderedModels.get(key) || ''
23+
}
24+
25+
/**
26+
* Store rendered model content and create a URI for it
27+
*/
28+
storeRenderedModel(modelName: string, content: string): vscode.Uri {
29+
const fileName = `${modelName} (rendered)`
30+
// Add a timestamp to make the URI unique for each render
31+
const timestamp = Date.now()
32+
// Use vscode.Uri.from for proper URI construction
33+
const uri = vscode.Uri.from({
34+
scheme: RenderedModelProvider.scheme,
35+
path: fileName,
36+
fragment: timestamp.toString(),
37+
})
38+
this.renderedModels.set(uri.toString(), content)
39+
this._onDidChange.fire(uri)
40+
return uri
41+
}
42+
43+
/**
44+
* Get the URI scheme for rendered models
45+
*/
46+
static getScheme(): string {
47+
return this.scheme
48+
}
49+
50+
/**
51+
* Clean up old rendered models to prevent memory leaks
52+
*/
53+
dispose() {
54+
this.renderedModels.clear()
55+
this._onDidChange.dispose()
56+
}
57+
}

vscode/extension/tests/render.spec.ts

Lines changed: 77 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,85 @@ test('Render works correctly', async () => {
3030

3131
// Check if the model is rendered by check if "`oi`.`order_id` AS `order_id`," is in the window
3232
await expect(window.locator('text="marketing"."customer_id" AS')).toBeVisible();
33+
await expect(window.locator('text=sushi.customers (rendered)')).toBeVisible();
3334

3435
await close();
3536
} finally {
3637
await fs.remove(tempDir);
3738
}
38-
});
39+
});
40+
41+
test('Render works correctly with model without a description', async () => {
42+
const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-'));
43+
await fs.copy(SUSHI_SOURCE_PATH, tempDir);
44+
45+
try {
46+
const { window, close } = await startVSCode(tempDir);
47+
48+
// Wait for the models folder to be visible
49+
await window.waitForSelector('text=models');
50+
51+
// Click on the models folder, excluding external_models
52+
await window.getByRole('treeitem', { name: 'models', exact: true }).locator('a').click();
53+
54+
// Open the latest_order model
55+
await window.getByRole('treeitem', { name: 'latest_order.sql', exact: true }).locator('a').click();
56+
57+
await window.waitForSelector('text=custom_full_with_custom_kind');
58+
await window.waitForSelector('text=Loaded SQLMesh Context')
59+
60+
// Render the model
61+
await window.keyboard.press(process.platform === 'darwin' ? 'Meta+Shift+P' : 'Control+Shift+P');
62+
await window.keyboard.type('Render Model');
63+
await window.keyboard.press('Enter');
64+
65+
// Check if the model is rendered correctly
66+
await expect(window.locator('text="orders"."id" AS "id",')).toBeVisible();
67+
await expect(window.locator('text=sushi.latest_order (rendered)')).toBeVisible();
68+
69+
await close();
70+
} finally {
71+
await fs.remove(tempDir);
72+
}
73+
});
74+
75+
test('Render works correctly with every rendered model opening a new tab', async () => {
76+
const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-'));
77+
await fs.copy(SUSHI_SOURCE_PATH, tempDir);
78+
79+
try {
80+
const { window, close } = await startVSCode(tempDir);
81+
82+
// Wait for the models folder to be visible
83+
await window.waitForSelector('text=models');
84+
await window.getByRole('treeitem', { name: 'models', exact: true }).locator('a').click();
85+
await window.getByRole('treeitem', { name: 'latest_order.sql', exact: true }).locator('a').click();
86+
await window.waitForSelector('text=custom_full_with_custom_kind');
87+
await window.waitForSelector('text=Loaded SQLMesh Context')
88+
89+
// Render the model
90+
await window.keyboard.press(process.platform === 'darwin' ? 'Meta+Shift+P' : 'Control+Shift+P');
91+
await window.keyboard.type('Render Model');
92+
await window.keyboard.press('Enter');
93+
94+
// Check if the model is rendered correctly
95+
await expect(window.locator('text=sushi.latest_order (rendered)')).toBeVisible();
96+
97+
// Open the customers model
98+
await window.getByRole('treeitem', { name: 'customers.sql', exact: true }).locator('a').click();
99+
await window.waitForSelector('text=grain');
100+
101+
// Render the customers model
102+
await window.keyboard.press(process.platform === 'darwin' ? 'Meta+Shift+P' : 'Control+Shift+P');
103+
await window.keyboard.type('Render Model');
104+
await window.keyboard.press('Enter');
105+
106+
// Assert both tabs exist
107+
await expect(window.locator('text=sushi.latest_order (rendered)')).toBeVisible();
108+
await expect(window.locator('text=sushi.customers (rendered)')).toBeVisible();
109+
110+
await close();
111+
} finally {
112+
await fs.remove(tempDir);
113+
}
114+
})

vscode/extension/tests/utils.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,8 @@ export const startVSCode = async (workspaceDir: string): Promise<{
4040
args,
4141
});
4242
const window = await electronApp.firstWindow();
43+
await window.waitForLoadState('domcontentloaded');
44+
await window.waitForLoadState('networkidle');
4345
return { window, close: async () => {
4446
await electronApp.close();
4547
await fs.removeSync(userDataDir);

0 commit comments

Comments
 (0)