Skip to content

Commit a6992e0

Browse files
benfdkingbenfdking
authored
feat(vscode): render model (#4533)
1 parent 780ebae commit a6992e0

10 files changed

Lines changed: 307 additions & 48 deletions

File tree

‎sqlmesh/lsp/context.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,10 @@
33
from sqlmesh.core.context import Context
44
import typing as t
55

6+
from sqlmesh.core.model.definition import SqlModel
7+
from sqlmesh.lsp.custom import RenderModelEntry
8+
from sqlmesh.lsp.uri import URI
9+
610

711
@dataclass
812
class ModelTarget:
@@ -49,3 +53,37 @@ def __init__(self, context: Context) -> None:
4953
**model_map,
5054
**audit_map,
5155
}
56+
57+
58+
def render_model(context: LSPContext, uri: URI) -> t.Iterator[RenderModelEntry]:
59+
target = context.map[uri.to_path()]
60+
if isinstance(target, AuditTarget):
61+
audit = context.context.standalone_audits[target.name]
62+
definition = audit.render_definition(
63+
include_python=False,
64+
render_query=True,
65+
)
66+
rendered_query = [render.sql(dialect=audit.dialect, pretty=True) for render in definition]
67+
yield RenderModelEntry(
68+
name=audit.name,
69+
fqn=audit.fqn,
70+
description=audit.description,
71+
rendered_query="\n\n".join(rendered_query),
72+
)
73+
if isinstance(target, ModelTarget):
74+
for name in target.names:
75+
model = context.context.get_model(name)
76+
if isinstance(model, SqlModel):
77+
rendered_query = [
78+
render.sql(dialect=model.dialect, pretty=True)
79+
for render in model.render_definition(
80+
include_python=False,
81+
render_query=True,
82+
)
83+
]
84+
yield RenderModelEntry(
85+
name=model.name,
86+
fqn=model.fqn,
87+
description=model.description,
88+
rendered_query="\n\n".join(rendered_query),
89+
)

‎sqlmesh/lsp/custom.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,3 +20,29 @@ class AllModelsResponse(PydanticModel):
2020

2121
models: t.List[str]
2222
keywords: t.List[str]
23+
24+
25+
RENDER_MODEL_FEATURE = "sqlmesh/render_model"
26+
27+
28+
class RenderModelRequest(PydanticModel):
29+
textDocumentUri: str
30+
31+
32+
class RenderModelEntry(PydanticModel):
33+
"""
34+
An entry in the rendered model.
35+
"""
36+
37+
name: str
38+
fqn: str
39+
description: str
40+
rendered_query: str
41+
42+
43+
class RenderModelResponse(PydanticModel):
44+
"""
45+
Response to render a model.
46+
"""
47+
48+
models: t.List[RenderModelEntry]

‎sqlmesh/lsp/main.py‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,15 @@
2020
ApiResponseGetModels,
2121
)
2222
from sqlmesh.lsp.completions import get_sql_completions
23-
from sqlmesh.lsp.context import LSPContext, ModelTarget
24-
from sqlmesh.lsp.custom import ALL_MODELS_FEATURE, AllModelsRequest, AllModelsResponse
23+
from sqlmesh.lsp.context import LSPContext, ModelTarget, render_model as render_model_context
24+
from sqlmesh.lsp.custom import (
25+
ALL_MODELS_FEATURE,
26+
RENDER_MODEL_FEATURE,
27+
AllModelsRequest,
28+
AllModelsResponse,
29+
RenderModelRequest,
30+
RenderModelResponse,
31+
)
2532
from sqlmesh.lsp.reference import (
2633
get_references,
2734
)
@@ -88,6 +95,12 @@ def all_models(ls: LanguageServer, params: AllModelsRequest) -> AllModelsRespons
8895
except Exception as e:
8996
return get_sql_completions(None, uri)
9097

98+
@self.server.feature(RENDER_MODEL_FEATURE)
99+
def render_model(ls: LanguageServer, params: RenderModelRequest) -> RenderModelResponse:
100+
uri = URI(params.textDocumentUri)
101+
context = self._context_get_or_load(uri)
102+
return RenderModelResponse(models=list(render_model_context(context, uri)))
103+
91104
@self.server.feature(API_FEATURE)
92105
def api(ls: LanguageServer, request: ApiRequest) -> t.Dict[str, t.Any]:
93106
ls.log_trace(f"API request: {request}")

