From 31cd1bdd5278f346b962b254c483e982a7ad48ba Mon Sep 17 00:00:00 2001 From: Benoit Perigaud <8754100+b-per@users.noreply.github.com> Date: Fri, 31 Oct 2025 14:23:00 +0100 Subject: [PATCH 1/3] Add tool to rename model using the LSP server --- ...cement or New Feature-20251031-141654.yaml | 3 + src/dbt_mcp/lsp/lsp_client.py | 248 ++++++++++++++++++ src/dbt_mcp/lsp/tools.py | 60 +++++ src/dbt_mcp/prompts/lsp/args/new_uri.md | 4 + src/dbt_mcp/prompts/lsp/args/old_uri.md | 4 + src/dbt_mcp/prompts/lsp/rename_model.md | 19 ++ src/dbt_mcp/tools/tool_names.py | 1 + tests/unit/lsp/test_lsp_client.py | 167 +++++++++++- tests/unit/lsp/test_lsp_tools.py | 104 ++++++++ 9 files changed, 609 insertions(+), 1 deletion(-) create mode 100644 .changes/unreleased/Enhancement or New Feature-20251031-141654.yaml create mode 100644 src/dbt_mcp/prompts/lsp/args/new_uri.md create mode 100644 src/dbt_mcp/prompts/lsp/args/old_uri.md create mode 100644 src/dbt_mcp/prompts/lsp/rename_model.md diff --git a/.changes/unreleased/Enhancement or New Feature-20251031-141654.yaml b/.changes/unreleased/Enhancement or New Feature-20251031-141654.yaml new file mode 100644 index 00000000..6021068e --- /dev/null +++ b/.changes/unreleased/Enhancement or New Feature-20251031-141654.yaml @@ -0,0 +1,3 @@ +kind: Enhancement or New Feature +body: Add ability to call the LSP to rename models +time: 2025-10-31T14:16:54.720743+01:00 diff --git a/src/dbt_mcp/lsp/lsp_client.py b/src/dbt_mcp/lsp/lsp_client.py index 16ef7d30..0d454ee1 100644 --- a/src/dbt_mcp/lsp/lsp_client.py +++ b/src/dbt_mcp/lsp/lsp_client.py @@ -7,7 +7,9 @@ import asyncio import logging +from pathlib import Path from typing import Any +from urllib.parse import unquote, urlparse from dbt_mcp.lsp.lsp_connection import LSPConnection, LspEventName @@ -148,3 +150,249 @@ async def _list_nodes( return {"nodes": result["nodes"]} return result + + def _uri_to_path(self, uri: str) -> Path: + """Convert a file:// URI to a Path object.""" + parsed = urlparse(uri) + # Decode percent-encoded characters and remove the leading slash on Windows + path_str = unquote(parsed.path) + return Path(path_str) + + def _apply_single_edit( + self, lines: list[str], edit: dict[str, Any] + ) -> None: + """Apply a single text edit to a list of lines. + + Args: + lines: List of lines (with line endings preserved) + edit: LSP TextEdit object with 'range' and 'newText' + """ + start_line_idx = edit["range"]["start"]["line"] + start_char = edit["range"]["start"]["character"] + end_line_idx = edit["range"]["end"]["line"] + end_char = edit["range"]["end"]["character"] + new_text = edit["newText"] + + # Handle single line edit + if start_line_idx == end_line_idx: + if start_line_idx < len(lines): + line = lines[start_line_idx] + # Strip line ending to work with just the text + line_ending = '' + if line.endswith('\r\n'): + line_ending = '\r\n' + line = line[:-2] + elif line.endswith('\n'): + line_ending = '\n' + line = line[:-1] + elif line.endswith('\r'): + line_ending = '\r' + line = line[:-1] + + # Apply the edit + lines[start_line_idx] = ( + line[:start_char] + new_text + line[end_char:] + line_ending + ) + else: + # Multi-line edit + start_line = lines[start_line_idx] + end_line = lines[end_line_idx] + + # Strip line endings + start_line_ending = '' + if start_line.endswith('\r\n'): + start_line_ending = '\r\n' + start_line = start_line[:-2] + elif start_line.endswith('\n'): + start_line_ending = '\n' + start_line = start_line[:-1] + elif start_line.endswith('\r'): + start_line_ending = '\r' + start_line = start_line[:-1] + + end_line_ending = '' + if end_line.endswith('\r\n'): + end_line_ending = '\r\n' + end_line = end_line[:-2] + elif end_line.endswith('\n'): + end_line_ending = '\n' + end_line = end_line[:-1] + elif end_line.endswith('\r'): + end_line_ending = '\r' + end_line = end_line[:-1] + + start_line_text = start_line[:start_char] + end_line_text = end_line[end_char:] + lines[start_line_idx] = start_line_text + new_text + end_line_text + end_line_ending + # Remove the lines in between + del lines[start_line_idx + 1:end_line_idx + 1] + + def _apply_text_edits(self, file_path: Path, edits: list[dict[str, Any]]) -> None: + """Apply text edits to a file. + + Args: + file_path: Path to the file to edit + edits: List of LSP TextEdit objects with 'range' and 'newText' + """ + if not file_path.exists(): + logger.warning(f"File not found for edits: {file_path}") + return + + # Read the file content + content = file_path.read_text() + + # Sort edits in reverse order (end to start) to avoid offset issues + # We sort by start position descending so we can apply from end to beginning + sorted_edits = sorted( + edits, + key=lambda e: ( + e["range"]["start"]["line"], + e["range"]["start"]["character"] + ), + reverse=True + ) + + # Convert content to list of lines for easier manipulation + lines = content.splitlines(keepends=True) + + # Apply each edit + for edit in sorted_edits: + self._apply_single_edit(lines, edit) + + # Write back to file + file_path.write_text("".join(lines)) + logger.info(f"Applied {len(edits)} edits to {file_path}") + + async def rename_model( + self, + old_uri: str, + new_uri: str, + apply_edits: bool = True, + timeout: float | None = None, + ) -> dict[str, Any]: + """Rename a dbt model file and update all references. + + This method: + 1. Asks the LSP server what edits are needed for the rename + 2. Applies those edits to update references in other files + 3. Performs the actual file rename on disk + 4. Notifies the LSP server that the rename is complete + + Args: + old_uri: The current file URI (e.g., "file:///path/to/model.sql") + new_uri: The new file URI (e.g., "file:///path/to/new_model.sql") + apply_edits: Whether to apply workspace edits and perform the rename (default: True) + timeout: Optional timeout for the request + + Returns: + Dictionary with: + - 'renamed': True if model was renamed + - 'old_path': Original file path + - 'new_path': New file path + - 'files_updated': List of files that had references updated + - 'error': Error message if something failed + """ + logger.info(f"Renaming model: {old_uri} -> {new_uri}") + + # Step 1: Ask LSP what edits are needed + params = { + "files": [ + { + "oldUri": old_uri, + "newUri": new_uri, + } + ] + } + + try: + async with asyncio.timeout(timeout or self.timeout): + result = await self.lsp_connection.send_request( + "workspace/willRenameFiles", + params, + ) + + # Handle None or empty result + if result is None: + result = {} + + if "error" in result and result["error"] is not None: + return {"error": result["error"]} + + if not apply_edits: + # Just return the workspace edits without applying them + return result + + # Step 2: Apply workspace edits + files_updated = [] + if result and "changes" in result: + # Handle WorkspaceEdit with 'changes' format + for file_uri, edits in result["changes"].items(): + try: + file_path = self._uri_to_path(file_uri) + self._apply_text_edits(file_path, edits) + files_updated.append(str(file_path)) + except Exception as e: + logger.error(f"Failed to apply edits to {file_uri}: {e}") + return { + "error": f"Failed to apply edits to {file_uri}: {str(e)}" + } + elif result and "documentChanges" in result: + # Handle WorkspaceEdit with 'documentChanges' format + for change in result["documentChanges"]: + if "textDocument" in change and "edits" in change: + file_uri = change["textDocument"]["uri"] + try: + file_path = self._uri_to_path(file_uri) + self._apply_text_edits(file_path, change["edits"]) + files_updated.append(str(file_path)) + except Exception as e: + logger.error(f"Failed to apply edits to {file_uri}: {e}") + return { + "error": f"Failed to apply edits to {file_uri}: {str(e)}" + } + + # Step 3: Perform the actual file rename + old_path = self._uri_to_path(old_uri) + new_path = self._uri_to_path(new_uri) + + if not old_path.exists(): + return {"error": f"Source file does not exist: {old_path}"} + + if new_path.exists(): + return {"error": f"Destination file already exists: {new_path}"} + + try: + # Ensure parent directory exists + new_path.parent.mkdir(parents=True, exist_ok=True) + old_path.rename(new_path) + logger.info(f"Renamed file: {old_path} -> {new_path}") + except Exception as e: + return {"error": f"Failed to rename file: {str(e)}"} + + # Step 4: Notify LSP of completion (didRenameFiles) + try: + self.lsp_connection.send_notification( + "workspace/didRenameFiles", + { + "files": [ + { + "oldUri": old_uri, + "newUri": new_uri, + } + ] + }, + ) + except Exception as e: + logger.warning(f"Failed to send didRenameFiles notification: {e}") + + return { + "renamed": True, + "old_path": str(old_path), + "new_path": str(new_path), + "files_updated": files_updated, + } + + except asyncio.TimeoutError: + return {"error": "Timeout waiting for LSP response"} + except Exception as e: + return {"error": f"Failed to rename file: {str(e)}"} diff --git a/src/dbt_mcp/lsp/tools.py b/src/dbt_mcp/lsp/tools.py index d9c7472a..9d09987b 100644 --- a/src/dbt_mcp/lsp/tools.py +++ b/src/dbt_mcp/lsp/tools.py @@ -111,6 +111,16 @@ async def wrapper(*args, **kwargs) -> Any: idempotent_hint=True, ), ), + ToolDefinition( + fn=call_with_lsp_client(rename_model), + description=get_prompt("lsp/rename_model"), + annotations=create_tool_annotations( + title="rename_model", + read_only_hint=False, + destructive_hint=True, # Modifies files on disk + idempotent_hint=False, + ), + ), ] @@ -163,6 +173,56 @@ async def get_column_lineage( return {"error": error_msg} +async def rename_model( + lsp_client: LSPClient, + old_uri: str = Field(description=get_prompt("lsp/args/old_uri")), + new_uri: str = Field(description=get_prompt("lsp/args/new_uri")), +) -> dict[str, Any]: + """Rename a dbt model in the project workspace and update all references. + + Args: + lsp_client: The LSP client instance + old_uri: The current file URI + new_uri: The new file URI + + Returns: + Dictionary with either: + - 'renamed': True, 'old_path', 'new_path', 'files_updated' on success + - 'error' key containing error message on failure + """ + try: + response = await lsp_client.rename_model( + old_uri=old_uri, + new_uri=new_uri, + apply_edits=True, + ) + + # Check for LSP-level errors + if "error" in response: + logger.error(f"Error renaming model: {response['error']}") + return {"error": response["error"]} + + if response.get("renamed"): + logger.info( + f"Successfully renamed {response['old_path']} to {response['new_path']}" + ) + if response.get("files_updated"): + logger.info( + f"Updated {len(response['files_updated'])} files with new references" + ) + + return response + + except TimeoutError: + error_msg = f"Timeout waiting for model rename (old: {old_uri}, new: {new_uri})" + logger.error(error_msg) + return {"error": error_msg} + except Exception as e: + error_msg = f"Failed to rename model from {old_uri} to {new_uri}: {str(e)}" + logger.error(error_msg) + return {"error": error_msg} + + async def cleanup_lsp_connection() -> None: """Clean up the LSP connection when shutting down.""" global _lsp_connection diff --git a/src/dbt_mcp/prompts/lsp/args/new_uri.md b/src/dbt_mcp/prompts/lsp/args/new_uri.md new file mode 100644 index 00000000..caf67be6 --- /dev/null +++ b/src/dbt_mcp/prompts/lsp/args/new_uri.md @@ -0,0 +1,4 @@ +The new model file URI (e.g., "file:///path/to/project/models/staging/stg_customers_new.sql"). + +This should be the absolute path to the new model file location using the file:// URI scheme. + diff --git a/src/dbt_mcp/prompts/lsp/args/old_uri.md b/src/dbt_mcp/prompts/lsp/args/old_uri.md new file mode 100644 index 00000000..3c690d25 --- /dev/null +++ b/src/dbt_mcp/prompts/lsp/args/old_uri.md @@ -0,0 +1,4 @@ +The current model file URI (e.g., "file:///path/to/project/models/staging/stg_customers.sql"). + +This should be the absolute path to the model file using the file:// URI scheme. + diff --git a/src/dbt_mcp/prompts/lsp/rename_model.md b/src/dbt_mcp/prompts/lsp/rename_model.md new file mode 100644 index 00000000..5357485d --- /dev/null +++ b/src/dbt_mcp/prompts/lsp/rename_model.md @@ -0,0 +1,19 @@ +Rename a dbt model in the project workspace and automatically update all references. + +This tool performs a complete model rename operation with the help of the LSP server: + +1. **Queries the LSP server** for all files that reference the model being renamed +2. **Updates references** in all affected files (e.g., updates model names in `ref()` calls) +3. **Renames the model file** on disk to the new location +4. **Notifies the LSP server** that the rename is complete + +The tool handles: +- Renaming model files (.sql files in your models directory) +- Updating all `ref('model_name')` references throughout the project +- Updating imports and dependencies in other models +- Preserving file encoding and line endings + +**Important**: This tool makes actual changes to your files. Make sure you have committed any unsaved work before using it. + +Use this when you need to safely rename a dbt model while ensuring all references to it are updated correctly throughout your dbt project. + diff --git a/src/dbt_mcp/tools/tool_names.py b/src/dbt_mcp/tools/tool_names.py index ccebc306..0b28fe39 100644 --- a/src/dbt_mcp/tools/tool_names.py +++ b/src/dbt_mcp/tools/tool_names.py @@ -60,6 +60,7 @@ class ToolName(Enum): # dbt LSP tools GET_COLUMN_LINEAGE = "get_column_lineage" + RENAME_MODEL = "rename_model" @classmethod def get_all_tool_names(cls) -> set[str]: diff --git a/tests/unit/lsp/test_lsp_client.py b/tests/unit/lsp/test_lsp_client.py index d02b35c6..691be518 100644 --- a/tests/unit/lsp/test_lsp_client.py +++ b/tests/unit/lsp/test_lsp_client.py @@ -1,6 +1,7 @@ """Tests for the DbtLspClient class.""" -from unittest.mock import MagicMock +from pathlib import Path +from unittest.mock import MagicMock, Mock, patch import pytest @@ -88,3 +89,167 @@ async def test_get_column_lineage_error(lsp_client, mock_lsp_connection): ) assert result == {"error": "LSP server error"} + + +@pytest.mark.asyncio +async def test_rename_model_success_with_apply_edits( + lsp_client, mock_lsp_connection, tmp_path +): + """Test successful model rename with edits applied.""" + # Create temporary files + old_file = tmp_path / "old_file.sql" + new_file = tmp_path / "new_file.sql" + ref_file = tmp_path / "ref_file.sql" + + old_file.write_text("SELECT * FROM table") + ref_file.write_text("SELECT * FROM {{ ref('old_file') }}") + + # Setup mock LSP response + mock_result = { + "changes": { + f"file://{ref_file}": [ + { + "range": { + "start": {"line": 0, "character": 22}, + "end": {"line": 0, "character": 30}, + }, + "newText": "new_file", + } + ] + } + } + mock_lsp_connection.send_request.return_value = mock_result + + # Execute + old_uri = f"file://{old_file}" + new_uri = f"file://{new_file}" + result = await lsp_client.rename_model(old_uri=old_uri, new_uri=new_uri) + + # Verify + assert result["renamed"] is True + assert result["old_path"] == str(old_file) + assert result["new_path"] == str(new_file) + assert str(ref_file) in result["files_updated"] + + # Verify file was renamed + assert not old_file.exists() + assert new_file.exists() + + # Verify reference was updated + assert ref_file.read_text() == "SELECT * FROM {{ ref('new_file') }}" + + # Verify didRenameFiles notification was sent + mock_lsp_connection.send_notification.assert_called_once() + + +@pytest.mark.asyncio +async def test_rename_model_without_apply_edits(lsp_client, mock_lsp_connection): + """Test model rename request without applying edits (dry-run).""" + # Setup mock + mock_result = { + "changes": { + "file:///path/to/other_file.sql": [ + { + "range": { + "start": {"line": 0, "character": 0}, + "end": {"line": 0, "character": 8}, + }, + "newText": "new_file", + } + ] + } + } + mock_lsp_connection.send_request.return_value = mock_result + + # Execute without applying edits + old_uri = "file:///path/to/old_model.sql" + new_uri = "file:///path/to/new_model.sql" + result = await lsp_client.rename_model( + old_uri=old_uri, new_uri=new_uri, apply_edits=False + ) + + # Verify - should just return the workspace edits + assert result == mock_result + assert "changes" in result + mock_lsp_connection.send_notification.assert_not_called() + + +@pytest.mark.asyncio +async def test_rename_model_error(lsp_client, mock_lsp_connection): + """Test model rename request with LSP error.""" + # Setup mock to return an error + mock_lsp_connection.send_request.return_value = {"error": "Model not found"} + + # Execute + result = await lsp_client.rename_model( + old_uri="file:///path/to/old_model.sql", + new_uri="file:///path/to/new_model.sql", + ) + + # Verify + assert result == {"error": "Model not found"} + + +@pytest.mark.asyncio +async def test_rename_model_source_not_exists(lsp_client, mock_lsp_connection, tmp_path): + """Test model rename when source file doesn't exist.""" + # Setup mock + mock_lsp_connection.send_request.return_value = {"changes": {}} + + # Execute with non-existent file + old_uri = f"file://{tmp_path / 'nonexistent.sql'}" + new_uri = f"file://{tmp_path / 'new.sql'}" + result = await lsp_client.rename_model(old_uri=old_uri, new_uri=new_uri) + + # Verify + assert "error" in result + assert "does not exist" in result["error"] + + +@pytest.mark.asyncio +async def test_rename_model_destination_exists(lsp_client, mock_lsp_connection, tmp_path): + """Test model rename when destination file already exists.""" + # Create both files + old_file = tmp_path / "old.sql" + new_file = tmp_path / "new.sql" + old_file.write_text("content") + new_file.write_text("existing") + + # Setup mock + mock_lsp_connection.send_request.return_value = {"changes": {}} + + # Execute + old_uri = f"file://{old_file}" + new_uri = f"file://{new_file}" + result = await lsp_client.rename_model(old_uri=old_uri, new_uri=new_uri) + + # Verify + assert "error" in result + assert "already exists" in result["error"] + + +@pytest.mark.asyncio +async def test_rename_model_none_response(lsp_client, mock_lsp_connection, tmp_path): + """Test model rename when LSP returns None response.""" + # Create the source file + old_file = tmp_path / "old.sql" + new_file = tmp_path / "new.sql" + old_file.write_text("content") + + # Setup mock to return None (simulating the user's issue) + mock_lsp_connection.send_request.return_value = None + + # Execute - should not raise "argument of type 'NoneType' is not iterable" + old_uri = f"file://{old_file}" + new_uri = f"file://{new_file}" + result = await lsp_client.rename_model(old_uri=old_uri, new_uri=new_uri) + + # Verify - should succeed with empty files_updated + assert result["renamed"] is True + assert result["old_path"] == str(old_file) + assert result["new_path"] == str(new_file) + assert result["files_updated"] == [] + + # Verify file was renamed + assert not old_file.exists() + assert new_file.exists() diff --git a/tests/unit/lsp/test_lsp_tools.py b/tests/unit/lsp/test_lsp_tools.py index 69a5aa82..1a1bd587 100644 --- a/tests/unit/lsp/test_lsp_tools.py +++ b/tests/unit/lsp/test_lsp_tools.py @@ -9,6 +9,7 @@ cleanup_lsp_connection, get_column_lineage, register_lsp_tools, + rename_model, ) from dbt_mcp.lsp.lsp_client import LSPClient from dbt_mcp.mcp.server import FastMCP @@ -66,6 +67,7 @@ async def test_register_lsp_tools_success( # Verify correct tools were registered tool_names = [tool.name for tool in await test_mcp_server.list_tools()] assert ToolName.GET_COLUMN_LINEAGE.value in tool_names + assert ToolName.RENAME_MODEL.value in tool_names @pytest.mark.asyncio @@ -203,3 +205,105 @@ async def test_get_column_lineage_generic_exception() -> None: assert "error" in result assert "Failed to get column lineage" in result["error"] assert "Connection lost" in result["error"] + + +@pytest.mark.asyncio +async def test_rename_model_success() -> None: + """Test successful model rename.""" + mock_lsp_client = AsyncMock(spec=LSPClient) + mock_lsp_client.rename_model = AsyncMock( + return_value={ + "renamed": True, + "old_path": "/path/to/old_model.sql", + "new_path": "/path/to/new_model.sql", + "files_updated": ["/path/to/ref_file.sql"], + } + ) + + old_uri = "file:///path/to/old_model.sql" + new_uri = "file:///path/to/new_model.sql" + result = await rename_model(mock_lsp_client, old_uri, new_uri) + + assert result["renamed"] is True + assert result["old_path"] == "/path/to/old_model.sql" + assert result["new_path"] == "/path/to/new_model.sql" + assert len(result["files_updated"]) == 1 + mock_lsp_client.rename_model.assert_called_once_with( + old_uri=old_uri, new_uri=new_uri, apply_edits=True + ) + + +@pytest.mark.asyncio +async def test_rename_model_with_multiple_updates() -> None: + """Test model rename that updates multiple files.""" + mock_lsp_client = AsyncMock(spec=LSPClient) + mock_lsp_client.rename_model = AsyncMock( + return_value={ + "renamed": True, + "old_path": "/path/to/old_model.sql", + "new_path": "/path/to/new_model.sql", + "files_updated": [ + "/path/to/ref_file1.sql", + "/path/to/ref_file2.sql", + "/path/to/ref_file3.sql", + ], + } + ) + + old_uri = "file:///path/to/old_model.sql" + new_uri = "file:///path/to/new_model.sql" + result = await rename_model(mock_lsp_client, old_uri, new_uri) + + assert result["renamed"] is True + assert len(result["files_updated"]) == 3 + + +@pytest.mark.asyncio +async def test_rename_model_error() -> None: + """Test model rename with error response.""" + mock_lsp_client = AsyncMock(spec=LSPClient) + mock_lsp_client.rename_model = AsyncMock( + return_value={"error": "Model not found"} + ) + + result = await rename_model( + mock_lsp_client, + "file:///invalid/old_model.sql", + "file:///invalid/new_model.sql", + ) + + assert "error" in result + assert "Model not found" in result["error"] + + +@pytest.mark.asyncio +async def test_rename_model_timeout() -> None: + """Test model rename with timeout error.""" + mock_lsp_client = AsyncMock(spec=LSPClient) + mock_lsp_client.rename_model = AsyncMock(side_effect=TimeoutError()) + + result = await rename_model( + mock_lsp_client, + "file:///path/to/old_model.sql", + "file:///path/to/new_model.sql", + ) + + assert "error" in result + assert "Timeout waiting for model rename" in result["error"] + + +@pytest.mark.asyncio +async def test_rename_model_generic_exception() -> None: + """Test model rename with generic exception.""" + mock_lsp_client = AsyncMock(spec=LSPClient) + mock_lsp_client.rename_model = AsyncMock(side_effect=Exception("Connection lost")) + + result = await rename_model( + mock_lsp_client, + "file:///path/to/old_model.sql", + "file:///path/to/new_model.sql", + ) + + assert "error" in result + assert "Failed to rename model" in result["error"] + assert "Connection lost" in result["error"] From 46c752db6b0b9351b3719ecdaf5ff85c160dcf24 Mon Sep 17 00:00:00 2001 From: Benoit Perigaud <8754100+b-per@users.noreply.github.com> Date: Fri, 31 Oct 2025 14:37:02 +0100 Subject: [PATCH 2/3] Running `task check` --- src/dbt_mcp/lsp/lsp_client.py | 88 +++++++++++++++---------------- tests/unit/lsp/test_lsp_client.py | 27 +++++----- tests/unit/lsp/test_lsp_tools.py | 4 +- 3 files changed, 59 insertions(+), 60 deletions(-) diff --git a/src/dbt_mcp/lsp/lsp_client.py b/src/dbt_mcp/lsp/lsp_client.py index 0d454ee1..80354617 100644 --- a/src/dbt_mcp/lsp/lsp_client.py +++ b/src/dbt_mcp/lsp/lsp_client.py @@ -158,11 +158,9 @@ def _uri_to_path(self, uri: str) -> Path: path_str = unquote(parsed.path) return Path(path_str) - def _apply_single_edit( - self, lines: list[str], edit: dict[str, Any] - ) -> None: + def _apply_single_edit(self, lines: list[str], edit: dict[str, Any]) -> None: """Apply a single text edit to a list of lines. - + Args: lines: List of lines (with line endings preserved) edit: LSP TextEdit object with 'range' and 'newText' @@ -172,23 +170,23 @@ def _apply_single_edit( end_line_idx = edit["range"]["end"]["line"] end_char = edit["range"]["end"]["character"] new_text = edit["newText"] - + # Handle single line edit if start_line_idx == end_line_idx: if start_line_idx < len(lines): line = lines[start_line_idx] # Strip line ending to work with just the text - line_ending = '' - if line.endswith('\r\n'): - line_ending = '\r\n' + line_ending = "" + if line.endswith("\r\n"): + line_ending = "\r\n" line = line[:-2] - elif line.endswith('\n'): - line_ending = '\n' + elif line.endswith("\n"): + line_ending = "\n" line = line[:-1] - elif line.endswith('\r'): - line_ending = '\r' + elif line.endswith("\r"): + line_ending = "\r" line = line[:-1] - + # Apply the edit lines[start_line_idx] = ( line[:start_char] + new_text + line[end_char:] + line_ending @@ -197,39 +195,37 @@ def _apply_single_edit( # Multi-line edit start_line = lines[start_line_idx] end_line = lines[end_line_idx] - + # Strip line endings - start_line_ending = '' - if start_line.endswith('\r\n'): - start_line_ending = '\r\n' + if start_line.endswith("\r\n"): start_line = start_line[:-2] - elif start_line.endswith('\n'): - start_line_ending = '\n' + elif start_line.endswith("\n"): start_line = start_line[:-1] - elif start_line.endswith('\r'): - start_line_ending = '\r' + elif start_line.endswith("\r"): start_line = start_line[:-1] - - end_line_ending = '' - if end_line.endswith('\r\n'): - end_line_ending = '\r\n' + + end_line_ending = "" + if end_line.endswith("\r\n"): + end_line_ending = "\r\n" end_line = end_line[:-2] - elif end_line.endswith('\n'): - end_line_ending = '\n' + elif end_line.endswith("\n"): + end_line_ending = "\n" end_line = end_line[:-1] - elif end_line.endswith('\r'): - end_line_ending = '\r' + elif end_line.endswith("\r"): + end_line_ending = "\r" end_line = end_line[:-1] - + start_line_text = start_line[:start_char] end_line_text = end_line[end_char:] - lines[start_line_idx] = start_line_text + new_text + end_line_text + end_line_ending + lines[start_line_idx] = ( + start_line_text + new_text + end_line_text + end_line_ending + ) # Remove the lines in between - del lines[start_line_idx + 1:end_line_idx + 1] + del lines[start_line_idx + 1 : end_line_idx + 1] def _apply_text_edits(self, file_path: Path, edits: list[dict[str, Any]]) -> None: """Apply text edits to a file. - + Args: file_path: Path to the file to edit edits: List of LSP TextEdit objects with 'range' and 'newText' @@ -237,28 +233,28 @@ def _apply_text_edits(self, file_path: Path, edits: list[dict[str, Any]]) -> Non if not file_path.exists(): logger.warning(f"File not found for edits: {file_path}") return - + # Read the file content content = file_path.read_text() - + # Sort edits in reverse order (end to start) to avoid offset issues # We sort by start position descending so we can apply from end to beginning sorted_edits = sorted( edits, key=lambda e: ( e["range"]["start"]["line"], - e["range"]["start"]["character"] + e["range"]["start"]["character"], ), - reverse=True + reverse=True, ) - + # Convert content to list of lines for easier manipulation lines = content.splitlines(keepends=True) - + # Apply each edit for edit in sorted_edits: self._apply_single_edit(lines, edit) - + # Write back to file file_path.write_text("".join(lines)) logger.info(f"Applied {len(edits)} edits to {file_path}") @@ -346,7 +342,9 @@ async def rename_model( self._apply_text_edits(file_path, change["edits"]) files_updated.append(str(file_path)) except Exception as e: - logger.error(f"Failed to apply edits to {file_uri}: {e}") + logger.error( + f"Failed to apply edits to {file_uri}: {e}" + ) return { "error": f"Failed to apply edits to {file_uri}: {str(e)}" } @@ -354,13 +352,13 @@ async def rename_model( # Step 3: Perform the actual file rename old_path = self._uri_to_path(old_uri) new_path = self._uri_to_path(new_uri) - + if not old_path.exists(): return {"error": f"Source file does not exist: {old_path}"} - + if new_path.exists(): return {"error": f"Destination file already exists: {new_path}"} - + try: # Ensure parent directory exists new_path.parent.mkdir(parents=True, exist_ok=True) @@ -392,7 +390,7 @@ async def rename_model( "files_updated": files_updated, } - except asyncio.TimeoutError: + except TimeoutError: return {"error": "Timeout waiting for LSP response"} except Exception as e: return {"error": f"Failed to rename file: {str(e)}"} diff --git a/tests/unit/lsp/test_lsp_client.py b/tests/unit/lsp/test_lsp_client.py index 691be518..3e8322d6 100644 --- a/tests/unit/lsp/test_lsp_client.py +++ b/tests/unit/lsp/test_lsp_client.py @@ -1,7 +1,6 @@ """Tests for the DbtLspClient class.""" -from pathlib import Path -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import MagicMock import pytest @@ -100,10 +99,10 @@ async def test_rename_model_success_with_apply_edits( old_file = tmp_path / "old_file.sql" new_file = tmp_path / "new_file.sql" ref_file = tmp_path / "ref_file.sql" - + old_file.write_text("SELECT * FROM table") ref_file.write_text("SELECT * FROM {{ ref('old_file') }}") - + # Setup mock LSP response mock_result = { "changes": { @@ -130,14 +129,14 @@ async def test_rename_model_success_with_apply_edits( assert result["old_path"] == str(old_file) assert result["new_path"] == str(new_file) assert str(ref_file) in result["files_updated"] - + # Verify file was renamed assert not old_file.exists() assert new_file.exists() - + # Verify reference was updated assert ref_file.read_text() == "SELECT * FROM {{ ref('new_file') }}" - + # Verify didRenameFiles notification was sent mock_lsp_connection.send_notification.assert_called_once() @@ -191,7 +190,9 @@ async def test_rename_model_error(lsp_client, mock_lsp_connection): @pytest.mark.asyncio -async def test_rename_model_source_not_exists(lsp_client, mock_lsp_connection, tmp_path): +async def test_rename_model_source_not_exists( + lsp_client, mock_lsp_connection, tmp_path +): """Test model rename when source file doesn't exist.""" # Setup mock mock_lsp_connection.send_request.return_value = {"changes": {}} @@ -207,14 +208,16 @@ async def test_rename_model_source_not_exists(lsp_client, mock_lsp_connection, t @pytest.mark.asyncio -async def test_rename_model_destination_exists(lsp_client, mock_lsp_connection, tmp_path): +async def test_rename_model_destination_exists( + lsp_client, mock_lsp_connection, tmp_path +): """Test model rename when destination file already exists.""" # Create both files old_file = tmp_path / "old.sql" new_file = tmp_path / "new.sql" old_file.write_text("content") new_file.write_text("existing") - + # Setup mock mock_lsp_connection.send_request.return_value = {"changes": {}} @@ -235,7 +238,7 @@ async def test_rename_model_none_response(lsp_client, mock_lsp_connection, tmp_p old_file = tmp_path / "old.sql" new_file = tmp_path / "new.sql" old_file.write_text("content") - + # Setup mock to return None (simulating the user's issue) mock_lsp_connection.send_request.return_value = None @@ -249,7 +252,7 @@ async def test_rename_model_none_response(lsp_client, mock_lsp_connection, tmp_p assert result["old_path"] == str(old_file) assert result["new_path"] == str(new_file) assert result["files_updated"] == [] - + # Verify file was renamed assert not old_file.exists() assert new_file.exists() diff --git a/tests/unit/lsp/test_lsp_tools.py b/tests/unit/lsp/test_lsp_tools.py index 1a1bd587..3178e677 100644 --- a/tests/unit/lsp/test_lsp_tools.py +++ b/tests/unit/lsp/test_lsp_tools.py @@ -262,9 +262,7 @@ async def test_rename_model_with_multiple_updates() -> None: async def test_rename_model_error() -> None: """Test model rename with error response.""" mock_lsp_client = AsyncMock(spec=LSPClient) - mock_lsp_client.rename_model = AsyncMock( - return_value={"error": "Model not found"} - ) + mock_lsp_client.rename_model = AsyncMock(return_value={"error": "Model not found"}) result = await rename_model( mock_lsp_client, From 16ff0972712c373d1648ad3b8ce33dbec8ff962a Mon Sep 17 00:00:00 2001 From: Benoit Perigaud <8754100+b-per@users.noreply.github.com> Date: Fri, 31 Oct 2025 14:42:41 +0100 Subject: [PATCH 3/3] Fix tests --- src/dbt_mcp/tools/policy.py | 3 +++ src/dbt_mcp/tools/toolsets.py | 1 + 2 files changed, 4 insertions(+) diff --git a/src/dbt_mcp/tools/policy.py b/src/dbt_mcp/tools/policy.py index b4b575a1..bb31365f 100644 --- a/src/dbt_mcp/tools/policy.py +++ b/src/dbt_mcp/tools/policy.py @@ -149,4 +149,7 @@ class ToolPolicy: ToolName.GET_COLUMN_LINEAGE.value: ToolPolicy( name=ToolName.GET_COLUMN_LINEAGE.value, behavior=ToolBehavior.METADATA ), + ToolName.RENAME_MODEL.value: ToolPolicy( + name=ToolName.RENAME_MODEL.value, behavior=ToolBehavior.METADATA + ), } diff --git a/src/dbt_mcp/tools/toolsets.py b/src/dbt_mcp/tools/toolsets.py index 2dbc3c1c..5cd4433f 100644 --- a/src/dbt_mcp/tools/toolsets.py +++ b/src/dbt_mcp/tools/toolsets.py @@ -73,5 +73,6 @@ class Toolset(Enum): }, Toolset.DBT_LSP: { ToolName.GET_COLUMN_LINEAGE, + ToolName.RENAME_MODEL, }, }