[WIP / DO NOT MERGE] Initial LSP implementation

pyproject.toml

@@ -42,6 +42,8 @@ dependencies = [
   "GitPython==3.1.43",
   "pip",
   "pydantic==2.9.1",
+  "pygls==1.3.1",
+  "pytest_asyncio==0.23.8"
 ]
 classifiers = [
   "Development Status :: 5 - Production/Stable",

snyk/requirements.txt

@@ -15,6 +15,8 @@ urllib3>=1.24.3,<2.3
 GitPython==3.1.43
 pip
 pydantic==2.9.1
+pygls==1.3.1
+pytest_asyncio==0.23.8
 coverage==7.6.1
 pre-commit>=3.5.0
 pytest==8.3.2

src/snowflake/cli/_app/commands_registration/__init__.py

@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from dataclasses import dataclass
+from typing import Callable, Optional
 
 from snowflake.cli.api.plugins.command import CommandSpec
 
@@ -21,6 +24,7 @@
 class LoadedCommandPlugin:
     plugin_name: str
     command_spec: CommandSpec
+    lsp_spec: Optional[Callable] = None
 
 
 @dataclass

src/snowflake/cli/_app/commands_registration/builtin_plugins.py

@@ -16,6 +16,7 @@
 from snowflake.cli._plugins.cortex import plugin_spec as cortex_plugin_spec
 from snowflake.cli._plugins.git import plugin_spec as git_plugin_spec
 from snowflake.cli._plugins.init import plugin_spec as init_plugin_spec
+from snowflake.cli._plugins.lsp import plugin_spec as lsp_plugin_spec
 from snowflake.cli._plugins.nativeapp import plugin_spec as nativeapp_plugin_spec
 from snowflake.cli._plugins.notebook import plugin_spec as notebook_plugin_spec
 from snowflake.cli._plugins.object import plugin_spec as object_plugin_spec
@@ -34,6 +35,7 @@ def get_builtin_plugin_name_to_plugin_spec():
         "spcs": spcs_plugin_spec,
         "nativeapp": nativeapp_plugin_spec,
         "object": object_plugin_spec,
+        "lsp": lsp_plugin_spec,
         "snowpark": snowpark_plugin_spec,
         "stage": stage_plugin_spec,
         "sql": sql_plugin_spec,

src/snowflake/cli/_app/commands_registration/command_plugins_loader.py

@@ -118,10 +118,10 @@ def _load_builtin_plugin_spec(
         self, plugin_name: str, plugin
     ) -> Optional[LoadedCommandPlugin]:
         command_spec = self._load_command_spec(plugin_name, plugin)
+        lsp_spec = self._load_lsp_spec(plugin_name, plugin)
         if command_spec:
             return LoadedBuiltInCommandPlugin(
-                plugin_name=plugin_name,
-                command_spec=command_spec,
+                plugin_name=plugin_name, command_spec=command_spec, lsp_spec=lsp_spec
             )
         else:
             return None
@@ -153,6 +153,19 @@ def _load_command_spec(plugin_name: str, plugin) -> Optional[CommandSpec]:
     def _is_external_plugin(plugin) -> bool:
         return isinstance(plugin, LoadedExternalCommandPlugin)
 
+    @staticmethod
+    def _load_lsp_spec(plugin_name: str, plugin) -> Optional[CommandSpec]:
+        try:
+            if getattr(plugin, "lsp_spec", None) is not None:
+                return plugin.lsp_spec()
+            return None
+        except Exception as ex:
+            log_exception(
+                f"Cannot load lsp specification from plugin [{plugin_name}]: {ex.__str__()}",
+                ex,
+            )
+            return None
+
 
 def load_only_builtin_command_plugins() -> List[LoadedCommandPlugin]:
     loader = CommandPluginsLoader()

src/snowflake/cli/_app/snow_connector.py

@@ -14,10 +14,11 @@
 
 from __future__ import annotations
 
+import json
 import contextlib
 import logging
 import os
-from typing import Dict, Optional
+from typing import Any, Dict, Optional
 
 import snowflake.connector
 from click.exceptions import ClickException
@@ -295,3 +296,27 @@ def prepare_private_key(private_key_pem, private_key_passphrase=None):
         format=PrivateFormat.PKCS8,
         encryption_algorithm=NoEncryption(),
     )