‎vscode/extension/package.json‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,8 +87,23 @@
8787
"command": "sqlmesh.signout",
8888
"title": "Sign out from Tobiko Cloud",
8989
"description": "SQLMesh"
90+
},
91+
{
92+
"command": "sqlmesh.renderModel",
93+
"title": "Render Model",
94+
"description": "Render the model in the current editor",
95+
"icon": "$(open-preview)"
9096
}
91-
]
97+
],
98+
"menus": {
99+
"editor/title": [
100+
{
101+
"command": "sqlmesh.renderModel",
102+
"when": "resourceExtname == .sql",
103+
"group": "navigation"
104+
}
105+
]
106+
}
92107
},
93108
"scripts": {
94109
"ci": "pnpm run lint && pnpm run compile",

‎vscode/extension/src/commands/renderModel.ts‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
import * as vscode from 'vscode'
2+
import { LSPClient } from '../lsp/lsp'
3+
import { isErr } from '@bus/result'
4+
import { RenderModelEntry } from '../lsp/custom'
5+
6+
export function renderModel(lspClient?: LSPClient) {
7+
return async () => {
8+
// Get the current active editor
9+
const activeEditor = vscode.window.activeTextEditor
10+
11+
if (!activeEditor) {
12+
vscode.window.showErrorMessage('No active editor found')
13+
return
14+
}
15+
16+
if (!lspClient) {
17+
vscode.window.showErrorMessage('LSP client not available')
18+
return
19+
}
20+
21+
// Get the current document URI
22+
const documentUri = activeEditor.document.uri.toString(true)
23+
24+
// Call the render model API
25+
const result = await lspClient.call_custom_method('sqlmesh/render_model', {
26+
textDocumentUri: documentUri,
27+
})
28+
29+
if (isErr(result)) {
30+
vscode.window.showErrorMessage(`Failed to render model: ${result.error}`)
31+
return
32+
}
33+
34+
// Check if we got any models
35+
if (!result.value.models || result.value.models.length === 0) {
36+
vscode.window.showInformationMessage(
37+
'No models found in the current file',
38+
)
39+
return
40+
}
41+
42+
// If multiple models, let user choose
43+
let selectedModel: RenderModelEntry
44+
if (result.value.models.length > 1) {
45+
const items = result.value.models.map((model: RenderModelEntry) => ({
46+
label: model.name,
47+
description: model.fqn,
48+
detail: model.description,
49+
model: model,
50+
}))
51+
52+
const selected = await vscode.window.showQuickPick(items, {
53+
placeHolder: 'Select a model to render',
54+
})
55+
56+
if (!selected) {
57+
return
58+
}
59+
60+
selectedModel = selected.model
61+
} else {
62+
selectedModel = result.value.models[0]
63+
}
64+
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+
})
70+
71+
// 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
81+
}
82+
83+
// Open the document in the editor as a preview (preview: true is default)
84+
await vscode.window.showTextDocument(document, {
85+
viewColumn: viewColumn,
86+
preview: true,
87+
preserveFocus: false,
88+
})
89+
}
90+
}

