sqlmesh 0.205.1.dev2__tar.gz → 0.205.1.dev4__tar.gz

{sqlmesh-0.205.1.dev2 → sqlmesh-0.205.1.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.205.1.dev2
+Version: 0.205.1.dev4
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.205.1.dev2 → sqlmesh-0.205.1.dev4}/sqlmesh/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.205.1.dev2'
-__version_tuple__ = version_tuple = (0, 205, 1, 'dev2')
+__version__ = version = '0.205.1.dev4'
+__version_tuple__ = version_tuple = (0, 205, 1, 'dev4')

{sqlmesh-0.205.1.dev2 → sqlmesh-0.205.1.dev4}/sqlmesh/lsp/context.py

@@ -3,10 +3,12 @@ from pathlib import Path
 from sqlmesh.core.context import Context
 import typing as t
 
+from sqlmesh.core.linter.rule import Range
 from sqlmesh.core.model.definition import SqlModel
 from sqlmesh.core.linter.definition import AnnotatedRuleViolation
-from sqlmesh.lsp.custom import ModelForRendering
+from sqlmesh.lsp.custom import ModelForRendering, TestEntry, RunTestResponse
 from sqlmesh.lsp.custom import AllModelsResponse, RenderModelEntry
+from sqlmesh.lsp.tests_ranges import get_test_ranges
 from sqlmesh.lsp.uri import URI
 from lsprotocol import types
 
@@ -63,6 +65,71 @@ class LSPContext:
             **audit_map,
         }
 
+    def list_workspace_tests(self) -> t.List[TestEntry]:
+        """List all tests in the workspace."""
+        tests = self.context.load_model_tests()
+
+        # Use a set to ensure unique URIs
+        unique_test_uris = {URI.from_path(test.path).value for test in tests}
+        test_uris: t.Dict[str, t.Dict[str, Range]] = {}
+        for uri in unique_test_uris:
+            test_ranges = get_test_ranges(URI(uri).to_path())
+            if uri not in test_uris:
+                test_uris[uri] = {}
+            test_uris[uri].update(test_ranges)
+        return [
+            TestEntry(
+                name=test.test_name,
+                uri=URI.from_path(test.path).value,
+                range=test_uris.get(URI.from_path(test.path).value, {}).get(test.test_name),
+            )
+            for test in tests
+        ]
+
+    def get_document_tests(self, uri: URI) -> t.List[TestEntry]:
+        """Get tests for a specific document.
+
+        Args:
+            uri: The URI of the file to get tests for.
+
+        Returns:
+            List of TestEntry objects for the specified document.
+        """
+        tests = self.context.load_model_tests(tests=[str(uri.to_path())])
+        test_ranges = get_test_ranges(uri.to_path())
+        return [
+            TestEntry(
+                name=test.test_name,
+                uri=URI.from_path(test.path).value,
+                range=test_ranges.get(test.test_name),
+            )
+            for test in tests
+        ]
+
+    def run_test(self, uri: URI, test_name: str) -> RunTestResponse:
+        """Run a specific test for a model.
+
+        Args:
+            uri: The URI of the file containing the test.
+            test_name: The name of the test to run.
+
+        Returns:
+            List of annotated rule violations from the test run.
+        """
+        path = uri.to_path()
+        results = self.context.test(
+            tests=[str(path)],
+            match_patterns=[test_name],
+        )
+        if results.testsRun != 1:
+            raise ValueError(f"Expected to run 1 test, but ran {results.testsRun} tests.")
+        if len(results.successes) == 1:
+            return RunTestResponse(success=True)
+        return RunTestResponse(
+            success=False,
+            error_message=str(results.failures[0][1]),
+        )
+
     def render_model(self, uri: URI) -> t.List[RenderModelEntry]:
         """Get rendered models for a file, using cache when available.
 

{sqlmesh-0.205.1.dev2 → sqlmesh-0.205.1.dev4}/sqlmesh/lsp/custom.py

@@ -1,5 +1,7 @@
 from lsprotocol import types
 import typing as t
+
+from sqlmesh.core.linter.rule import Range
 from sqlmesh.utils.pydantic import PydanticModel
 
 
@@ -143,3 +145,52 @@ class FormatProjectResponse(CustomMethodResponseBaseClass):
     """
 
     pass
