snowflake-cli 2.8.2__py3-none-any.whl

snowflake/cli/app/cli_app.py

@@ -0,0 +1,243 @@
+# 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
+import platform
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+
+import click
+import typer
+from click import Context
+from snowflake.cli import __about__
+from snowflake.cli.api import Api, api_provider
+from snowflake.cli.api.config import config_init
+from snowflake.cli.api.output.formats import OutputFormat
+from snowflake.cli.api.output.types import CollectionResult
+from snowflake.cli.api.secure_path import SecurePath
+from snowflake.cli.app.api_impl.plugin.plugin_config_provider_impl import (
+    PluginConfigProviderImpl,
+)
+from snowflake.cli.app.commands_registration.commands_registration_with_callbacks import (
+    CommandsRegistrationWithCallbacks,
+)
+from snowflake.cli.app.dev.commands_structure import generate_commands_structure
+from snowflake.cli.app.dev.docs.generator import generate_docs
+from snowflake.cli.app.dev.pycharm_remote_debug import (
+    setup_pycharm_remote_debugger_if_provided,
+)
+from snowflake.cli.app.main_typer import SnowCliMainTyper
+from snowflake.cli.app.printing import MessageResult, print_result
+from snowflake.connector.config_manager import CONFIG_MANAGER
+
+log = logging.getLogger(__name__)
+
+_api = Api(plugin_config_provider=PluginConfigProviderImpl())
+api_provider.register_api(_api)
+
+_commands_registration = CommandsRegistrationWithCallbacks(_api.plugin_config_provider)
+
+
+@dataclass
+class AppContextHolder:
+    # needed to access the context from tests
+    app_context: Optional[Context] = None
+
+
+app_context_holder = AppContextHolder()
+
+
+def _exit_with_cleanup():
+    _commands_registration.reset_running_instance_registration_state()
+    raise typer.Exit()
+
+
+def _do_not_execute_on_completion(callback):
+    def enriched_callback(value):
+        if click.get_current_context().resilient_parsing:
+            return
+        callback(value)
+
+    return enriched_callback
+
+
+def _commands_registration_callback(value: bool):
+    if value:
+        _commands_registration.register_commands_if_ready_and_not_registered_yet()
+    # required to make the tests working
+    # because a single test can execute multiple commands using always the same "app" instance
+    _commands_registration.reset_running_instance_registration_state()
+    app_context_holder.app_context = click.get_current_context()
+
+
+@_commands_registration.before
+def _config_init_callback(configuration_file: Optional[Path]):
+    config_init(configuration_file)
+
+
+@_commands_registration.before
+def _disable_external_command_plugins_callback(value: bool):
+    if value:
+        _commands_registration.disable_external_command_plugins()
+
+
+@_do_not_execute_on_completion
+@_commands_registration.after
+def _docs_callback(value: bool):
+    if value:
+        ctx = click.get_current_context()
+        generate_docs(SecurePath("gen_docs"), ctx.command)
+        _exit_with_cleanup()
+
+
+@_do_not_execute_on_completion
+@_commands_registration.after
+def _commands_structure_callback(value: bool):
+    if value:
+        ctx = click.get_current_context()
+        generate_commands_structure(ctx.command).print_node()
+        _exit_with_cleanup()
+
+
+@_do_not_execute_on_completion
+def _version_callback(value: bool):
+    if value:
+        print_result(MessageResult(f"Snowflake CLI version: {__about__.VERSION}"))
+        _exit_with_cleanup()
+
+
+from snowflake.cli.api.config import get_feature_flags_section
+
+
+@_do_not_execute_on_completion
+def _info_callback(value: bool):
+    if value:
+        result = CollectionResult(
+            [
+                {"key": "version", "value": __about__.VERSION},
+                {
+                    "key": "default_config_file_path",
+                    "value": str(CONFIG_MANAGER.file_path),
+                },
+                {"key": "python_version", "value": sys.version},
+                {"key": "system_info", "value": platform.platform()},
+                {"key": "feature_flags", "value": get_feature_flags_section()},
+            ],
+        )
+        print_result(result, output_format=OutputFormat.JSON)
+        _exit_with_cleanup()
+
+
+def app_factory() -> SnowCliMainTyper:
+    app = SnowCliMainTyper()
+
+    @app.callback(invoke_without_command=True)
+    def default(
+        ctx: typer.Context,
+        version: bool = typer.Option(
+            None,
+            "--version",
+            help="Shows version of the Snowflake CLI",
+            callback=_version_callback,
+            is_eager=True,
+        ),
+        docs: bool = typer.Option(
+            None,
+            "--docs",
+            hidden=True,
+            help="Generates Snowflake CLI documentation",
+            callback=_docs_callback,
+            is_eager=True,
+        ),
+        structure: bool = typer.Option(
+            None,
+            "--structure",
+            hidden=True,
+            help="Prints Snowflake CLI structure of commands",
+            callback=_commands_structure_callback,
+            is_eager=True,
+        ),
+        info: bool = typer.Option(
+            None,
+            "--info",
+            help="Shows information about the Snowflake CLI",
+            callback=_info_callback,
+        ),
+        configuration_file: Path = typer.Option(
+            None,
+            "--config-file",
+            help="Specifies Snowflake CLI configuration file that should be used",
+            exists=True,
+            dir_okay=False,
+            is_eager=True,
+            callback=_config_init_callback,
+        ),
+        pycharm_debug_library_path: str = typer.Option(
+            None,
+            "--pycharm-debug-library-path",
+            hidden=True,
+        ),
+        pycharm_debug_server_host: str = typer.Option(
+            "localhost",
+            "--pycharm-debug-server-host",
+            hidden=True,
+        ),
+        pycharm_debug_server_port: int = typer.Option(
+            12345,
+            "--pycharm-debug-server-port",
+            hidden=True,
+        ),
+        disable_external_command_plugins: bool = typer.Option(
+            None,
+            "--disable-external-command-plugins",
+            help="Disable external command plugins",
+            callback=_disable_external_command_plugins_callback,
+            is_eager=True,
+            hidden=True,
+        ),
+        # THIS OPTION SHOULD BE THE LAST OPTION IN THE LIST!
+        # ---
+        # This is a hidden artificial option used only to guarantee execution of commands registration
+        # and make this guaranty not dependent on other callbacks.
+        # Commands registration is invoked as soon as all callbacks
+        # decorated with "_commands_registration.before" are executed
+        # but if there are no such callbacks (at the result of possible future changes)
+        # then we need to invoke commands registration manually.
+        #
+        # This option is also responsible for resetting registration state for test purposes.
+        commands_registration: bool = typer.Option(
+            True,
+            "--commands-registration",
+            help="Commands registration",
+            hidden=True,
+            is_eager=True,
+            callback=_commands_registration_callback,
+        ),
+    ) -> None:
+        """
+        Snowflake CLI tool for developers.
+        """
+        if not ctx.invoked_subcommand:
+            typer.echo(ctx.get_help())
+        setup_pycharm_remote_debugger_if_provided(
+            pycharm_debug_library_path=pycharm_debug_library_path,
+            pycharm_debug_server_host=pycharm_debug_server_host,
+            pycharm_debug_server_port=pycharm_debug_server_port,
+        )
+
+    return app

