snowflake-cli-labs 2.8.1__py3-none-any.whl → 3.0.0__py3-none-any.whl

snowflake/cli/api/commands/overrideable_parameter.py

@@ -0,0 +1,143 @@
+# 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 abc import ABC, abstractmethod
+from inspect import signature
+from typing import Any, List, Optional, Tuple
+
+import typer
+from click import ClickException
+from snowflake.cli.api.exceptions import IncompatibleParametersError
+
+
+class _OverrideableParameter(ABC):
+    """
+    Class that allows you to generate instances of typer.models.OptionInfo with some default properties while allowing
+    specific values to be overridden.
+
+    Custom parameters:
+    - mutually_exclusive (Tuple[str]|List[str]): A list of parameter names that this Option is not compatible with. If this Option has
+     a truthy value and any of the other parameters in the mutually_exclusive list has a truthy value, a
+     ClickException will be thrown. Note that mutually_exclusive can contain an option's own name but does not require
+     it.
+    """
+
+    def __init__(
+        self,
+        default: Any = ...,
+        *param_decls: str,
+        mutually_exclusive: Optional[List[str] | Tuple[str]] = None,
+        **kwargs,
+    ):
+        self.default = default
+        self.param_decls = param_decls
+        self.mutually_exclusive = mutually_exclusive
+        self.kwargs = kwargs
+
+    def __call__(self, **kwargs) -> typer.models.ParameterInfo:
+        """
+        Returns a typer.models.OptionInfo instance initialized with the specified default values along with any overrides
+        from kwargs. Note that if you are overriding param_decls, you must pass an iterable of strings, you cannot use
+        positional arguments like you can with typer.Option. Does not modify the original instance.
+        """
+        default = kwargs.get("default", self.default)
+        param_decls = kwargs.get("param_decls", self.param_decls)
+        mutually_exclusive = kwargs.get("mutually_exclusive", self.mutually_exclusive)
+        if not isinstance(param_decls, list) and not isinstance(param_decls, tuple):
+            raise TypeError("param_decls must be a list or tuple")
+        passed_kwargs = self.kwargs.copy()
+        passed_kwargs.update(kwargs)
+        if passed_kwargs.get("callback", None) or mutually_exclusive:
+            passed_kwargs["callback"] = self._callback_factory(
+                passed_kwargs.get("callback", None), mutually_exclusive
+            )
+        for non_kwarg in ["default", "param_decls", "mutually_exclusive"]:
+            passed_kwargs.pop(non_kwarg, None)
+
+        return self.get_parameter(default, *param_decls, **passed_kwargs)
+
+    @abstractmethod
+    def get_parameter(
+        self, default: Any = None, *param_decls: str, **kwargs
+    ) -> typer.models.ParameterInfo:
+        pass
+
+    class InvalidCallbackSignature(ClickException):
+        def __init__(self, callback):
+            super().__init__(
+                f"Signature {signature(callback)} is not valid for an OverrideableOption callback function. Must have "
+                f"at most one parameter with each of the following types: (typer.Context, typer.CallbackParam, "
+                f"Any Other Type)"
+            )
+
+    def _callback_factory(
+        self, callback, mutually_exclusive: Optional[List[str] | Tuple[str]]
+    ):
+        callback = callback if callback else lambda x: x
+
+        # inspect existing_callback to make sure signature is valid
+        existing_params = signature(callback).parameters
+        # at most one parameter with each type in [typer.Context, typer.CallbackParam, any other type]
+        limits = [
+            lambda x: x == typer.Context,
+            lambda x: x == typer.CallbackParam,
+            lambda x: x != typer.Context and x != typer.CallbackParam,
+        ]
+        for limit in limits:
+            if len([v for v in existing_params.values() if limit(v.annotation)]) > 1:
+                raise self.InvalidCallbackSignature(callback)
+
+        def generated_callback(ctx: typer.Context, param: typer.CallbackParam, value):
+            if mutually_exclusive:
+                for name in mutually_exclusive:
+                    if value and ctx.params.get(
+                        name, False
+                    ):  # if the current parameter is set to True and a previous parameter is also Truthy
+                        curr_opt = param.opts[0]
+                        other_opt = [x for x in ctx.command.params if x.name == name][
+                            0
+                        ].opts[0]
+                        raise IncompatibleParametersError([curr_opt, other_opt])
+
+            # pass args to existing callback based on its signature (this is how Typer infers callback args)
+            passed_params = {}
+            for existing_param in existing_params:
+                annotation = existing_params[existing_param].annotation
+                if annotation == typer.Context:
+                    passed_params[existing_param] = ctx
+                elif annotation == typer.CallbackParam:
+                    passed_params[existing_param] = param
+                else:
+                    passed_params[existing_param] = value
+            return callback(**passed_params)
+
+        return generated_callback
+
+
+class OverrideableArgument(_OverrideableParameter):
+    def get_parameter(
+        self, default: Any = ..., *param_decls: str, **kwargs
+    ) -> typer.models.ArgumentInfo:
+        return typer.Argument(default, *param_decls, **kwargs)
+
+
+# OverrideableOption doesn't work with flags with type List[Any] and default None, because typer executes the callback
+# function which converts the default value iterating over it, but None is not iterable.
+class OverrideableOption(_OverrideableParameter):
+    def get_parameter(
+        self, default: Any = ..., *param_decls: str, **kwargs
+    ) -> typer.models.OptionInfo:
+        return typer.Option(default, *param_decls, **kwargs)

