feat: first step of daemon mode (not ready)

[ci skip]

Author: benfdking

sqlmesh/cli/daemon_connector.py

@@ -53,6 +53,9 @@ def _validate_lock_file(lock_file_path: Path) -> LockFile:
     return read_file
 
 
+import socket
+
+
 class DaemonConnector:
     """Connects to the LSP daemon via socket to execute commands."""
 
@@ -61,37 +64,25 @@ def __init__(self, project_path: Path, lock_file: LockFile):
         self.lock_file = lock_file
         self.renderer = JanitorStateRenderer()
 
-    def _open_connection(self) -> t.Any:
-        """Open connection (named pipe or Unix socket) for communication."""
-        if self.lock_file.communication is None:
-            raise ValueError("Lock file does not contain communication information")
-
-        comm_mode = self.lock_file.communication.type
-
-        if isinstance(comm_mode, DaemonCommunicationModeUnixSocket):
-            socket_path = comm_mode.socket
-
-            # Check if it's a Unix domain socket or named pipe
-            if socket_path.endswith(".sock"):
-                # Unix domain socket
-                import socket
-
-                sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-                sock.connect(socket_path)
-
-                # Convert to file-like object
-                import io
-
-                return io.BufferedRWPair(
-                    io.BufferedReader(io.FileIO(sock.fileno(), mode="r", closefd=False)),
-                    io.BufferedWriter(io.FileIO(sock.fileno(), mode="w", closefd=False)),
-                )
-            else:
-                # Named pipe
-                return open(socket_path, "r+b")
+    def _open_connection(self) -> tuple[t.BinaryIO, t.BinaryIO]:
+        lock_file = self.lock_file
+        communication = lock_file.communication
+        print(f"Using communication mode: {communication}")
+        if communication is None:
+            raise ValueError("not correct")
+
+        if isinstance(communication.type, DaemonCommunicationModeUnixSocket):
+            print("Opening Unix socket connection...")
+            sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+            sock.connect(communication.type.socket)
+            print(f"Connected to Unix socket at {communication.type.socket}")
+            rfile = sock.makefile("rb", buffering=0)
+            wfile = sock.makefile("wb", buffering=0)
+            print("Connected to daemon via Unix socket.")
+            return rfile, wfile
         else:
-            raise ValueError(f"Only Unix socket communication is supported")
-
+            raise ValueError("Only Unix socket communication is supported")
+        
     def _send_jsonrpc_request(self, connection: t.Any, method: str, params: dict) -> str:
         """Send a JSON-RPC request over the connection and return the request ID."""
         request_id = str(uuid.uuid4())
@@ -152,61 +143,45 @@ def _read_jsonrpc_response(self, connection: t.Any, expected_id: str) -> t.Any:
         return message.get("result", {})
 
     def call_janitor(self, ignore_ttl: bool = False) -> bool:
-        """Call the janitor command through the LSP daemon."""
-        connection = None
+        rfile = wfile = None
         try:
-            connection = self._open_connection()
+            rfile, wfile = self._open_connection()
 
-            # Send the janitor request via JSON-RPC
             request = LSPCLICallRequest(
                 arguments=["janitor"] + (["--ignore-ttl"] if ignore_ttl else [])
             )
+            request_id = self._send_jsonrpc_request(wfile, "sqlmesh/cli/call", request.model_dump())
 
-            request_id = self._send_jsonrpc_request(
-                connection, "sqlmesh/cli/call", request.model_dump()
-            )
-
-            # Listen for notifications and the final response
             with self.renderer as renderer:
                 while True:
                     try:
-                        # Read any JSON-RPC message (response or notification)
-                        message_data = self._read_jsonrpc_message(connection)
-
-                        # Check if it's the response to our request
+                        message_data = self._read_jsonrpc_message(rfile)
                         if "id" in message_data and message_data["id"] == request_id:
-                            # This is the response to our request
                             result = message_data.get("result", {})
                             if result.get("state") == "finished":
                                 return True
                             elif result.get("state") == "error":
                                 print(f"Error from daemon: {result.get('message', 'Unknown error')}")
                                 return False
-
-                        # Check if it's a notification with updates
                         elif message_data.get("method") == "sqlmesh/cli/update":
                             params = message_data.get("params", {})
                             if params.get("state") == "ongoing":