snowflake/cli/app/commands_registration/__init__.py

@@ -0,0 +1,33 @@
+# 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 dataclasses import dataclass
+
+from snowflake.cli.api.plugins.command import CommandSpec
+
+
+@dataclass
+class LoadedCommandPlugin:
+    plugin_name: str
+    command_spec: CommandSpec
+
+
+@dataclass
+class LoadedBuiltInCommandPlugin(LoadedCommandPlugin):
+    pass
+
+
+@dataclass
+class LoadedExternalCommandPlugin(LoadedCommandPlugin):
+    pass

snowflake/cli/app/commands_registration/builtin_plugins.py

@@ -0,0 +1,54 @@
+# 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.plugins.connection import plugin_spec as connection_plugin_spec
+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.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
+
+# TODO 3.0: remove this import
+from snowflake.cli.plugins.object_stage_deprecated import (
+    plugin_spec as object_stage_deprecated_plugin_spec,
+)
+from snowflake.cli.plugins.snowpark import plugin_spec as snowpark_plugin_spec
+from snowflake.cli.plugins.spcs import plugin_spec as spcs_plugin_spec
+from snowflake.cli.plugins.sql import plugin_spec as sql_plugin_spec
+from snowflake.cli.plugins.stage import plugin_spec as stage_plugin_spec
+from snowflake.cli.plugins.streamlit import plugin_spec as streamlit_plugin_spec
+from snowflake.cli.plugins.workspace import plugin_spec as workspace_plugin_spec
+
+
+# plugin name to plugin spec
+def get_builtin_plugin_name_to_plugin_spec():
+    plugin_specs = {
+        "connection": connection_plugin_spec,
+        "spcs": spcs_plugin_spec,
+        "nativeapp": nativeapp_plugin_spec,
+        "object": object_plugin_spec,
+        "snowpark": snowpark_plugin_spec,
+        "stage": stage_plugin_spec,
+        "sql": sql_plugin_spec,
+        "streamlit": streamlit_plugin_spec,
+        "git": git_plugin_spec,
+        "notebook": notebook_plugin_spec,
+        "object-stage-deprecated": object_stage_deprecated_plugin_spec,
+        "cortex": cortex_plugin_spec,
+        "init": init_plugin_spec,
+        "workspace": workspace_plugin_spec,
+    }
+
+    return plugin_specs