snowflake/cli/api/commands/snow_typer.py

@@ -20,6 +20,7 @@ from functools import wraps
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import typer
+from click import ClickException
 from snowflake.cli.api.commands.decorators import (
     global_options,
     global_options_with_connection,
@@ -29,10 +30,11 @@ from snowflake.cli.api.commands.execution_metadata import (
     ExecutionStatus,
 )
 from snowflake.cli.api.commands.flags import DEFAULT_CONTEXT_SETTINGS
-from snowflake.cli.api.commands.typer_pre_execute import run_pre_execute_commands
 from snowflake.cli.api.exceptions import CommandReturnTypeError
 from snowflake.cli.api.output.types import CommandResult
 from snowflake.cli.api.sanitizers import sanitize_for_terminal
+from snowflake.cli.api.sql_execution import SqlExecutionMixin
+from snowflake.connector import DatabaseError
 
 log = logging.getLogger(__name__)
 
@@ -46,6 +48,7 @@ class SnowTyper(typer.Typer):
             pretty_exceptions_show_locals=False,
             no_args_is_help=True,
             add_completion=True,
+            rich_markup_mode="markdown",
         )
 
     @staticmethod
@@ -70,6 +73,7 @@ class SnowTyper(typer.Typer):
         requires_global_options: bool = True,
         requires_connection: bool = False,
         is_enabled: Callable[[], bool] | None = None,