+
+
+LIST_WORKSPACE_TESTS_FEATURE = "sqlmesh/list_workspace_tests"
+
+
+class ListWorkspaceTestsRequest(CustomMethodRequestBaseClass):
+    """
+    Request to list all tests in the current project.
+    """
+
+    pass
+
+
+class TestEntry(PydanticModel):
+    """
+    An entry representing a test in the workspace.
+    """
+
+    name: str
+    uri: str
+    range: Range
+
+
+class ListWorkspaceTestsResponse(CustomMethodResponseBaseClass):
+    tests: t.List[TestEntry]
+
+
+LIST_DOCUMENT_TESTS_FEATURE = "sqlmesh/list_document_tests"
+
+
+class ListDocumentTestsRequest(CustomMethodRequestBaseClass):
+    textDocument: types.TextDocumentIdentifier
+
+
+class ListDocumentTestsResponse(CustomMethodResponseBaseClass):
+    tests: t.List[TestEntry]
+
+
+RUN_TEST_FEATURE = "sqlmesh/run_test"
+
+
+class RunTestRequest(CustomMethodRequestBaseClass):
+    textDocument: types.TextDocumentIdentifier
+    testName: str
+
+
+class RunTestResponse(CustomMethodResponseBaseClass):
+    success: bool
+    error_message: t.Optional[str] = None

{sqlmesh-0.205.1.dev2 → sqlmesh-0.205.1.dev4}/sqlmesh/lsp/main.py

@@ -46,6 +46,15 @@ from sqlmesh.lsp.custom import (
     FormatProjectRequest,
     FormatProjectResponse,
     CustomMethod,
+    LIST_WORKSPACE_TESTS_FEATURE,
+    ListWorkspaceTestsRequest,
+    ListWorkspaceTestsResponse,
+    LIST_DOCUMENT_TESTS_FEATURE,
+    ListDocumentTestsRequest,
+    ListDocumentTestsResponse,
+    RUN_TEST_FEATURE,
+    RunTestRequest,
+    RunTestResponse,
 )
 from sqlmesh.lsp.errors import ContextFailedError, context_error_to_diagnostic
 from sqlmesh.lsp.helpers import to_lsp_range, to_sqlmesh_position
@@ -127,11 +136,58 @@ class SQLMeshLanguageServer:
             API_FEATURE: self._custom_api,
             SUPPORTED_METHODS_FEATURE: self._custom_supported_methods,
             FORMAT_PROJECT_FEATURE: self._custom_format_project,
+            LIST_WORKSPACE_TESTS_FEATURE: self._list_workspace_tests,
+            LIST_DOCUMENT_TESTS_FEATURE: self._list_document_tests,
+            RUN_TEST_FEATURE: self._run_test,
         }
 
         # Register LSP features (e.g., formatting, hover, etc.)
         self._register_features()
 
+    def _list_workspace_tests(
+        self,
+        ls: LanguageServer,
+        params: ListWorkspaceTestsRequest,
+    ) -> ListWorkspaceTestsResponse:
+        """List all tests in the current workspace."""
+        try:
+            context = self._context_get_or_load()
+            tests = context.list_workspace_tests()
+            return ListWorkspaceTestsResponse(tests=tests)
+        except Exception as e:
+            ls.log_trace(f"Error listing workspace tests: {e}")
+            return ListWorkspaceTestsResponse(tests=[])
+
+    def _list_document_tests(
+        self,
+        ls: LanguageServer,
+        params: ListDocumentTestsRequest,
+    ) -> ListDocumentTestsResponse:
+        """List tests for a specific document."""
+        try:
+            uri = URI(params.textDocument.uri)
+            context = self._context_get_or_load(uri)
+            tests = context.get_document_tests(uri)
+            return ListDocumentTestsResponse(tests=tests)
+        except Exception as e:
+            ls.log_trace(f"Error listing document tests: {e}")
+            return ListDocumentTestsResponse(tests=[])
+
+    def _run_test(
+        self,
+        ls: LanguageServer,
+        params: RunTestRequest,
+    ) -> RunTestResponse:
+        """Run a specific test."""
+        try:
+            uri = URI(params.textDocument.uri)
+            context = self._context_get_or_load(uri)
+            result = context.run_test(uri, params.testName)
+            return result
+        except Exception as e:
+            ls.log_trace(f"Error running test: {e}")
+            return RunTestResponse(success=False, response_error=str(e))
+
     # All the custom LSP methods are registered here and prefixed with _custom
     def _custom_all_models(self, ls: LanguageServer, params: AllModelsRequest) -> AllModelsResponse:
         uri = URI(params.textDocument.uri)