snowflake/cli/app/commands_registration/command_plugins_loader.py

@@ -0,0 +1,169 @@
+# 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 typing import Dict, List, Optional
+
+import pluggy
+from snowflake.cli.api.plugins.command import (
+    SNOWCLI_COMMAND_PLUGIN_NAMESPACE,
+    CommandPath,
+    CommandSpec,
+    plugin_hook_specs,
+)
+from snowflake.cli.app.commands_registration import (
+    LoadedBuiltInCommandPlugin,
+    LoadedCommandPlugin,
+    LoadedExternalCommandPlugin,
+)
+from snowflake.cli.app.commands_registration.builtin_plugins import (
+    get_builtin_plugin_name_to_plugin_spec,
+)
+from snowflake.cli.app.commands_registration.exception_logging import exception_logging
+
+log = logging.getLogger(__name__)
+log_exception = exception_logging(log)
+
+
+class CommandPluginsLoader:
+    def __init__(self):
+        plugin_manager = pluggy.PluginManager(SNOWCLI_COMMAND_PLUGIN_NAMESPACE)
+        plugin_manager.add_hookspecs(plugin_hook_specs)
+        self._plugin_manager = plugin_manager
+        self._loaded_plugins: Dict[str, LoadedCommandPlugin] = {}
+        self._loaded_command_paths: Dict[CommandPath, LoadedCommandPlugin] = {}
+
+    def register_builtin_plugins(self) -> None:
+        for plugin_name, plugin in get_builtin_plugin_name_to_plugin_spec().items():
+            try:
+                self._plugin_manager.register(plugin=plugin, name=plugin_name)
+            except Exception as ex:
+                log_exception(
+                    f"Cannot register plugin [{plugin_name}]: {ex.__str__()}", ex
+                )
+
+    def register_external_plugins(self, plugin_names: List[str]) -> None:
+        for plugin_name in plugin_names:
+            try:
+                self._plugin_manager.load_setuptools_entrypoints(
+                    SNOWCLI_COMMAND_PLUGIN_NAMESPACE, plugin_name
+                )
+            except Exception as ex:
+                log_exception(
+                    f"Cannot register plugin [{plugin_name}]: {ex.__str__()}", ex
+                )
+
+    def load_all_registered_plugins(self) -> List[LoadedCommandPlugin]:
+        for plugin_name, plugin in self._plugin_manager.list_name_plugin():
+            self._load_plugin(plugin_name, plugin)
+        return list(self._loaded_plugins.values())
+
+    def _load_plugin(self, plugin_name: str, plugin) -> Optional[LoadedCommandPlugin]:
+        already_loaded_plugin = self._loaded_plugins.get(plugin_name)
+        if already_loaded_plugin:
+            return already_loaded_plugin
+        return self._load_new_plugin(plugin_name, plugin)
+
+    def _load_new_plugin(
+        self, plugin_name: str, plugin
+    ) -> Optional[LoadedCommandPlugin]:
+        loaded_plugin = self._load_plugin_spec(plugin_name, plugin)
+        if not loaded_plugin:
+            return None
+        other_plugin_with_the_same_command_path = self._loaded_command_paths.get(
+            loaded_plugin.command_spec.full_command_path
+        )
+        if other_plugin_with_the_same_command_path:
+            log.error(
+                "Cannot load plugin [%s] "
+                "because it defines the same command [%s] "
+                "as already loaded plugin [%s].",
+                plugin_name,
+                loaded_plugin.command_spec.full_command_path,
+                other_plugin_with_the_same_command_path.plugin_name,
+            )
+            return None
+        self._loaded_plugins[plugin_name] = loaded_plugin
+        self._loaded_command_paths[
+            loaded_plugin.command_spec.full_command_path
+        ] = loaded_plugin
+
+        if self._is_external_plugin(loaded_plugin):
+            log.info("Loaded external plugin: %s", plugin_name)
+
+        return loaded_plugin
+
+    def _load_plugin_spec(
+        self, plugin_name: str, plugin
+    ) -> Optional[LoadedCommandPlugin]:
+        if plugin_name in get_builtin_plugin_name_to_plugin_spec().keys():
+            return self._load_builtin_plugin_spec(plugin_name, plugin)
+        else:
+            return self._load_external_plugin_spec(plugin_name, plugin)
+
+    def _load_builtin_plugin_spec(
+        self, plugin_name: str, plugin
+    ) -> Optional[LoadedCommandPlugin]:
+        command_spec = self._load_command_spec(plugin_name, plugin)
+        if command_spec:
+            return LoadedBuiltInCommandPlugin(
+                plugin_name=plugin_name,
+                command_spec=command_spec,
+            )
+        else:
+            return None
+
+    def _load_external_plugin_spec(
+        self, plugin_name: str, plugin
+    ) -> Optional[LoadedCommandPlugin]:
+        command_spec = self._load_command_spec(plugin_name, plugin)
+        if command_spec:
+            return LoadedExternalCommandPlugin(
+                plugin_name=plugin_name,
+                command_spec=command_spec,
+            )
+        else:
+            return None
+
+    @staticmethod
+    def _load_command_spec(plugin_name: str, plugin) -> Optional[CommandSpec]:
+        try:
+            return plugin.command_spec()
+        except Exception as ex:
+            log_exception(
+                f"Cannot load command specification from plugin [{plugin_name}]: {ex.__str__()}",
+                ex,
+            )
+            return None
+
+    @staticmethod
+    def _is_external_plugin(plugin) -> bool:
+        return isinstance(plugin, LoadedExternalCommandPlugin)
+
+
+def load_only_builtin_command_plugins() -> List[LoadedCommandPlugin]:
+    loader = CommandPluginsLoader()
+    loader.register_builtin_plugins()
+    return loader.load_all_registered_plugins()
+
+
+def load_builtin_and_external_command_plugins(
+    external_plugin_names: List[str],
+) -> List[LoadedCommandPlugin]:
+    loader = CommandPluginsLoader()
+    loader.register_builtin_plugins()
+    loader.register_external_plugins(external_plugin_names)
+    return loader.load_all_registered_plugins()