‎vscode/extension/src/extension.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import { AuthenticationProviderTobikoCloud } from './auth/auth'
1212
import { signOut } from './commands/signout'
1313
import { signIn } from './commands/signin'
1414
import { signInSpecifyFlow } from './commands/signinSpecifyFlow'
15+
import { renderModel } from './commands/renderModel'
1516
import { isErr } from '@bus/result'
1617
import {
1718
handleNotSginedInError,
@@ -64,6 +65,13 @@ export async function activate(context: vscode.ExtensionContext) {
6465

6566
lspClient = new LSPClient()
6667

68+
context.subscriptions.push(
69+
vscode.commands.registerCommand(
70+
'sqlmesh.renderModel',
71+
renderModel(lspClient),
72+
),
73+
)
74+
6775
context.subscriptions.push(
6876
vscode.languages.registerCompletionItemProvider(
6977
selector,

‎vscode/extension/src/lsp/custom.ts‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,32 @@ export interface AllModelsMethod {
44
response: AllModelsResponse
55
}
66

7+
export interface RenderModelMethod {
8+
method: 'sqlmesh/render_model'
9+
request: RenderModelRequest
10+
response: RenderModelResponse
11+
}
12+
13+
interface RenderModelRequest {
14+
textDocumentUri: string
15+
}
16+
17+
interface RenderModelResponse {
18+
models: RenderModelEntry[]
19+
}
20+
21+
export interface RenderModelEntry {
22+
name: string
23+
fqn: string
24+
description: string
25+
rendered_query: string
26+
}
27+
728
// @eslint-disable-next-line @typescript-eslint/consistent-type-definition
8-
export type CustomLSPMethods = AllModelsMethod | AbstractAPICall
29+
export type CustomLSPMethods =
30+
| AllModelsMethod
31+
| AbstractAPICall
32+
| RenderModelMethod
933

1034
interface AllModelsRequest {
1135
textDocument: {

‎vscode/extension/tests/lineage.spec.ts‎

Lines changed: 1 addition & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -2,13 +2,8 @@ import { test, _electron as electron, expect, ElectronApplication, Page } from '
22
import path from 'path';
33
import fs from 'fs-extra';
44
import os from 'os';
5+
import { startVSCode, SUSHI_SOURCE_PATH } from './utils';
56

6-
// Absolute path to the VS Code executable you downloaded in step 1.
7-
const VS_CODE_EXE = fs.readJsonSync('.vscode-test/paths.json').executablePath;
8-
// Where your extension lives on disk
9-
const EXT_PATH = path.resolve(__dirname, '..');
10-
// Where the sushi project lives which we copy from
11-
const SUSHI_SOURCE_PATH = path.join(__dirname, '..', '..', '..', 'examples', 'sushi');
127

138
/**
149
* Helper function to launch VS Code and test lineage with given project path config
@@ -26,44 +21,6 @@ async function testLineageWithProjectPath(
2621
await expect(loadedContextText.first()).toBeVisible({ timeout: 10_000 });
2722
}
2823

29-
/**
30-
* Launch VS Code and return the window and a function to close the app.
31-
* @param workspaceDir The workspace directory to open.
32-
* @returns The window and a function to close the app.
33-
*/
34-
export const startVSCode = async (workspaceDir: string): Promise<{
35-
window: Page,
36-
close: () => Promise<void>,
37-
}> => {
38-
const userDataDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-user-data-'));
39-
const ciArgs = process.env.CI ? [
40-
'--disable-gpu',
41-
'--headless',
42-
'--no-sandbox',
43-
'--disable-dev-shm-usage',
44-
'--window-position=-10000,0',
45-
] : [];
46-
const args = [
47-
...ciArgs,
48-
`--extensionDevelopmentPath=${EXT_PATH}`,
49-
'--disable-workspace-trust',
50-
'--disable-telemetry',
51-
`--user-data-dir=${userDataDir}`,
52-
workspaceDir,
53-
];
54-
const electronApp = await electron.launch({
55-
executablePath: VS_CODE_EXE,
56-
args,
57-
});
58-
const window = await electronApp.firstWindow();
59-
await window.waitForLoadState('domcontentloaded');
60-
await window.waitForLoadState('networkidle');
61-
await window.waitForTimeout(2_000);
62-
return { window, close: async () => {
63-
await electronApp.close();
64-
await fs.remove(userDataDir);
65-
} };
66-
}
6724

6825
test('Lineage panel renders correctly - no project path config (default)', async () => {
6926
const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-'));

‎vscode/extension/tests/render.spec.ts‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
import { test, expect } from '@playwright/test';
2+
import path from 'path';
3+
import fs from 'fs-extra';
4+
import os from 'os';
5+
import { startVSCode, SUSHI_SOURCE_PATH } from './utils';
6+
7+
test('Render works correctly', async () => {
8+
const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'vscode-test-sushi-'));
9+
await fs.copy(SUSHI_SOURCE_PATH, tempDir);
10+
11+
try {
12+
const { window, close } = await startVSCode(tempDir);
13+
14+
// Wait for the models folder to be visible
15+
await window.waitForSelector('text=models');
16+
17+
// Click on the models folder, excluding external_models
18+
await window.getByRole('treeitem', { name: 'models', exact: true }).locator('a').click();
19+
20+
// Open the customer_revenue_lifetime model
21+
await window.getByRole('treeitem', { name: 'customers.sql', exact: true }).locator('a').click();
22+
23+
await window.waitForSelector('text=grain');
24+
await window.waitForSelector('text=Loaded SQLMesh Context')
25+
26+
// Render the model
27+
await window.keyboard.press(process.platform === 'darwin' ? 'Meta+Shift+P' : 'Control+Shift+P');
28+
await window.keyboard.type('Render Model');
29+
await window.keyboard.press('Enter');
30+
31+
// Check if the model is rendered by check if "`oi`.`order_id` AS `order_id`," is in the window
32+
await expect(window.locator('text="marketing"."customer_id" AS')).toBeVisible();
33+
34+
await close();
35+
} finally {
36+
await fs.remove(tempDir);
37+
}
38+
});

0 commit comments

Comments
 (0)