-                                # Parse the janitor state and render it
                                 message = params.get("message", {})
                                 if "state" in message:
                                     janitor_state = JanitorState.model_validate(message)
                                     renderer.render(janitor_state.state)
-
                     except Exception as stream_error:
-                        # If we can't read more messages, assume we're done
                         print(f"Stream ended: {stream_error}")
                         break
-
             return True
-
         except Exception as e:
             print(f"Failed to communicate with daemon: {e}")
             return False
         finally:
-            if connection:
-                connection.close()
-
+            try:
+                if rfile: rfile.close()
+            finally:
+                if wfile: wfile.close()
 
 def get_daemon_connector(project_path: Path) -> t.Optional[DaemonConnector]:
     """Get a daemon connector if a valid lock file exists."""

sqlmesh/lsp/main.py

@@ -3,6 +3,8 @@
 
 from itertools import chain
 import logging
+import socket
+import sys
 import typing as t
 from pathlib import Path
 import urllib.parse
@@ -466,14 +468,16 @@ def _custom_cli_call(self, ls: LanguageServer, params: LSPCLICallRequest) -> Soc
         """Handle CLI call requests from the daemon connector."""
         try:
             context = self._context_get_or_load()
-            if not context or not hasattr(context, 'context'):
+            if not context or not hasattr(context, "context"):
                 return SocketMessage(state=SocketMessageError(message="No context available"))
 
-            arguments = params.argumens if hasattr(params, 'argumens') else params.arguments
+            arguments = params.argumens if hasattr(params, "argumens") else params.arguments
 
             # For now, only support janitor command
             if not arguments or arguments[0] != "janitor":
-                return SocketMessage(state=SocketMessageError(message="Only 'janitor' command is supported"))
+                return SocketMessage(
+                    state=SocketMessageError(message="Only 'janitor' command is supported")
+                )
 
             # Parse janitor arguments
             ignore_ttl = "--ignore-ttl" in arguments
@@ -490,10 +494,7 @@ def send_state(state: JanitorState):
                     notification = {
                         "jsonrpc": "2.0",
                         "method": "sqlmesh/cli/update",
-                        "params": {
-                            "state": "ongoing",
-                            "message": state.state.model_dump()
-                        }
+                        "params": {"state": "ongoing", "message": state.state.model_dump()},
                     }
 
                     # Send notification through the server's transport
@@ -503,7 +504,7 @@ def send_state(state: JanitorState):
                     full_message = header.encode("utf-8") + message.encode("utf-8")
 
                     # Write directly to the server's output stream
-                    if hasattr(ls.lsp, '_writer') and ls.lsp._writer:
+                    if hasattr(ls.lsp, "_writer") and ls.lsp._writer:
                         ls.lsp._writer.write(full_message)
                         ls.lsp._writer.flush()
 
@@ -1284,47 +1285,15 @@ def _uri_to_path(uri: str) -> Path:
         """Convert a URI to a path."""
         return URI(uri).to_path()
 
-    def start(self, socket_path: t.Optional[str] = None) -> None:
+    def start(self, rfile: t.Optional[t.Any], wfile: t.Optional[t.Any]) -> None:
         """Start the server with Unix socket (for both VS Code and CLI)."""
         logging.basicConfig(level=logging.DEBUG)
 
-        if socket_path:
-            import os
-            import socket
-            from pathlib import Path
-
-            # Write lock file with socket information
-            project_path = self.workspace_folders[0] if self.workspace_folders else Path.cwd()
-            lock_path = return_lock_file_path(project_path)
-            lock_path.parent.mkdir(parents=True, exist_ok=True)
-
-            lock_file = generate_lock_file()
-            lock_file.communication = DaemonCommunicationMode(
-                type=DaemonCommunicationModeUnixSocket(socket=socket_path)
-            )
-            lock_path.write_text(lock_file.model_dump_json(indent=2))
-
-            # Remove existing socket file if it exists
-            if os.path.exists(socket_path):
-                os.unlink(socket_path)
-
-            # Create and bind to the Unix domain socket
-            server_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-            server_sock.bind(socket_path)
-            server_sock.listen(1)
-
-            # Accept the first connection (from VS Code)
-            client_sock, _ = server_sock.accept()
-
-            # Convert socket to file-like objects for reading and writing
-            import io
-            reader = io.BufferedReader(io.FileIO(client_sock.fileno(), mode='r', closefd=False))
-            writer = io.BufferedWriter(io.FileIO(client_sock.fileno(), mode='w', closefd=False))
-
-            self.server.start_io(reader, writer)
-        else:
-            # Use standard I/O
+        if rfile is None or wfile is None:
             self.server.start_io()