snowflake/cli/app/commands_registration/commands_registration_with_callbacks.py

@@ -0,0 +1,105 @@
+# 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 dataclasses import dataclass
+from typing import Callable, List
+
+from snowflake.cli.api.plugins.plugin_config import PluginConfigProvider
+from snowflake.cli.app.commands_registration.command_plugins_loader import (
+    load_builtin_and_external_command_plugins,
+    load_only_builtin_command_plugins,
+)
+from snowflake.cli.app.commands_registration.threadsafe import ThreadsafeCounter
+from snowflake.cli.app.commands_registration.typer_registration import (
+    register_commands_from_plugins,
+)
+
+
+@dataclass
+class CommandRegistrationConfig:
+    enable_external_command_plugins: bool
+
+
+class CommandsRegistrationWithCallbacks:
+    def __init__(self, plugin_config_provider: PluginConfigProvider):
+        self._plugin_config_provider = plugin_config_provider
+        self._counter_of_callbacks_required_before_registration: ThreadsafeCounter = (
+            ThreadsafeCounter(0)
+        )
+        self._counter_of_callbacks_invoked_before_registration: ThreadsafeCounter = (
+            ThreadsafeCounter(0)
+        )
+        self._callbacks_after_registration: List[Callable[[], None]] = []
+        self._commands_registration_config: CommandRegistrationConfig = (
+            CommandRegistrationConfig(enable_external_command_plugins=True)
+        )
+        self._commands_already_registered: bool = False
+
+    def register_commands_if_ready_and_not_registered_yet(self):
+        all_required_callbacks_executed = (
+            self._counter_of_callbacks_required_before_registration.value
+            == self._counter_of_callbacks_invoked_before_registration.value
+        )
+        if all_required_callbacks_executed and not self._commands_already_registered:
+            self._register_commands_from_plugins()
+
+    def _register_commands_from_plugins(self) -> None:
+        if self._commands_registration_config.enable_external_command_plugins:
+            self._register_builtin_and_enabled_external_plugin_commands()
+        else:
+            self._register_only_builtin_plugin_commands()
+
+        self._commands_already_registered = True
+        for callback in self._callbacks_after_registration:
+            callback()
+
+    @staticmethod
+    def _register_only_builtin_plugin_commands() -> None:
+        loaded_command_plugins = load_only_builtin_command_plugins()
+        register_commands_from_plugins(loaded_command_plugins)
+
+    def _register_builtin_and_enabled_external_plugin_commands(self):
+        enabled_external_plugins = (
+            self._plugin_config_provider.get_enabled_plugin_names()
+        )
+        loaded_command_plugins = load_builtin_and_external_command_plugins(
+            enabled_external_plugins
+        )
+        register_commands_from_plugins(loaded_command_plugins)
+
+    def disable_external_command_plugins(self):
+        self._commands_registration_config.enable_external_command_plugins = False
+
+    def before(self, callback):
+        def enriched_callback(value):
+            self._counter_of_callbacks_invoked_before_registration.increment()
+            callback(value)
+            self.register_commands_if_ready_and_not_registered_yet()
+
+        self._counter_of_callbacks_required_before_registration.increment()
+        return enriched_callback
+
+    def after(self, callback):
+        def delayed_callback(value):
+            self._callbacks_after_registration.append(lambda: callback(value))
+
+        return delayed_callback
+
+    def reset_running_instance_registration_state(self):
+        self._commands_already_registered = False
+        self._counter_of_callbacks_invoked_before_registration.set(0)
+        self._callbacks_after_registration.clear()
+        self._commands_registration_config.enable_external_command_plugins = True