+        require_warehouse: bool = False,
         **kwargs,
     ):
         """
@@ -96,15 +100,15 @@ class SnowTyper(typer.Typer):
             def command_callable_decorator(*args, **kw):
                 """Wrapper around command callable. This is what happens at "runtime"."""
                 execution = ExecutionMetadata()
-                self.pre_execute(execution)
+                self.pre_execute(execution, require_warehouse=require_warehouse)
                 try:
                     result = command_callable(*args, **kw)
                     self.process_result(result)
                     execution.complete(ExecutionStatus.SUCCESS)
-                except Exception as err:
+                except BaseException as err:
                     execution.complete(ExecutionStatus.FAILURE)
-                    self.exception_handler(err, execution)
-                    raise
+                    exception = self.exception_handler(err, execution)
+                    raise exception
                 finally:
                     self.post_execute(execution)
 
@@ -115,22 +119,25 @@ class SnowTyper(typer.Typer):
         return custom_command
 
     @staticmethod
-    def pre_execute(execution: ExecutionMetadata):
+    def pre_execute(execution: ExecutionMetadata, require_warehouse: bool = False):
         """
         Callback executed before running any command callable (after context execution).
         Pay attention to make this method safe to use if performed operations are not necessary
         for executing the command in proper way.
         """
-        from snowflake.cli.app.telemetry import log_command_usage
+        from snowflake.cli._app.telemetry import log_command_usage
 
         log.debug("Executing command pre execution callback")
-        run_pre_execute_commands()
         log_command_usage(execution)
+        if require_warehouse and not SqlExecutionMixin().session_has_warehouse():
+            raise ClickException(
+                "The command requires warehouse. No warehouse found in current connection."
+            )
 
     @staticmethod
     def process_result(result):
         """Command result processor"""
-        from snowflake.cli.app.printing import print_result
+        from snowflake.cli._app.printing import print_result
 
         # Because we still have commands like "logs" that do not return anything.
         # We should improve it in future.
@@ -145,10 +152,13 @@ class SnowTyper(typer.Typer):
         """
         Callback executed on command execution error.
         """
-        from snowflake.cli.app.telemetry import log_command_execution_error
+        from snowflake.cli._app.telemetry import log_command_execution_error
 
         log.debug("Executing command exception callback")
         log_command_execution_error(exception, execution)
+        if isinstance(exception, DatabaseError):
+            return ClickException(exception.msg)
+        return exception
 
     @staticmethod
     def post_execute(execution: ExecutionMetadata):
@@ -156,7 +166,7 @@ class SnowTyper(typer.Typer):
         Callback executed after running any command callable. Pay attention to make this method safe to
         use if performed operations are not necessary for executing the command in proper way.
         """
-        from snowflake.cli.app.telemetry import flush_telemetry, log_command_result
+        from snowflake.cli._app.telemetry import flush_telemetry, log_command_result
 
         log.debug("Executing command post execution callback")
         log_command_result(execution)

snowflake/cli/api/commands/utils.py

@@ -0,0 +1,18 @@
+from typing import List, Optional
+
+from click import ClickException
+from snowflake.cli.api.commands.common import Variable
+
+
+def parse_key_value_variables(variables: Optional[List[str]]) -> List[Variable]:
+    """Util for parsing key=value input. Useful for commands accepting multiple input options."""
+    result: List[Variable] = []
+    if not variables:
+        return result
+    for p in variables:
+        if "=" not in p:
+            raise ClickException(f"Invalid variable: '{p}'")
+
+        key, value = p.split("=", 1)
+        result.append(Variable(key.strip(), value.strip()))
+    return result

snowflake/cli/api/config.py

@@ -30,7 +30,10 @@ from snowflake.cli.api.exceptions import (
     UnsupportedConfigSectionTypeError,
 )
 from snowflake.cli.api.secure_path import SecurePath
-from snowflake.cli.api.secure_utils import file_permissions_are_strict
+from snowflake.cli.api.secure_utils import (
+    file_permissions_are_strict,
+    windows_get_not_whitelisted_users_with_access,
+)
 from snowflake.cli.api.utils.types import try_cast_to_bool
 from snowflake.connector.compat import IS_WINDOWS
 from snowflake.connector.config_manager import CONFIG_MANAGER
@@ -77,7 +80,7 @@ class ConnectionConfig:
     warehouse: Optional[str] = None
     role: Optional[str] = None
     authenticator: Optional[str] = None
-    private_key_path: Optional[str] = None
+    private_key_file: Optional[str] = None
     token_file_path: Optional[str] = None
 
     _other_settings: dict = field(default_factory=lambda: {})
