Skip to content

fix(lsp): return context error in case it failed #4787

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
benfdking merged 1 commit into from
Jun 23, 2025
Merged

fix(lsp): return context error in case it failed #4787

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
141 changes: 106 additions & 35 deletions sqlmesh/lsp/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,33 @@
from sqlmesh.lsp.uri import URI
from web.server.api.endpoints.lineage import column_lineage, model_lineage
from web.server.api.endpoints.models import get_models
from typing import Union
from dataclasses import dataclass


@dataclass
class NoContext:
"""State when no context has been attempted to load."""

pass


@dataclass
class ContextLoaded:
"""State when context has been successfully loaded."""

lsp_context: LSPContext


@dataclass
class ContextFailed:
"""State when context failed to load with an error message."""

error_message: str
context: t.Optional[Context] = None


ContextState = Union[NoContext, ContextLoaded, ContextFailed]
Copy link
Copy Markdown
Collaborator

@themisvaltinos themisvaltinos Jun 23, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

im wondering whether this is the best approach since now since only ContextLoaded has the lsp_context we would need to have either type guards or assertions in quite a lot of places



class SQLMeshLanguageServer:
Expand All @@ -72,7 +99,7 @@ def __init__(
"""
self.server = LanguageServer(server_name, version)
self.context_class = context_class
self.lsp_context: t.Optional[LSPContext] = None
self.context_state: ContextState = NoContext()
self.workspace_folders: t.List[Path] = []

self.has_raised_loading_error: bool = False
Expand Down Expand Up @@ -257,16 +284,53 @@ def did_open(ls: LanguageServer, params: types.DidOpenTextDocumentParams) -> Non
@self.server.feature(types.TEXT_DOCUMENT_DID_SAVE)
def did_save(ls: LanguageServer, params: types.DidSaveTextDocumentParams) -> None:
uri = URI(params.text_document.uri)
if self.lsp_context is None:
if isinstance(self.context_state, NoContext):
return

context = self.lsp_context.context
context.load()
self.lsp_context = LSPContext(context)
if isinstance(self.context_state, ContextFailed):
if self.context_state.context:
try:
self.context_state.context.load()
self.context_state = ContextLoaded(
lsp_context=LSPContext(self.context_state.context)
)
except Exception as e:
ls.log_trace(f"Error loading context: {e}")
if not isinstance(self.context_state, ContextFailed):
raise Exception("Context state should be failed")
self.context_state = ContextFailed(
error_message=str(e), context=self.context_state.context
)
return
else:
# If there's no context, try to create one from scratch
try:
self._ensure_context_for_document(uri)
# If successful, context_state will be ContextLoaded
if isinstance(self.context_state, ContextLoaded):
ls.show_message(
"Successfully loaded SQLMesh context",
types.MessageType.Info,
)
except Exception as e:
ls.log_trace(f"Still cannot load context: {e}")
return

# Reload the context if was successfully
try:
context = self.context_state.lsp_context.context
context.load()
self.context_state = ContextLoaded(lsp_context=LSPContext(context))
except Exception as e:
ls.log_trace(f"Error loading context: {e}")
self.context_state = ContextFailed(
error_message=str(e), context=self.context_state.lsp_context.context
)
return

# Only publish diagnostics if client doesn't support pull diagnostics
if not self.client_supports_pull_diagnostics:
diagnostics = self.lsp_context.lint_model(uri)
diagnostics = self.context_state.lsp_context.lint_model(uri)
ls.publish_diagnostics(
params.text_document.uri,
SQLMeshLanguageServer._diagnostics_to_lsp_diagnostics(diagnostics),
Expand Down Expand Up @@ -443,11 +507,8 @@ def prepare_rename_handler(
"""Prepare for rename operation by checking if the symbol can be renamed."""
try:
uri = URI(params.text_document.uri)
self._ensure_context_for_document(uri)
if self.lsp_context is None:
raise RuntimeError(f"No context found for document: {uri}")

result = prepare_rename(self.lsp_context, uri, params.position)
context = self._context_get_or_load(uri)
result = prepare_rename(context, uri, params.position)
return result
except Exception as e:
ls.log_trace(f"Error preparing rename: {e}")
Expand All @@ -460,13 +521,8 @@ def rename_handler(
"""Perform rename operation on the symbol at the given position."""
try:
uri = URI(params.text_document.uri)
self._ensure_context_for_document(uri)
if self.lsp_context is None:
raise RuntimeError(f"No context found for document: {uri}")

workspace_edit = rename_symbol(
self.lsp_context, uri, params.position, params.new_name
)
context = self._context_get_or_load(uri)
workspace_edit = rename_symbol(context, uri, params.position, params.new_name)
return workspace_edit
except Exception as e:
ls.show_message(f"Error performing rename: {e}", types.MessageType.Error)
Expand All @@ -479,11 +535,8 @@ def document_highlight_handler(
"""Highlight all occurrences of the symbol at the given position."""
try:
uri = URI(params.text_document.uri)
self._ensure_context_for_document(uri)
if self.lsp_context is None:
raise RuntimeError(f"No context found for document: {uri}")

highlights = get_document_highlights(self.lsp_context, uri, params.position)
context = self._context_get_or_load(uri)
highlights = get_document_highlights(context, uri, params.position)
return highlights
except Exception as e:
ls.log_trace(f"Error getting document highlights: {e}")
Expand Down Expand Up @@ -670,11 +723,13 @@ def _get_diagnostics_for_uri(self, uri: URI) -> t.Tuple[t.List[types.Diagnostic]
return [], 0

def _context_get_or_load(self, document_uri: t.Optional[URI] = None) -> LSPContext:
if self.lsp_context is None:
if isinstance(self.context_state, ContextFailed):
raise RuntimeError(self.context_state.error_message)
if isinstance(self.context_state, NoContext):
self._ensure_context_for_document(document_uri)
if self.lsp_context is None:
raise RuntimeError("No context found able to get or load")
return self.lsp_context
if not isinstance(self.context_state, ContextLoaded):
raise RuntimeError("Context is not loaded")
return self.context_state.lsp_context

def _ensure_context_for_document(
self,
Expand All @@ -692,10 +747,10 @@ def _ensure_context_for_document(
self._ensure_context_in_folder(document_folder)
return

return self._ensure_context_in_folder()
self._ensure_context_in_folder()

def _ensure_context_in_folder(self, folder_path: t.Optional[Path] = None) -> None:
if self.lsp_context is not None:
if not isinstance(self.context_state, NoContext):
return

# If not found in the provided folder, search through all workspace folders
Expand Down Expand Up @@ -729,7 +784,7 @@ def _ensure_context_in_folder(self, folder_path: t.Optional[Path] = None) -> Non
def _create_lsp_context(self, paths: t.List[Path]) -> t.Optional[LSPContext]:
"""Create a new LSPContext instance using the configured context class.

On success, sets self.lsp_context and returns the created context.
On success, sets self.context_state to ContextLoaded and returns the created context.

Args:
paths: List of paths to pass to the context constructor
Expand All @@ -738,14 +793,22 @@ def _create_lsp_context(self, paths: t.List[Path]) -> t.Optional[LSPContext]:
A new LSPContext instance wrapping the created context, or None if creation fails
"""
try:
if self.lsp_context is None:
if isinstance(self.context_state, NoContext):
context = self.context_class(paths=paths)
loaded_sqlmesh_message(self.server, paths[0])
elif isinstance(self.context_state, ContextFailed):
if self.context_state.context:
context = self.context_state.context
context.load()
else:
# If there's no context (initial creation failed), try creating again
context = self.context_class(paths=paths)
loaded_sqlmesh_message(self.server, paths[0])
else:
self.lsp_context.context.load()
context = self.lsp_context.context
self.lsp_context = LSPContext(context)
return self.lsp_context
context = self.context_state.lsp_context.context
context.load()
self.context_state = ContextLoaded(lsp_context=LSPContext(context))
return self.context_state.lsp_context
except Exception as e:
# Only show the error message once
if not self.has_raised_loading_error:
Expand All @@ -756,6 +819,14 @@ def _create_lsp_context(self, paths: t.List[Path]) -> t.Optional[LSPContext]:
self.has_raised_loading_error = True

self.server.log_trace(f"Error creating context: {e}")
# Store the error in context state so subsequent requests show the actual error
# Try to preserve any partially loaded context if it exists
context = None
if isinstance(self.context_state, ContextLoaded):
context = self.context_state.lsp_context.context
elif isinstance(self.context_state, ContextFailed) and self.context_state.context:
context = self.context_state.context
self.context_state = ContextFailed(error_message=str(e), context=context)
return None

@staticmethod
Expand Down
64 changes: 62 additions & 2 deletions vscode/extension/tests/broken_project.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,66 @@ test('bad project, double model', async ({}) => {
}
})

test('working project, then broken through adding double model, then refixed', async ({}) => {
const tempDir = await fs.mkdtemp(
path.join(os.tmpdir(), 'vscode-test-tcloud-'),
)
await fs.copy(SUSHI_SOURCE_PATH, tempDir)

const { window, close } = await startVSCode(tempDir)
try {
// First, verify the project is working correctly
await window.waitForSelector('text=models')

// Open the lineage view to confirm it loads properly
await openLineageView(window)
await window.waitForSelector('text=Loaded SQLMesh context')

// Read the customers.sql file
const customersSql = await fs.readFile(
path.join(tempDir, 'models', 'customers.sql'),
'utf8',
)

// Add a duplicate model to break the project
await fs.writeFile(
path.join(tempDir, 'models', 'customers_duplicated.sql'),
customersSql,
)

// Open the customers model to trigger the error
await window
.getByRole('treeitem', { name: 'models', exact: true })
.locator('a')
.click()
await window
.getByRole('treeitem', { name: 'customers.sql', exact: true })
.locator('a')
.click()
// Save to refresh the context
await window.keyboard.press('Control+S')
await window.keyboard.press('Meta+S')

// Wait for the error to appear
// TODO: Selector doesn't work in the linage view
// await window.waitForSelector('text=Error')

// Remove the duplicated model to fix the project
await fs.remove(path.join(tempDir, 'models', 'customers_duplicated.sql'))

// Save again to refresh the context
await window.keyboard.press('Control+S')
await window.keyboard.press('Meta+S')

// Wait for the error to go away and context to reload
// TODO: Selector doesn't work in the linage view
// await window.waitForSelector('text=raw.demographics')
} finally {
await close()
await fs.remove(tempDir)
}
})

test('bad project, double model, then fixed', async ({}) => {
const tempDir = await fs.mkdtemp(
path.join(os.tmpdir(), 'vscode-test-tcloud-'),
Expand Down Expand Up @@ -86,7 +146,8 @@ test('bad project, double model, then fixed', async ({}) => {
await openLineageView(window)

// Wait for the error to go away
await window.waitForSelector('text=Loaded SQLMesh context')
// TODO: Selector doesn't work in the linage view
// await window.waitForSelector('text=raw.demographics')
} finally {
await close()
await fs.remove(tempDir)
Expand Down Expand Up @@ -119,7 +180,6 @@ test('bad project, double model, check lineage', async ({}) => {
await openLineageView(window)

await window.waitForSelector('text=Error creating context')

await window.waitForSelector('text=Error:')

await window.waitForTimeout(1000)
Expand Down