sqlmesh-0.205.1.dev4/sqlmesh/lsp/tests_ranges.py

@@ -0,0 +1,65 @@
+"""
+Provides helper functions to get ranges of tests in SQLMesh LSP.
+"""
+
+from pathlib import Path
+
+from sqlmesh.core.linter.rule import Range, Position
+from ruamel import yaml
+from ruamel.yaml.comments import CommentedMap
+import typing as t
+
+
+def get_test_ranges(
+    path: Path,
+) -> t.Dict[str, Range]:
+    """
+    Test files are yaml files with a stucture of dict to test information. This returns a dictionary
+    with the test name as the key and the range of the test in the file as the value.
+    """
+    test_ranges: t.Dict[str, Range] = {}
+
+    with open(path, "r", encoding="utf-8") as file:
+        content = file.read()
+
+    # Parse YAML to get line numbers
+    yaml_obj = yaml.YAML()
+    yaml_obj.preserve_quotes = True
+    data = yaml_obj.load(content)
+
+    if not isinstance(data, dict):
+        raise ValueError("Invalid test file format: expected a dictionary at the top level.")
+
+    # For each top-level key (test name), find its range
+    for test_name in data:
+        if isinstance(data, CommentedMap) and test_name in data.lc.data:
+            # Get line and column info from ruamel yaml
+            line_info = data.lc.data[test_name]
+            start_line = line_info[0]  # 0-based line number
+            start_col = line_info[1]  # 0-based column number
+
+            # Find the end of this test by looking for the next test or end of file
+            lines = content.splitlines()
+            end_line = start_line
+
+            # Find where this test ends by looking for the next top-level key
+            # or the end of the file
+            for i in range(start_line + 1, len(lines)):
+                line = lines[i]
+                # Check if this line starts a new top-level key (no leading spaces)
+                if line and not line[0].isspace() and ":" in line:
+                    end_line = i - 1
+                    break
+            else:
+                # This test goes to the end of the file
+                end_line = len(lines) - 1
+
+            # Create the range
+            test_ranges[test_name] = Range(
+                start=Position(line=start_line, character=start_col),
+                end=Position(
+                    line=end_line, character=len(lines[end_line]) if end_line < len(lines) else 0
+                ),
+            )
+
+    return test_ranges

{sqlmesh-0.205.1.dev2 → sqlmesh-0.205.1.dev4}/sqlmesh.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlmesh
-Version: 0.205.1.dev2
+Version: 0.205.1.dev4
 Summary: Next-generation data transformation framework
 Author-email: "TobikoData Inc." <engineering@tobikodata.com>
 License:                                  Apache License

{sqlmesh-0.205.1.dev2 → sqlmesh-0.205.1.dev4}/sqlmesh.egg-info/SOURCES.txt

@@ -370,6 +370,7 @@ sqlmesh/lsp/hints.py
 sqlmesh/lsp/main.py
 sqlmesh/lsp/reference.py
 sqlmesh/lsp/rename.py
+sqlmesh/lsp/tests_ranges.py
 sqlmesh/lsp/uri.py
 sqlmesh/migrations/__init__.py
 sqlmesh/migrations/v0001_init.py
@@ -768,6 +769,7 @@ vscode/extension/src/completion/completion.ts
 vscode/extension/src/lsp/custom.ts
 vscode/extension/src/lsp/lsp.ts
 vscode/extension/src/providers/renderedModelProvider.ts
+vscode/extension/src/tests/tests.ts
 vscode/extension/src/utilities/config.ts
 vscode/extension/src/utilities/errors.ts
 vscode/extension/src/utilities/exec.ts
@@ -806,6 +808,7 @@ vscode/extension/tests/render.spec.ts
 vscode/extension/tests/stop.spec.ts
 vscode/extension/tests/tcloud.spec.ts
 vscode/extension/tests/tcloud_utils.ts
+vscode/extension/tests/tests.spec.ts
 vscode/extension/tests/utils.ts
 vscode/extension/tests/utils_code_server.ts
 vscode/extension/tests/venv_naming.spec.ts