+        else:
+            self.server.start_io(rfile, wfile)
+    
 
 
 def loaded_sqlmesh_message(ls: LanguageServer) -> None:
@@ -1338,12 +1307,36 @@ def main() -> None:
     import argparse
 
     parser = argparse.ArgumentParser()
-    parser.add_argument('--socket', help='Unix socket path for communication', default=None)
+    parser.add_argument("--pipe", help="Unix socket path for communication", default=None)
     args = parser.parse_args()
 
     # Example instantiator that just uses the same signature as your original `Context` usage.
+    file = args.pipe
     sqlmesh_server = SQLMeshLanguageServer(context_class=Context)
-    sqlmesh_server.start(socket_path=args.socket)
+    if args.pipe:
+        # Connect to the pipe the VS Code client created
+        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        try:
+            sock.connect(args.pipe)
+        except OSError as e:
+            print(f"Failed to connect to pipe {args.pipe}: {e}", file=sys.stderr)
+            sys.exit(1)
+
+        # Wrap the socket as file-like binary streams for pygls
+        rfile = sock.makefile("rb", buffering=0)
+        wfile = sock.makefile("wb", buffering=0)
+
+        try:
+            sqlmesh_server.start(rfile, wfile)
+        finally:
+            try: rfile.close()
+            except: pass
+            try: wfile.close()
+            except: pass
+            try: sock.close()
+            except: pass
+    else:
+        sqlmesh_server.start()
 
 
 if __name__ == "__main__":

vscode/extension/src/lsp/lsp.ts

@@ -17,6 +17,9 @@ import {
 } from '../utilities/errors'
 import { CustomLSPMethods } from './custom'
 import { resolveProjectPath } from '../utilities/config'
+import * as path from 'path'
+import * as os from 'os'
+import * as crypto from 'crypto'
 
 type SupportedMethodsState =
   | { type: 'not-fetched' }
@@ -97,27 +100,20 @@ export class LSPClient implements Disposable {
 
     const workspacePath = sqlmesh.value.workspacePath
 
-    // Create a unique socket path for this workspace
-    const os = require('os')
-    const path = require('path')
-    const crypto = require('crypto')
-    const tmpDir = os.tmpdir()
-    const workspaceHash = crypto.createHash('md5').update(workspacePath).digest('hex').substring(0, 8)
-    const socketPath = path.join(tmpDir, `sqlmesh_${workspaceHash}.sock`)
-
     // Add --socket argument to the LSP server
-    const argsWithSocket = [...sqlmesh.value.args, '--socket', socketPath]
-
+    const socketPath = socketAddressForWorkspace(workspacePath)
+    const argsWithSocket = [...sqlmesh.value.args]
+      // '--socket', socketPath]
     const serverOptions: ServerOptions = {
       run: {
         command: sqlmesh.value.bin,
-        transport: TransportKind.ipc,
+        transport: TransportKind.pipe,
         options: { cwd: workspacePath, env: sqlmesh.value.env },
         args: argsWithSocket,
       },
       debug: {
         command: sqlmesh.value.bin,
-        transport: TransportKind.ipc,
+        transport: TransportKind.pipe,
         options: { cwd: workspacePath, env: sqlmesh.value.env },
         args: argsWithSocket,
       },
@@ -279,3 +275,17 @@ export class LSPClient implements Disposable {
     }
   }
 }
+
+function socketAddressForWorkspace(workspacePath: string) {
+  const hash = crypto.createHash('md5').update(workspacePath).digest('hex').slice(0, 8)
+  const base = `sqlmesh_${hash}`
+
+  if (process.platform === 'win32') {
+    // Windows wants a Named Pipe path, not a filesystem .sock file
+    // NOTE: double backslashes are required in JS strings
+    return `\\\\.\\pipe\\${base}`
+  }
+
+  // POSIX: use a short dir (tmp) to avoid the ~108 byte sun_path limit
+  return path.join(os.tmpdir(), `${base}.sock`)
+}