snowflake/cli/app/commands_registration/exception_logging.py

@@ -0,0 +1,26 @@
+# 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 logging import Logger
+from typing import Callable
+
+from snowflake.cli.api.cli_global_context import cli_context
+
+
+def exception_logging(logger: Logger) -> Callable[[str, Exception], None]:
+    def log_error(msg: str, exception: Exception) -> None:
+        exc_info = exception if cli_context.enable_tracebacks else None
+        logger.error(msg=msg, exc_info=exc_info)
+
+    return log_error

snowflake/cli/app/commands_registration/threadsafe.py

@@ -0,0 +1,48 @@
+# 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 threading
+from typing import Callable, Generic, TypeVar
+
+T = TypeVar("T")
+
+
+class ThreadsafeValue(Generic[T]):
+    def __init__(self, value: T):
+        self._value = value
+        self._lock = threading.Lock()
+
+    def set(self, new_value: T) -> T:  # noqa: A003
+        return self.transform(lambda _: new_value)
+
+    def transform(self, f: Callable[[T], T]) -> T:
+        with self._lock:
+            new_value = f(self._value)
+            self._value = new_value
+            return new_value
+
+    @property
+    def value(self) -> T:
+        with self._lock:
+            return self._value
+
+
+class ThreadsafeCounter(ThreadsafeValue[int]):
+    def increment(self, d=1) -> int:
+        return self.transform(lambda v: v + d)
+
+    def decrement(self, d=1) -> int:
+        return self.increment(-d)