sqlmesh-0.205.1.dev4/tests/lsp/test_context.py

@@ -0,0 +1,63 @@
+from pathlib import Path
+
+from sqlmesh.core.context import Context
+from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.uri import URI
+
+
+def test_lsp_context():
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    assert lsp_context is not None
+    assert lsp_context.context is not None
+    assert lsp_context.map is not None
+
+    # find one model in the map
+    active_customers_key = next(
+        key for key in lsp_context.map.keys() if key.name == "active_customers.sql"
+    )
+
+    # Check that the value is a ModelInfo with the expected model name
+    assert isinstance(lsp_context.map[active_customers_key], ModelTarget)
+    assert "sushi.active_customers" in lsp_context.map[active_customers_key].names
+
+
+def test_lsp_context_list_workspace_tests():
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    # List workspace tests
+    tests = lsp_context.list_workspace_tests()
+
+    # Check that the tests are returned correctly
+    assert len(tests) == 3
+    assert any(test.name == "test_order_items" for test in tests)
+
+
+def test_lsp_context_get_document_tests():
+    test_path = Path.cwd() / "examples/sushi/tests/test_order_items.yaml"
+    uri = URI.from_path(test_path)
+
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+    tests = lsp_context.get_document_tests(uri)
+
+    assert len(tests) == 1
+    assert tests[0].uri == uri.value
+    assert tests[0].name == "test_order_items"
+
+
+def test_lsp_context_run_test():
+    test_path = Path.cwd() / "examples/sushi/tests/test_order_items.yaml"
+    uri = URI.from_path(test_path)
+
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    # Run the test
+    result = lsp_context.run_test(uri, "test_order_items")
+
+    # Check that the result is not None and has the expected properties
+    assert result is not None
+    assert result.success is True

{sqlmesh-0.205.1.dev2 → sqlmesh-0.205.1.dev4}/vscode/extension/src/commands/stop.ts

@@ -11,7 +11,7 @@ export const stop = (lspClient: LSPClient | undefined) => {
       return
     }
 
-    await lspClient.stop()
+    await lspClient.stop(true)
     await window.showInformationMessage('LSP server stopped')
     traceInfo('LSP server stopped successfully')
   }

{sqlmesh-0.205.1.dev2 → sqlmesh-0.205.1.dev4}/vscode/extension/src/extension.ts

@@ -21,6 +21,10 @@ import { selector, completionProvider } from './completion/completion'
 import { LineagePanel } from './webviews/lineagePanel'
 import { RenderedModelProvider } from './providers/renderedModelProvider'
 import { sleep } from './utilities/sleep'
+import {
+  controller as testController,
+  setupTestController,
+} from './tests/tests'
 
 let lspClient: LSPClient | undefined
 
@@ -115,10 +119,10 @@ export async function activate(context: vscode.ExtensionContext) {
     }),
   )
 