+
+#########
+
+# FIXME: circular logic
+# _MEMOIZED_CONNECTIONS: Dict[str, SnowflakeConnection] = {}
+
+# def open_and_memoize(args: ConnectionArgs = {}):
+#     """
+#     Finds an existing open connection, or connects to Snowflake with a given set of connection parameters.
+#     Connections are treated as temporary unless otherwise specified in the connection arguments.
+#     """
+#     key = json.dumps(args)
+#     if key not in _MEMOIZED_CONNECTIONS:
+#         _MEMOIZED_CONNECTIONS[key] = connect_to_snowflake(temporary_connection=True, **args)
+#     return _MEMOIZED_CONNECTIONS[key]
+
+
+# def close_memoized_connections():
+#     """
+#     Closes and forgets all connections made through get_memoized_connection.
+#     """
+#     for key in _MEMOIZED_CONNECTIONS.keys():
+#         _MEMOIZED_CONNECTIONS[key].close()
+#         del _MEMOIZED_CONNECTIONS[key]

src/snowflake/cli/_plugins/nativeapp/lsp_commands.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from pygls.server import LanguageServer
+from snowflake.cli.api.cli_global_context import get_cli_context
+from snowflake.cli.api.output.types import MessageResult
+from snowflake.cli.api.project.definition_manager import DefinitionManager
+from snowflake.cli.plugins.lsp.server import lsp_plugin
+from snowflake.cli.plugins.nativeapp.manager import NativeAppManagern
+
+
+@lsp_plugin(
+    name="nativeapp",
+    capabilities={
+        "openApplication": True,
+    },
+)
+def nade_lsp_plugin(server: LanguageServer):
+    # FIXME: can't parametrize iter_lsp_plugins() if this is top-level ?
+    from snowflake.cli.plugins.lsp.interface import workspace_command
+
+    @workspace_command(server, "openApplication")
+    def open_app() -> MessageResult:
+
+        ctx = get_cli_context()
+
+        dm = DefinitionManager(ctx.project_root)
+        project_definition = getattr(dm.project_definition, "native_app", None)
+        project_root = dm.project_root
+        manager = NativeAppManager(
+            project_definition=project_definition,
+            project_root=project_root,
+            connection=ctx.connection,
+        )
+        if manager.get_existing_app_info():
+            url = manager.get_snowsight_url()
+            return MessageResult(f"{url}")
+        else:
+            return MessageResult(
+                'Snowflake Native App not yet deployed! Please run "runApplication" first.'
+            )

src/snowflake/cli/_plugins/nativeapp/plugin_spec.py

@@ -17,8 +17,11 @@
     SNOWCLI_ROOT_COMMAND_PATH,
     CommandSpec,
     CommandType,
+    lsp_plugin_hook_impl,
     plugin_hook_impl,
 )
+from snowflake.cli._plugins.nativeapp import commands
+from snowflake.cli._plugins.nativeapp.lsp_commands import nade_lsp_plugin
 
 
 @plugin_hook_impl
@@ -28,3 +31,8 @@ def command_spec():
         command_type=CommandType.COMMAND_GROUP,
         typer_instance=commands.app.create_instance(),
     )
+
+
+@lsp_plugin_hook_impl
+def lsp_spec():
+    return nade_lsp_plugin

src/snowflake/cli/api/plugins/command/__init__.py

@@ -28,6 +28,8 @@
 
 plugin_hook_spec = pluggy.HookspecMarker(SNOWCLI_COMMAND_PLUGIN_NAMESPACE)
 plugin_hook_impl = pluggy.HookimplMarker(SNOWCLI_COMMAND_PLUGIN_NAMESPACE)
+lsp_plugin_hook_spec = pluggy.HookspecMarker(SNOWCLI_COMMAND_PLUGIN_NAMESPACE)
+lsp_plugin_hook_impl = pluggy.HookimplMarker(SNOWCLI_COMMAND_PLUGIN_NAMESPACE)
 
 
 class CommandPath:

src/snowflake/cli/api/plugins/command/plugin_hook_specs.py

@@ -12,10 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from snowflake.cli.api.plugins.command import plugin_hook_spec
+from snowflake.cli.api.plugins.command import lsp_plugin_hook_spec, plugin_hook_spec
 
 
 @plugin_hook_spec
 def command_spec():
     """Command spec"""
     pass
+
+
+@lsp_plugin_hook_spec
+def lsp_spec():
+    """LSP spec"""
+    pass

src/snowflake/cli/plugins/lsp/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

src/snowflake/cli/plugins/lsp/commands.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+
+from snowflake.cli.api.commands.snow_typer import SnowTyper
+from snowflake.cli.api.output.types import CommandResult, MessageResult
+from snowflake.cli.plugins.lsp.server import (
+    start_lsp_server,
+)
+
+app = SnowTyper(name="lsp", help="Manages a Snowflake LSP server.", hidden=True)
+
+log = logging.getLogger(__name__)
+
+
+@app.command("start")
+def lsp_start(
+    **options,
+) -> CommandResult:
+    """
+    Starts the LSP language server in the foreground.
+    """
+    start_lsp_server()
+    return MessageResult(f"LSP server process ended.")

src/snowflake/cli/plugins/lsp/interface.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union, get_type_hints, is_typeddict
+from pydantic import TypeAdapter, ValidationError
+from typing_extensions import TypedDict, NotRequired
+from snowflake.cli.api.cli_global_context import CliContextArguments, fork_cli_context
+
+from pygls.server import LanguageServer
+
+
+ORIGINAL_FUNCTION_KEY = "__lsp_original_function__"
+
+
+TypeDef = Dict[str, Union[str, 'TypeDef']]
+
+class CommandArguments(CliContextArguments):
+    """
+    The arguments that can be passed to a workspace command.
+
+    """
+    args: NotRequired[List[Any]]
+    kwargs: NotRequired[Dict[str, Any]]
+
+
+def workspace_command(
+    server: LanguageServer,
+    name: str,
+    requires_connection: bool = False,
+    requires_project: bool = False,
+):
+    """
+    Wrap a function with pygls' @server.command.
+    Ensures that the command invocation provides a valid connection / project context
+    (if required) as well as ensuring arguments are in the required format.
+    """
+    def _decorator(func):
+        @server.command(name)
+        def wrapper(params: List[CommandArguments]):
+            if len(params) > 1:
+                raise ValueError("Expected exactly one CommandArguments object")
+
+            try:
+                args = TypeAdapter(CommandArguments).validate_python(params[0] if len(params) == 1 else {})
+
+                if requires_connection and "connection" not in args:
+                    raise ValueError("connection missing, but requires_connection=True")
+
+                if requires_project and "project_path" not in args:
+                    raise ValueError("project_path missing, but requires_connection=True")
+
+                # TODO: validation of args.args / args.kwargs based on shape of actual command...
+
+                with fork_cli_context(**args):
+                    return func(*args.get("args", []), **args.get("kwargs", {}))
+
+            except ValidationError as exc: 
+                raise ValueError(f"ERROR: Invalid schema: {exc}")
+
+
+        setattr(wrapper, ORIGINAL_FUNCTION_KEY, func)
+        return wrapper
+
+    return _decorator

src/snowflake/cli/plugins/lsp/plugin_spec.py

@@ -0,0 +1,30 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from snowflake.cli.api.plugins.command import (
+    SNOWCLI_ROOT_COMMAND_PATH,
+    CommandSpec,
+    CommandType,
+    plugin_hook_impl,
+)
+from snowflake.cli.plugins.lsp import commands
+
+
+@plugin_hook_impl
+def command_spec():
+    return CommandSpec(
+        parent_command_path=SNOWCLI_ROOT_COMMAND_PATH,
+        command_type=CommandType.COMMAND_GROUP,
+        typer_instance=commands.app,
+    )