@@ -115,7 +118,7 @@ def config_init(config_file: Optional[Path]):
     Initializes the app configuration. Config provided via cli flag takes precedence.
     If config file does not exist we create an empty one.
     """
-    from snowflake.cli.app.loggers import create_initial_loggers
+    from snowflake.cli._app.loggers import create_initial_loggers
 
     if config_file:
         CONFIG_MANAGER.file_path = config_file
@@ -127,12 +130,21 @@ def config_init(config_file: Optional[Path]):
     create_initial_loggers()
 
 
-def add_connection(name: str, connection_config: ConnectionConfig):
-    set_config_value(
-        section=CONNECTIONS_SECTION,
-        key=name,
-        value=connection_config.to_dict_of_all_non_empty_values(),
-    )
+def add_connection_to_proper_file(name: str, connection_config: ConnectionConfig):
+    if CONNECTIONS_FILE.exists():
+        existing_connections = _read_connections_toml()
+        existing_connections.update(
+            {name: connection_config.to_dict_of_all_non_empty_values()}
+        )
+        _update_connections_toml(existing_connections)
+        return CONNECTIONS_FILE
+    else:
+        set_config_value(
+            section=CONNECTIONS_SECTION,
+            key=name,
+            value=connection_config.to_dict_of_all_non_empty_values(),
+        )
+        return CONFIG_MANAGER.file_path
 
 
 _DEFAULT_LOGS_CONFIG = {
@@ -160,6 +172,18 @@ def _read_config_file():
                 message="Bad owner or permissions.*",
                 module="snowflake.connector.config_manager",
             )
+
+            if not file_permissions_are_strict(CONFIG_MANAGER.file_path):
+                users = ", ".join(
+                    windows_get_not_whitelisted_users_with_access(
+                        CONFIG_MANAGER.file_path
+                    )
+                )
+                warnings.warn(
+                    f"Unauthorized users ({users}) have access to configuration file {CONFIG_MANAGER.file_path}.\n"
+                    f'Run `icacls "{CONFIG_MANAGER.file_path}" /deny <USER_ID>:F` on those users to restrict permissions.'
+                )
+
         try:
             CONFIG_MANAGER.read_config()
         except ConfigSourceError as exception:
@@ -281,7 +305,7 @@ def _initialise_config(config_file: Path) -> None:
 
 
 def get_env_variable_name(*path, key: str) -> str:
-    return "SNOWFLAKE_" + "_".join(p.upper() for p in path) + f"_{key.upper()}"
+    return ("_".join(["snowflake", *path, key])).upper()
 
 
 def get_env_value(*path, key: str) -> str | None:
@@ -322,8 +346,6 @@ def _dump_config(conf_file_cache: Dict):
 
 
 def _check_default_config_files_permissions() -> None:
-    if IS_WINDOWS:
-        return
     if CONNECTIONS_FILE.exists() and not file_permissions_are_strict(CONNECTIONS_FILE):
         raise ConfigFileTooWidePermissionsError(CONNECTIONS_FILE)
     if CONFIG_FILE.exists() and not file_permissions_are_strict(CONFIG_FILE):
@@ -346,3 +368,13 @@ def get_feature_flags_section() -> Dict[str, bool | Literal["UNKNOWN"]]:
             return "UNKNOWN"
 
     return {k: _bool_or_unknown(v) for k, v in flags.items()}
+
+
+def _read_connections_toml() -> dict:
+    with open(CONNECTIONS_FILE, "r") as f:
+        return tomlkit.loads(f.read()).unwrap()
+
+
+def _update_connections_toml(connections: dict):
+    with open(CONNECTIONS_FILE, "w") as f:
+        f.write(tomlkit.dumps(connections))

snowflake/cli/api/connections.py

@@ -0,0 +1,216 @@
+# 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 asyncio
+import logging
+import re
+import warnings
+from dataclasses import asdict, dataclass, field, fields, replace
+from pathlib import Path
+from typing import Optional
+
+from snowflake.cli.api.config import get_default_connection_name
+from snowflake.cli.api.exceptions import InvalidSchemaError
+from snowflake.connector import SnowflakeConnection
+from snowflake.connector.compat import IS_WINDOWS
+
+logger = logging.getLogger(__name__)
+
+schema_pattern = re.compile(r".+\..+")
+
+
+@dataclass
+class ConnectionContext:
+    # FIXME: can reduce duplication using config.ConnectionConfig
+    connection_name: Optional[str] = None
+    host: Optional[str] = None
+    port: Optional[int] = None
+    account: Optional[str] = None
+    database: Optional[str] = None
+    role: Optional[str] = None
+    schema: Optional[str] = None
+    user: Optional[str] = None
+    password: Optional[str] = field(default=None, repr=False)
+    authenticator: Optional[str] = None
+    private_key_file: Optional[str] = None
+    warehouse: Optional[str] = None
+    mfa_passcode: Optional[str] = None
+    enable_diag: Optional[bool] = False
+    diag_log_path: Optional[Path] = None
+    diag_allowlist_path: Optional[Path] = None
+    temporary_connection: bool = False
+    session_token: Optional[str] = None
+    master_token: Optional[str] = None
+    token_file_path: Optional[Path] = None
+
+    VALIDATED_FIELD_NAMES = ["schema"]
+
+    def present_values_as_dict(self) -> dict:
+        """Dictionary representation of this ConnectionContext for values that are not None"""
+        return {k: v for (k, v) in asdict(self).items() if v is not None}
+
+    def clone(self) -> ConnectionContext:
+        return replace(self)
+
+    def update(self, **updates):
+        """
+        Given a dictionary of property (key, value) mappings, update properties
+        of this context object with equivalent names to the keys.
+
+        Raises KeyError if a non-property is specified as a key.
+        """
+        field_map = {field.name: field for field in fields(self)}
+        for key, value in updates.items():
+            # ensure key represents a property
+            if key not in field_map:
+                raise KeyError(f"{key} is not a field of {self.__class__.__name__}")
+            setattr(self, key, value)
+
+    def __repr__(self) -> str:
+        """Minimal repr where None values have their keys omitted."""
+        items = [f"{k}={repr(v)}" for (k, v) in self.present_values_as_dict().items()]
+        return f"{self.__class__.__name__}({', '.join(items)})"
+
+    def __setattr__(self, prop, val):
+        """Runs registered validators before setting fields."""
+        if prop in self.VALIDATED_FIELD_NAMES:
+            validate = getattr(self, f"validate_{prop}")
+            validate(val)
+        super().__setattr__(prop, val)
+
+    def validate_schema(self, value: Optional[str]):
+        if (
+            value
+            and not (value.startswith('"') and value.endswith('"'))
+            # if schema is fully qualified name (db.schema)
+            and schema_pattern.match(value)
+        ):
+            raise InvalidSchemaError(value)
+
+    def validate_and_complete(self):
+        """
+        Ensure we can create a connection from this context.
+        """
+        if not self.temporary_connection and not self.connection_name:
+            self.connection_name = get_default_connection_name()
+
+    def build_connection(self):
+        from snowflake.cli._app.snow_connector import connect_to_snowflake
+
+        # Ignore warnings about bad owner or permissions on Windows
+        # Telemetry omit our warning filter from config.py
+        if IS_WINDOWS:
+            warnings.filterwarnings(
+                action="ignore",
+                message="Bad owner or permissions.*",
+                module="snowflake.connector.config_manager",
+            )
+
+        return connect_to_snowflake(**self.present_values_as_dict())
+
+
+class OpenConnectionCache:
+    """
+    A connection cache that transparently manages SnowflakeConnection objects
+    and is keyed by ConnectionContext objects, e.g. cache[ctx].execute_string(...).
+    Connections are automatically closed after CONNECTION_CLEANUP_SEC, but
+    are guaranteed to be open (if config is valid) when returned by the cache.
+    """
+
+    connections: dict[str, SnowflakeConnection]
+    cleanup_futures: dict[str, asyncio.TimerHandle]
+
+    CONNECTION_CLEANUP_SEC: float = 10.0 * 60
+    """Connections are closed this many seconds after the last time they are accessed."""
+
+    def __init__(self):
+        self.connections = {}
+        self.cleanup_futures = {}
+
+    def __getitem__(self, ctx):
+        if not isinstance(ctx, ConnectionContext):
+            raise ValueError(
+                f"Expected key to be ConnectionContext but got {repr(ctx)}"
+            )
+        key = repr(ctx)
+        if not self._has_open_connection(key):
+            self._insert(key, ctx)
+        self._touch(key)
+        return self.connections[key]
+
+    def clear(self):
+        """Closes all connections and resets the cache to its initial state."""
+        connection_keys = list(self.connections.keys())
+        for key in connection_keys:
+            self._cleanup(key)
+
+        # if any orphaned futures still exist, clean them up too
+        for future in self.cleanup_futures.values():
+            future.cancel()
+        self.cleanup_futures.clear()
+
+    def _has_open_connection(self, key: str):
+        return key in self.connections
+
+    def _insert(self, key: str, ctx: ConnectionContext):
+        try:
+            # N.B. build_connection ultimately calls connect_to_snowflake, which
+            # interpolates in connection dicts (from config) and environment variables.
+            # This means that we could return a stale (incorrect) connection for the
+            # given ConnectionContext if get_env_value or get_connection_dict would
+            # have returned different values (i.e. env / config have changed).
+            self.connections[key] = ctx.build_connection()
+        except Exception:
+            logger.debug(
+                "ConnectionCache: failed to connect using %s; not caching.", key
+            )
+            raise
+
+    def _cancel_cleanup_future_if_exists(self, key: str):
+        if key in self.cleanup_futures:
+            self.cleanup_futures.pop(key).cancel()
+
+    def _touch(self, key: str):
+        """
+        Extend the lifetime of the cached connection at the given key.
+        """
+        loop = None
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            # Python 3.11+ will throw when no event loop;
+            # Python 3.10 will issue a DeprecationWarning and return None
+            pass
+
+        if not loop:
+            logger.debug(
+                "ConnectionCache: no event loop; connections will close at exit."
+            )
+            return
+
+        self._cancel_cleanup_future_if_exists(key)
+        self.cleanup_futures[key] = loop.call_later(
+            self.CONNECTION_CLEANUP_SEC, lambda: self._cleanup(key)
+        )
+
+    def _cleanup(self, key: str):
+        """Closes the cached connection at the given key."""
+        if key not in self.connections:
+            logger.debug("Cleaning up connection %s, but not found in cache!", key)
+
+        # doesn't cancel in-flight async queries
+        self._cancel_cleanup_future_if_exists(key)
+        self.connections.pop(key).close()

snowflake/cli/api/console/abc.py

@@ -20,7 +20,10 @@ from typing import Callable, Iterator, Optional
 
 from rich import print as rich_print
 from rich.text import Text
-from snowflake.cli.api.cli_global_context import _CliGlobalContextAccess, cli_context
+from snowflake.cli.api.cli_global_context import (
+    _CliGlobalContextAccess,
+    get_cli_context,
+)
 
 
 class AbstractConsole(ABC):
@@ -34,14 +37,16 @@ class AbstractConsole(ABC):
     """
 
     _print_fn: Callable[[str], None]
-    _cli_context: _CliGlobalContextAccess
     _in_phase: bool
 
     def __init__(self):
         super().__init__()
-        self._cli_context = cli_context
         self._in_phase = False
 
+    @property
+    def _cli_context(self) -> _CliGlobalContextAccess:
+        return get_cli_context()
+
     @property
     def is_silent(self) -> bool:
         """Returns information whether intermediate output is muted."""

snowflake/cli/api/constants.py

@@ -77,3 +77,14 @@ VALID_SCOPES = ["database", "schema", "compute-pool"]
 DEFAULT_SIZE_LIMIT_MB = 128
 
 SF_REST_API_URL_PREFIX = "/api/v2"
+
+PROJECT_TEMPLATE_VARIABLE_OPENING = "<%"
+PROJECT_TEMPLATE_VARIABLE_CLOSING = "%>"
+
+INIT_TEMPLATE_VARIABLE_OPENING = "<!"
+INIT_TEMPLATE_VARIABLE_CLOSING = "!>"
+
+SNOWPARK_SHARED_MIXIN = "snowpark_shared"
+
+DEFAULT_ENV_FILE = "environment.yml"
+DEFAULT_PAGES_DIR = "pages"