-  const restart = async () => {
+  const restart = async (invokedByUser = false) => {
     if (lspClient) {
       traceVerbose('Restarting LSP client')
-      const restartResult = await lspClient.restart()
+      const restartResult = await lspClient.restart(invokedByUser)
       if (isErr(restartResult)) {
         return handleError(
           authProvider,
@@ -128,9 +132,10 @@ export async function activate(context: vscode.ExtensionContext) {
         )
       }
       context.subscriptions.push(lspClient)
+      context.subscriptions.push(setupTestController(lspClient))
     } else {
       lspClient = new LSPClient()
-      const result = await lspClient.start()
+      const result = await lspClient.start(invokedByUser)
       if (isErr(result)) {
         return handleError(
           authProvider,
@@ -140,6 +145,7 @@ export async function activate(context: vscode.ExtensionContext) {
         )
       } else {
         context.subscriptions.push(lspClient)
+        context.subscriptions.push(setupTestController(lspClient))
       }
     }
   }
@@ -159,7 +165,7 @@ export async function activate(context: vscode.ExtensionContext) {
       await restart()
     }),
     registerCommand(`sqlmesh.restart`, async () => {
-      await restart()
+      await restart(true)
     }),
     registerCommand(`sqlmesh.stop`, stop(lspClient)),
     registerCommand(`sqlmesh.printEnvironment`, printEnvironment()),
@@ -175,6 +181,8 @@ export async function activate(context: vscode.ExtensionContext) {
     )
   } else {
     context.subscriptions.push(lspClient)
+    context.subscriptions.push(setupTestController(lspClient))
+    context.subscriptions.push(testController)
   }
 
   if (lspClient && !lspClient.hasCompletionCapability()) {

{sqlmesh-0.205.1.dev2 → sqlmesh-0.205.1.dev4}/vscode/extension/src/lsp/custom.ts

@@ -32,6 +32,9 @@ export type CustomLSPMethods =
   | AllModelsForRenderMethod
   | SupportedMethodsMethod
   | FormatProjectMethod
+  | ListWorkspaceTests
+  | ListDocumentTests
+  | RunTest
 
 interface AllModelsRequest {
   textDocument: {
@@ -111,3 +114,65 @@ interface FormatProjectResponse extends BaseResponse {}
 interface BaseResponse {
   response_error?: string
 }
+
+export interface ListWorkspaceTests {
+  method: 'sqlmesh/list_workspace_tests'
+  request: ListWorkspaceTestsRequest
+  response: ListWorkspaceTestsResponse
+}
+
+type ListWorkspaceTestsRequest = object
+
+interface Position {
+  line: number
+  character: number
+}
+
+interface Range {
+  start: Position
+  end: Position
+}
+
+interface TestEntry {
+  name: string
+  uri: string
+  range: Range
+}
+
+interface ListWorkspaceTestsResponse extends BaseResponse {
+  tests: TestEntry[]
+}
+
+export interface ListDocumentTests {
+  method: 'sqlmesh/list_document_tests'
+  request: ListDocumentTestsRequest
+  response: ListDocumentTestsResponse
+}
+
+export interface DocumentIdentifier {
+  uri: string
+}
+
+export interface ListDocumentTestsRequest {
+  textDocument: DocumentIdentifier
+}
+
+export interface ListDocumentTestsResponse extends BaseResponse {
+  tests: TestEntry[]
+}
+
+export interface RunTest {
+  method: 'sqlmesh/run_test'
+  request: RunTestRequest
+  response: RunTestResponse
+}
+
+export interface RunTestRequest {
+  textDocument: DocumentIdentifier
+  testName: string
+}
+
+export interface RunTestResponse extends BaseResponse {
+  success: boolean
+  error_message?: string
+}

{sqlmesh-0.205.1.dev2 → sqlmesh-0.205.1.dev4}/vscode/extension/src/lsp/lsp.ts

@@ -36,6 +36,13 @@ export class LSPClient implements Disposable {
    */
   private supportedMethodsState: SupportedMethodsState = { type: 'not-fetched' }
 
+  /**
+   * Explicitly stopped remembers whether the LSP client has been explicitly stopped
+   * by the user. This is used to prevent the client from being restarted unless the user
+   * explicitly calls the `restart` method.
+   */
+  private explicitlyStopped = false
+
   constructor() {
     this.client = undefined
   }
@@ -54,7 +61,15 @@ export class LSPClient implements Disposable {
     return completion !== undefined
   }
 
-  public async start(): Promise<Result<undefined, ErrorType>> {
+  public async start(
+    overrideStoppedByUser = false,
+  ): Promise<Result<undefined, ErrorType>> {
+    if (this.explicitlyStopped && !overrideStoppedByUser) {
+      traceInfo(
+        'LSP client has been explicitly stopped by user, not starting again.',
+      )
+      return ok(undefined)
+    }
     if (!outputChannel) {
       outputChannel = window.createOutputChannel('sqlmesh-lsp')
     }
@@ -124,18 +139,24 @@ export class LSPClient implements Disposable {
     return ok(undefined)
   }
 
-  public async restart(): Promise<Result<undefined, ErrorType>> {
+  public async restart(
+    overrideByUser = false,
+  ): Promise<Result<undefined, ErrorType>> {
     await this.stop()
-    return await this.start()
+    return await this.start(overrideByUser)
   }
 
-  public async stop() {
+  public async stop(stoppedByUser = false): Promise<void> {
     if (this.client) {
       await this.client.stop()
       this.client = undefined
       // Reset supported methods state when the client stops
       this.supportedMethodsState = { type: 'not-fetched' }
     }
+    if (stoppedByUser) {
+      this.explicitlyStopped = true
+      traceInfo('SQLMesh LSP client stopped by user.')
+    }
   }
 
   public async dispose() {