diff --git a/src/snowflake/cli/api/commands/snow_typer.py b/src/snowflake/cli/api/commands/snow_typer.py new file mode 100644 index 0000000000..edffab99d1 --- /dev/null +++ b/src/snowflake/cli/api/commands/snow_typer.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +import logging +from functools import wraps +from typing import Optional + +import typer +from snowflake.cli.api.commands.decorators import ( + global_options, + global_options_with_connection, +) +from snowflake.cli.api.commands.flags import DEFAULT_CONTEXT_SETTINGS +from snowflake.cli.api.exceptions import CommandReturnTypeError +from snowflake.cli.api.output.types import CommandResult +from snowflake.cli.app.printing import print_result +from snowflake.cli.app.telemetry import flush_telemetry, log_command_usage + +log = logging.getLogger(__name__) + + +class SnowTyper(typer.Typer): + def __init__(self, /, **kwargs): + super().__init__( + **kwargs, + context_settings=DEFAULT_CONTEXT_SETTINGS, + pretty_exceptions_show_locals=False, + ) + + @wraps(typer.Typer.command) + def command( + self, + name: Optional[str] = None, + requires_global_options: bool = True, + requires_connection: bool = False, + **kwargs, + ): + """ + Custom implementation of Typer.command that adds ability to execute additional + logic before and after execution as well as process the result and act on possible + errors. + """ + + def custom_command(command_callable): + """Custom command wrapper similar to Typer.command.""" + if requires_connection: + command_callable = global_options_with_connection(command_callable) + elif requires_global_options: + command_callable = global_options(command_callable) + + @wraps(command_callable) + def command_callable_decorator(*args, **kw): + """Wrapper around command callable. This is what happens at "runtime".""" + self.pre_execute() + try: + result = command_callable(*args, **kw) + return self.process_result(result) + except Exception as err: + self.exception_handler(err) + raise + finally: + self.post_execute() + + return super(SnowTyper, self).command(name=name, **kwargs)( + command_callable_decorator + ) + + return custom_command + + @staticmethod + def pre_execute(): + """ + 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. + """ + log.debug("Executing command pre execution callback") + log_command_usage() + + @staticmethod + def process_result(result): + """Command result processor""" + if not isinstance(result, CommandResult): + raise CommandReturnTypeError(type(result)) + print_result(result) + + @staticmethod + def exception_handler(exception: Exception): + """ + Callback executed on command execution error. + """ + log.debug("Executing command exception callback") + + @staticmethod + def post_execute(): + """ + 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. + """ + log.debug("Executing command post execution callback") + flush_telemetry() diff --git a/src/snowflake/cli/api/utils/error_handling.py b/src/snowflake/cli/api/utils/error_handling.py new file mode 100644 index 0000000000..9f1e14b8ff --- /dev/null +++ b/src/snowflake/cli/api/utils/error_handling.py @@ -0,0 +1,9 @@ +from contextlib import contextmanager + + +@contextmanager +def ignore_exceptions(): + try: + yield + except: + pass diff --git a/src/snowflake/cli/app/main_typer.py b/src/snowflake/cli/app/main_typer.py index 24059a760d..eabcb0574a 100644 --- a/src/snowflake/cli/app/main_typer.py +++ b/src/snowflake/cli/app/main_typer.py @@ -1,9 +1,11 @@ +from __future__ import annotations + import sys +import typer from rich import print as rich_print from snowflake.cli.api.cli_global_context import cli_context from snowflake.cli.api.commands.flags import DEFAULT_CONTEXT_SETTINGS, DebugOption -from typer import Typer def _handle_exception(exception: Exception): @@ -17,7 +19,7 @@ def _handle_exception(exception: Exception): raise SystemExit(1) -class SnowCliMainTyper(Typer): +class SnowCliMainTyper(typer.Typer): """ Top-level SnowCLI Typer. It contains global exception handling. diff --git a/src/snowflake/cli/app/snow_connector.py b/src/snowflake/cli/app/snow_connector.py index 66a7b3fbf8..ccb61fa50d 100644 --- a/src/snowflake/cli/app/snow_connector.py +++ b/src/snowflake/cli/app/snow_connector.py @@ -5,7 +5,6 @@ import os from typing import Dict, Optional -import click import snowflake.connector from click.exceptions import ClickException from snowflake.cli.api.config import get_connection, get_default_connection @@ -13,6 +12,7 @@ InvalidConnectionConfiguration, SnowflakeConnectionError, ) +from snowflake.cli.app.telemetry import command_info from snowflake.connector import SnowflakeConnection from snowflake.connector.errors import DatabaseError, ForbiddenError @@ -47,7 +47,7 @@ def connect_to_snowflake(temporary_connection: bool = False, connection_name: Op # for cases when external browser and json format are used. with contextlib.redirect_stdout(None): return snowflake.connector.connect( - application=_find_command_path(), + application=command_info(), **connection_parameters, ) except ForbiddenError as err: @@ -69,14 +69,6 @@ def _update_connection_details_with_private_key(connection_parameters: Dict): return connection_parameters -def _find_command_path(): - ctx = click.get_current_context(silent=True) - if ctx: - # Example: SNOWCLI.WAREHOUSE.STATUS - return ".".join(["SNOWCLI", *ctx.command_path.split(" ")[1:]]).upper() - return "SNOWCLI" - - def _load_pem_to_der(private_key_path: str) -> bytes: """ Given a private key file path (in PEM format), decode key data into DER diff --git a/src/snowflake/cli/app/telemetry.py b/src/snowflake/cli/app/telemetry.py new file mode 100644 index 0000000000..6c9e8de2c1 --- /dev/null +++ b/src/snowflake/cli/app/telemetry.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +import platform +import sys +from enum import Enum, unique +from typing import Any, Dict, Union + +import click +from snowflake.cli.__about__ import VERSION +from snowflake.cli.api.cli_global_context import cli_context +from snowflake.cli.api.output.formats import OutputFormat +from snowflake.cli.api.utils.error_handling import ignore_exceptions +from snowflake.connector.telemetry import ( + TelemetryData, + TelemetryField, +) +from snowflake.connector.time_util import get_time_millis + + +@unique +class CLITelemetryField(Enum): + # Basic information + SOURCE = "source" + VERSION_CLI = "version_cli" + VERSION_PYTHON = "version_python" + VERSION_OS = "version_os" + # Command execution context + COMMAND = "command" + COMMAND_GROUP = "command_group" + COMMAND_FLAGS = "command_flags" + COMMAND_OUTPUT_TYPE = "command_output_type" + # Information + EVENT = "event" + ERROR_MSG = "error_msg" + ERROR_TYPE = "error_type" + + +class TelemetryEvent(Enum): + CMD_EXECUTION = "executing_command" + + +TelemetryDict = Dict[Union[CLITelemetryField, TelemetryField], Any] + + +def _find_command_info() -> TelemetryDict: + ctx = click.get_current_context() + command_path = ctx.command_path.split(" ")[1:] + return { + CLITelemetryField.COMMAND: command_path, + CLITelemetryField.COMMAND_GROUP: command_path[0], + CLITelemetryField.COMMAND_FLAGS: { + k: ctx.get_parameter_source(k).name # type: ignore[attr-defined] + for k, v in ctx.params.items() + if v # noqa + }, + CLITelemetryField.COMMAND_OUTPUT_TYPE: ctx.params.get( + "format", OutputFormat.TABLE + ).value, + } + + +def command_info() -> str: + info = _find_command_info() + return ("SNOWCLI." + ".".join(info[CLITelemetryField.COMMAND])).upper() + + +def python_version() -> str: + py_ver = sys.version_info + return f"{py_ver.major}.{py_ver.minor}.{py_ver.micro}" + + +class CLITelemetryClient: + def __init__(self, ctx): + self._ctx = ctx + + @staticmethod + def generate_telemetry_data_dict( + telemetry_payload: TelemetryDict, + ) -> Dict[str, Any]: + data = { + CLITelemetryField.SOURCE: "snowcli", + CLITelemetryField.VERSION_CLI: VERSION, + CLITelemetryField.VERSION_OS: platform.platform(), + CLITelemetryField.VERSION_PYTHON: python_version(), + **_find_command_info(), + **telemetry_payload, + } + # To map Enum to string, so we don't have to use .value every time + return {getattr(k, "value", k): v for k, v in data.items()} # type: ignore[arg-type] + + @property + def _telemetry(self): + return self._ctx.connection._telemetry # noqa + + def send(self, payload: TelemetryDict): + if self._telemetry: + message = self.generate_telemetry_data_dict(payload) + telemetry_data = TelemetryData.from_telemetry_data_dict( + from_dict=message, timestamp=get_time_millis() + ) + self._telemetry.try_add_log_to_batch(telemetry_data) + + def flush(self): + self._telemetry.send_batch() + + +_telemetry = CLITelemetryClient(ctx=cli_context) + + +@ignore_exceptions() +def log_command_usage(): + _telemetry.send({TelemetryField.KEY_TYPE: TelemetryEvent.CMD_EXECUTION.value}) + + +@ignore_exceptions() +def flush_telemetry(): + _telemetry.flush() diff --git a/src/snowflake/cli/plugins/connection/commands.py b/src/snowflake/cli/plugins/connection/commands.py index 1d4f725796..f868df79bc 100644 --- a/src/snowflake/cli/plugins/connection/commands.py +++ b/src/snowflake/cli/plugins/connection/commands.py @@ -5,8 +5,8 @@ import typer from click import ClickException from click.types import StringParamType -from snowflake.cli.api.commands.decorators import global_options, with_output -from snowflake.cli.api.commands.flags import DEFAULT_CONTEXT_SETTINGS, ConnectionOption +from snowflake.cli.api.commands.flags import ConnectionOption +from snowflake.cli.api.commands.snow_typer import SnowTyper from snowflake.cli.api.config import ( add_connection, connection_exists, @@ -20,8 +20,7 @@ ) from snowflake.connector.config_manager import CONFIG_MANAGER -app = typer.Typer( - context_settings=DEFAULT_CONTEXT_SETTINGS, +app = SnowTyper( name="connection", help="Manages connections to Snowflake.", ) @@ -45,8 +44,6 @@ def _mask_password(connection_params: dict): @app.command(name="list") -@with_output -@global_options def list_connections(**options) -> CommandResult: """ Lists configured connections. @@ -71,8 +68,6 @@ def callback(value: str): @app.command() -@global_options -@with_output def add( connection_name: str = typer.Option( None, @@ -210,9 +205,7 @@ def add( ) -@app.command() -@global_options -@with_output +@app.command(requires_connection=False) def test(connection: str = ConnectionOption, **options) -> CommandResult: """ Tests the connection to Snowflake. diff --git a/src/snowflake/cli/plugins/sql/commands.py b/src/snowflake/cli/plugins/sql/commands.py index 18b92c27b1..92f47ce296 100644 --- a/src/snowflake/cli/plugins/sql/commands.py +++ b/src/snowflake/cli/plugins/sql/commands.py @@ -2,20 +2,15 @@ from typing import Optional import typer -from snowflake.cli.api.commands.decorators import ( - global_options_with_connection, - with_output, -) +from snowflake.cli.api.commands.snow_typer import SnowTyper from snowflake.cli.api.output.types import CommandResult, MultipleResults, QueryResult from snowflake.cli.plugins.sql.manager import SqlManager # simple Typer with defaults because it won't become a command group as it contains only one command -app = typer.Typer() +app = SnowTyper() -@app.command(name="sql") -@with_output -@global_options_with_connection +@app.command(name="sql", requires_connection=True) def execute_sql( query: Optional[str] = typer.Option( None, diff --git a/tests/__snapshots__/test_snow_typer.ambr b/tests/__snapshots__/test_snow_typer.ambr new file mode 100644 index 0000000000..6bb87d1c67 --- /dev/null +++ b/tests/__snapshots__/test_snow_typer.ambr @@ -0,0 +1,100 @@ +# serializer version: 1 +# name: test_command_with_connection_options + ''' + + Usage: snow cmd_with_connection_options [OPTIONS] NAME + + ╭─ Arguments ──────────────────────────────────────────────────────────────────╮ + │ * name TEXT [default: None] [required] │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + ╭─ Options ────────────────────────────────────────────────────────────────────╮ + │ --help -h Show this message and exit. │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + ╭─ Connection configuration ───────────────────────────────────────────────────╮ + │ --connection,--environment -c TEXT Name of the connection, as defined │ + │ in your `config.toml`. Default: │ + │ `default`. │ + │ --account,--accountname TEXT Name assigned to your Snowflake │ + │ account. Overrides the value │ + │ specified for the connection. │ + │ --user,--username TEXT Username to connect to Snowflake. │ + │ Overrides the value specified for │ + │ the connection. │ + │ --password TEXT Snowflake password. Overrides the │ + │ value specified for the │ + │ connection. │ + │ --authenticator TEXT Snowflake authenticator. Overrides │ + │ the value specified for the │ + │ connection. │ + │ --private-key-path TEXT Snowflake private key path. │ + │ Overrides the value specified for │ + │ the connection. │ + │ --database,--dbname TEXT Database to use. Overrides the │ + │ value specified for the │ + │ connection. │ + │ --schema,--schemaname TEXT Database schema to use. Overrides │ + │ the value specified for the │ + │ connection. │ + │ --role,--rolename TEXT Role to use. Overrides the value │ + │ specified for the connection. │ + │ --warehouse TEXT Warehouse to use. Overrides the │ + │ value specified for the │ + │ connection. │ + │ --temporary-connection -x Uses connection defined with │ + │ command line parameters, instead │ + │ of one defined in config │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + ╭─ Global configuration ───────────────────────────────────────────────────────╮ + │ --format [TABLE|JSON] Specifies the output format. │ + │ [default: TABLE] │ + │ --verbose -v Displays log entries for log levels `info` │ + │ and higher. │ + │ --debug Displays log entries for log levels `debug` │ + │ and higher; debug logs contains additional │ + │ information. │ + │ --silent Turns off intermediate output to console. │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + + + ''' +# --- +# name: test_command_with_global_options + ''' + + Usage: snow cmd_with_global_options [OPTIONS] NAME + + ╭─ Arguments ──────────────────────────────────────────────────────────────────╮ + │ * name TEXT [default: None] [required] │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + ╭─ Options ────────────────────────────────────────────────────────────────────╮ + │ --help -h Show this message and exit. │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + ╭─ Global configuration ───────────────────────────────────────────────────────╮ + │ --format [TABLE|JSON] Specifies the output format. │ + │ [default: TABLE] │ + │ --verbose -v Displays log entries for log levels `info` │ + │ and higher. │ + │ --debug Displays log entries for log levels `debug` │ + │ and higher; debug logs contains additional │ + │ information. │ + │ --silent Turns off intermediate output to console. │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + + + ''' +# --- +# name: test_command_without_any_options + ''' + + Usage: snow simple_cmd [OPTIONS] NAME + + ╭─ Arguments ──────────────────────────────────────────────────────────────────╮ + │ * name TEXT [default: None] [required] │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + ╭─ Options ────────────────────────────────────────────────────────────────────╮ + │ --help -h Show this message and exit. │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + + + ''' +# --- diff --git a/tests/api/__init__.py b/tests/api/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/api/commands/__init__.py b/tests/api/commands/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/api/commands/__snapshots__/test_snow_typer.ambr b/tests/api/commands/__snapshots__/test_snow_typer.ambr new file mode 100644 index 0000000000..6bb87d1c67 --- /dev/null +++ b/tests/api/commands/__snapshots__/test_snow_typer.ambr @@ -0,0 +1,100 @@ +# serializer version: 1 +# name: test_command_with_connection_options + ''' + + Usage: snow cmd_with_connection_options [OPTIONS] NAME + + ╭─ Arguments ──────────────────────────────────────────────────────────────────╮ + │ * name TEXT [default: None] [required] │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + ╭─ Options ────────────────────────────────────────────────────────────────────╮ + │ --help -h Show this message and exit. │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + ╭─ Connection configuration ───────────────────────────────────────────────────╮ + │ --connection,--environment -c TEXT Name of the connection, as defined │ + │ in your `config.toml`. Default: │ + │ `default`. │ + │ --account,--accountname TEXT Name assigned to your Snowflake │ + │ account. Overrides the value │ + │ specified for the connection. │ + │ --user,--username TEXT Username to connect to Snowflake. │ + │ Overrides the value specified for │ + │ the connection. │ + │ --password TEXT Snowflake password. Overrides the │ + │ value specified for the │ + │ connection. │ + │ --authenticator TEXT Snowflake authenticator. Overrides │ + │ the value specified for the │ + │ connection. │ + │ --private-key-path TEXT Snowflake private key path. │ + │ Overrides the value specified for │ + │ the connection. │ + │ --database,--dbname TEXT Database to use. Overrides the │ + │ value specified for the │ + │ connection. │ + │ --schema,--schemaname TEXT Database schema to use. Overrides │ + │ the value specified for the │ + │ connection. │ + │ --role,--rolename TEXT Role to use. Overrides the value │ + │ specified for the connection. │ + │ --warehouse TEXT Warehouse to use. Overrides the │ + │ value specified for the │ + │ connection. │ + │ --temporary-connection -x Uses connection defined with │ + │ command line parameters, instead │ + │ of one defined in config │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + ╭─ Global configuration ───────────────────────────────────────────────────────╮ + │ --format [TABLE|JSON] Specifies the output format. │ + │ [default: TABLE] │ + │ --verbose -v Displays log entries for log levels `info` │ + │ and higher. │ + │ --debug Displays log entries for log levels `debug` │ + │ and higher; debug logs contains additional │ + │ information. │ + │ --silent Turns off intermediate output to console. │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + + + ''' +# --- +# name: test_command_with_global_options + ''' + + Usage: snow cmd_with_global_options [OPTIONS] NAME + + ╭─ Arguments ──────────────────────────────────────────────────────────────────╮ + │ * name TEXT [default: None] [required] │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + ╭─ Options ────────────────────────────────────────────────────────────────────╮ + │ --help -h Show this message and exit. │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + ╭─ Global configuration ───────────────────────────────────────────────────────╮ + │ --format [TABLE|JSON] Specifies the output format. │ + │ [default: TABLE] │ + │ --verbose -v Displays log entries for log levels `info` │ + │ and higher. │ + │ --debug Displays log entries for log levels `debug` │ + │ and higher; debug logs contains additional │ + │ information. │ + │ --silent Turns off intermediate output to console. │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + + + ''' +# --- +# name: test_command_without_any_options + ''' + + Usage: snow simple_cmd [OPTIONS] NAME + + ╭─ Arguments ──────────────────────────────────────────────────────────────────╮ + │ * name TEXT [default: None] [required] │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + ╭─ Options ────────────────────────────────────────────────────────────────────╮ + │ --help -h Show this message and exit. │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + + + ''' +# --- diff --git a/tests/api/commands/test_snow_typer.py b/tests/api/commands/test_snow_typer.py new file mode 100644 index 0000000000..d3a22a1410 --- /dev/null +++ b/tests/api/commands/test_snow_typer.py @@ -0,0 +1,171 @@ +from functools import partial +from unittest import mock +from unittest.mock import MagicMock + +import pytest +import typer + +from snowflake.cli.api.commands.snow_typer import SnowTyper +from snowflake.cli.api.output.types import MessageResult + +from typer.testing import CliRunner + + +def class_factory( + pre_execute=None, + result_handler=None, + exception_handler=None, + post_execute=None, +): + class _CustomTyper(SnowTyper): + @staticmethod + def pre_execute(): + if pre_execute: + pre_execute() + + @staticmethod + def post_execute(): + if post_execute: + post_execute() + + @staticmethod + def process_result(result): + if result_handler: + result_handler(result) + + @staticmethod + def exception_handler(err): + if exception_handler: + exception_handler(err) + + return _CustomTyper + + +def app_factory(typer_cls): + app = typer_cls(name="snow") + + @app.command("simple_cmd", requires_global_options=False, requires_connection=False) + def simple_cmd(name: str = typer.Argument()): + return MessageResult(f"hello {name}") + + @app.command("fail_cmd", requires_global_options=False, requires_connection=False) + def fail_cmd(name: str = typer.Argument()): + raise Exception("err") + + @app.command( + "cmd_with_global_options", + requires_global_options=True, + requires_connection=False, + ) + def cmd_with_global_options(name: str = typer.Argument()): + return MessageResult(f"hello {name}") + + @app.command("cmd_with_connection_options", requires_connection=True) + def cmd_with_connection_options(name: str = typer.Argument()): + return MessageResult(f"hello {name}") + + return app + + +@pytest.fixture +def cli(): + def mock_cli(app): + return partial(CliRunner().invoke, app) + + return mock_cli + + +def test_no_callbacks(cli): + result = cli(app_factory(class_factory()))(["simple_cmd", "Norman"]) + assert result.exit_code == 0, result.output + assert result.output == "" + + +def test_result_callbacks(cli): + result = cli(app_factory(class_factory(result_handler=lambda x: print(x.message))))( + ["simple_cmd", "Norman"] + ) + assert result.exit_code == 0, result.output + assert result.output.strip() == "hello Norman" + + +def test_pre_callback_green_path(cli): + pre_execute = MagicMock() + post_execute = MagicMock() + exception_callback = MagicMock() + result = cli( + app_factory( + class_factory( + pre_execute=pre_execute, + post_execute=post_execute, + exception_handler=exception_callback, + ) + ) + )(["simple_cmd", "Norman"]) + assert result.exit_code == 0, result.output + + assert pre_execute.called + assert post_execute.called + assert not exception_callback.called + + +def test_pre_callback_error_path(cli): + pre_execute = MagicMock() + post_execute = MagicMock() + exception_callback = MagicMock() + result_handler = MagicMock() + + result = cli( + app_factory( + class_factory( + pre_execute=pre_execute, + post_execute=post_execute, + exception_handler=exception_callback, + result_handler=result_handler, + ) + ) + )(["fail_cmd", "Norman"]) + assert result.exit_code == 1, result.output + + assert pre_execute.called + assert post_execute.called + assert not result_handler.called + assert exception_callback.called + assert len(exception_callback.call_args_list) == 1 + + +def test_command_without_any_options(cli, snapshot): + result = cli(app_factory(SnowTyper))(["simple_cmd", "--help"]) + assert result.output == snapshot + + +def test_command_with_global_options(cli, snapshot): + result = cli(app_factory(SnowTyper))(["cmd_with_global_options", "--help"]) + assert result.output == snapshot + + +def test_command_with_connection_options(cli, snapshot): + result = cli(app_factory(SnowTyper))(["cmd_with_connection_options", "--help"]) + assert result.output == snapshot + + +@mock.patch("snowflake.cli.api.commands.snow_typer.log_command_usage") +def test_snow_typer_pre_execute_sends_telemetry(mock_log_command_usage, cli): + result = cli(app_factory(SnowTyper))(["simple_cmd", "Norma"]) + assert result.exit_code == 0 + mock_log_command_usage.assert_called_once_with() + + +@mock.patch("snowflake.cli.api.commands.snow_typer.flush_telemetry") +def test_snow_typer_post_execute_sends_telemetry(mock_flush_telemetry, cli): + result = cli(app_factory(SnowTyper))(["simple_cmd", "Norma"]) + assert result.exit_code == 0 + mock_flush_telemetry.assert_called_once_with() + + +@mock.patch("snowflake.cli.api.commands.snow_typer.print_result") +def test_snow_typer_result_callback_sends_telemetry(mock_print_result, cli): + result = cli(app_factory(SnowTyper))(["simple_cmd", "Norma"]) + assert result.exit_code == 0 + assert mock_print_result.call_count == 1 + assert mock_print_result.call_args.args[0].message == "hello Norma" diff --git a/tests/app/test_telemetry.py b/tests/app/test_telemetry.py new file mode 100644 index 0000000000..eaa7193425 --- /dev/null +++ b/tests/app/test_telemetry.py @@ -0,0 +1,41 @@ +from unittest import mock + +from snowflake.cli.__about__ import VERSION +from snowflake.connector.version import VERSION as DRIVER_VERSION + + +@mock.patch( + "snowflake.cli.app.telemetry.python_version", +) +@mock.patch("snowflake.cli.app.telemetry.platform.platform") +@mock.patch("snowflake.cli.app.telemetry.get_time_millis") +@mock.patch("snowflake.connector.connect") +def test_executing_command_sends_telemetry_data( + mock_conn, mock_time, mock_platform, mock_version, runner +): + mock_time.return_value = "123" + mock_platform.return_value = "FancyOS" + mock_version.return_value = "2.3.4" + + result = runner.invoke(["connection", "test"], catch_exceptions=False) + assert result.exit_code == 0, result.output + + # The method is called with a TelemetryData type, so we cast it to dict for simpler comparison + assert mock_conn.return_value._telemetry.try_add_log_to_batch.call_args.args[ + 0 + ].to_dict() == { + "message": { + "driver_type": "PythonConnector", + "driver_version": ".".join(str(s) for s in DRIVER_VERSION[:3]), + "source": "snowcli", + "version_cli": VERSION, + "version_os": "FancyOS", + "version_python": "2.3.4", + "command": ["connection", "test"], + "command_group": "connection", + "command_flags": {"format": "DEFAULT"}, + "command_output_type": "TABLE", + "type": "executing_command", + }, + "timestamp": "123", + } diff --git a/tests/test_connection.py b/tests/test_connection.py index 82458d6d26..e98d445469 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -235,7 +235,7 @@ def test_connection_test(mock_connect, runner): assert "Host" in result.output assert "Password" not in result.output assert "password" not in result.output - mock_connect.assert_called_once_with(connection_name="full") + mock_connect.assert_called_with(connection_name="full") @mock.patch("snowflake.connector.connect") diff --git a/tests/test_loaded_modules.py b/tests/test_loaded_modules.py index 89f4a29654..4bd390c887 100644 --- a/tests/test_loaded_modules.py +++ b/tests/test_loaded_modules.py @@ -1,8 +1,6 @@ import pytest import sys -from tests.testing_utils.fixtures import * - @pytest.mark.loaded_modules def test_loaded_modules(runner): diff --git a/tests/test_snow_connector.py b/tests/test_snow_connector.py index 4b6ebc92ac..7fff427fce 100644 --- a/tests/test_snow_connector.py +++ b/tests/test_snow_connector.py @@ -23,9 +23,9 @@ def __repr__(self): ], ) @mock.patch("snowflake.connector.connect") -@mock.patch("snowflake.cli.app.snow_connector.click") +@mock.patch("snowflake.cli.app.snow_connector.command_info") def test_command_context_is_passed_to_snowflake_connection( - mock_click, mock_connect, cmd, expected, test_snowcli_config + mock_command_info, mock_connect, cmd, expected, test_snowcli_config ): from snowflake.cli.app.snow_connector import connect_to_snowflake from snowflake.cli.api.config import config_init @@ -34,7 +34,7 @@ def test_command_context_is_passed_to_snowflake_connection( mock_ctx = mock.Mock() mock_ctx.command_path = cmd - mock_click.get_current_context.return_value = mock_ctx + mock_command_info.return_value = expected connect_to_snowflake() diff --git a/tests/testing_utils/result_assertions.py b/tests/testing_utils/result_assertions.py index 16f688829b..b6e9456ba0 100644 --- a/tests/testing_utils/result_assertions.py +++ b/tests/testing_utils/result_assertions.py @@ -4,7 +4,7 @@ def assert_that_result_is_usage_error( result: Result, expected_error_message: str ) -> None: - assert result.exit_code == 2, result.exit_code + assert result.exit_code == 2, result.output assert expected_error_message in result.output, result.output assert isinstance(result.exception, SystemExit) assert "traceback" not in result.output.lower()