From 660236a12990c1cdbe8c06b63993c00e83bb48a7 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Mon, 5 Feb 2024 09:33:50 +0100 Subject: [PATCH 1/5] Add SnowTyper together with telemetry --- src/snowflake/cli/api/utils/error_handling.py | 9 ++ src/snowflake/cli/app/main_typer.py | 94 +++++++++++++- src/snowflake/cli/app/snow_connector.py | 12 +- src/snowflake/cli/app/telemetry.py | 119 ++++++++++++++++++ .../cli/plugins/connection/commands.py | 15 +-- src/snowflake/cli/plugins/sql/commands.py | 11 +- tests/test_connection.py | 2 +- tests/test_snow_connector.py | 6 +- tests/testing_utils/result_assertions.py | 2 +- 9 files changed, 234 insertions(+), 36 deletions(-) create mode 100644 src/snowflake/cli/api/utils/error_handling.py create mode 100644 src/snowflake/cli/app/telemetry.py 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..ebeb91e33d --- /dev/null +++ b/src/snowflake/cli/api/utils/error_handling.py @@ -0,0 +1,9 @@ +from contextlib import contextmanager + + +@contextmanager +def safe(): + try: + yield + except: + pass diff --git a/src/snowflake/cli/app/main_typer.py b/src/snowflake/cli/app/main_typer.py index 24059a760d..8fa7fa7e90 100644 --- a/src/snowflake/cli/app/main_typer.py +++ b/src/snowflake/cli/app/main_typer.py @@ -1,9 +1,24 @@ +from __future__ import annotations + +import logging import sys +from functools import wraps +from typing import Optional +import typer from rich import print as rich_print from snowflake.cli.api.cli_global_context import cli_context +from snowflake.cli.api.commands.decorators import ( + global_options, + global_options_with_connection, +) from snowflake.cli.api.commands.flags import DEFAULT_CONTEXT_SETTINGS, DebugOption -from typer import Typer +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__) def _handle_exception(exception: Exception): @@ -17,7 +32,82 @@ def _handle_exception(exception: Exception): raise SystemExit(1) -class SnowCliMainTyper(Typer): +class SnowTyper(typer.Typer): + def __init__(self, *args, **kwargs): + super().__init__( + *args, + **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, + ): + cls_super = super() + + def custom_command(command_callable): + 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): + self.pre_execute_callback() + try: + result = command_callable(*args, **kw) + return self.process_result(result) + except Exception as err: + self.exception_execute_callback(err) + raise + finally: + self.post_execute_callback() + + return cls_super.command(name=name, **kwargs)(command_callable_decorator) + + return custom_command + + @staticmethod + def pre_execute_callback(): + """ + 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_execute_callback(exception: Exception): + """ + Callback executed on command execution error. + """ + log.debug("Executing command exception callback") + + @staticmethod + def post_execute_callback(): + """ + 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() + + +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..bc87ce8b1e --- /dev/null +++ b/src/snowflake/cli/app/telemetry.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +import platform +import sys +from enum import Enum, unique +from typing import Any, Dict, Union + +import click +from rich import print_json +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 safe +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() + ) + print_json(data=telemetry_data.to_dict()) + self._telemetry.try_add_log_to_batch(telemetry_data) + + def flush(self): + self._telemetry.send_batch() + + +_telemetry = CLITelemetryClient(ctx=cli_context) + + +@safe() +def log_command_usage(): + _telemetry.send({TelemetryField.KEY_TYPE: TelemetryEvent.CMD_EXECUTION.value}) + + +@safe() +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..01aa75c484 100644 --- a/src/snowflake/cli/plugins/connection/commands.py +++ b/src/snowflake/cli/plugins/connection/commands.py @@ -5,8 +5,7 @@ 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.config import ( add_connection, connection_exists, @@ -18,10 +17,10 @@ MessageResult, ObjectResult, ) +from snowflake.cli.app.main_typer import SnowTyper 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..3ea1f19c2c 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.output.types import CommandResult, MultipleResults, QueryResult +from snowflake.cli.app.main_typer import SnowTyper 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/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_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() From d71cd99b9008b1e87c6714dc856c5f81a17ae45b Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Mon, 5 Feb 2024 11:55:30 +0100 Subject: [PATCH 2/5] Add tests --- src/snowflake/cli/api/commands/snow_typer.py | 99 +++++++++++++ src/snowflake/cli/app/main_typer.py | 85 ----------- src/snowflake/cli/app/telemetry.py | 2 - .../cli/plugins/connection/commands.py | 2 +- src/snowflake/cli/plugins/sql/commands.py | 2 +- tests/__snapshots__/test_snow_typer.ambr | 100 +++++++++++++ tests/test_loaded_modules.py | 2 - tests/test_snow_typer.py | 140 ++++++++++++++++++ 8 files changed, 341 insertions(+), 91 deletions(-) create mode 100644 src/snowflake/cli/api/commands/snow_typer.py create mode 100644 tests/__snapshots__/test_snow_typer.ambr create mode 100644 tests/test_snow_typer.py 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..678e923d12 --- /dev/null +++ b/src/snowflake/cli/api/commands/snow_typer.py @@ -0,0 +1,99 @@ +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. + """ + cls_super = super() + + 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_callback() + try: + result = command_callable(*args, **kw) + return self.process_result(result) + except Exception as err: + self.exception_execute_callback(err) + raise + finally: + self.post_execute_callback() + + return cls_super.command(name=name, **kwargs)(command_callable_decorator) + + return custom_command + + @staticmethod + def pre_execute_callback(): + """ + 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_execute_callback(exception: Exception): + """ + Callback executed on command execution error. + """ + log.debug("Executing command exception callback") + + @staticmethod + def post_execute_callback(): + """ + 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/app/main_typer.py b/src/snowflake/cli/app/main_typer.py index 8fa7fa7e90..301011e514 100644 --- a/src/snowflake/cli/app/main_typer.py +++ b/src/snowflake/cli/app/main_typer.py @@ -2,21 +2,11 @@ import logging import sys -from functools import wraps -from typing import Optional import typer from rich import print as rich_print from snowflake.cli.api.cli_global_context import cli_context -from snowflake.cli.api.commands.decorators import ( - global_options, - global_options_with_connection, -) from snowflake.cli.api.commands.flags import DEFAULT_CONTEXT_SETTINGS, DebugOption -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__) @@ -32,81 +22,6 @@ def _handle_exception(exception: Exception): raise SystemExit(1) -class SnowTyper(typer.Typer): - def __init__(self, *args, **kwargs): - super().__init__( - *args, - **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, - ): - cls_super = super() - - def custom_command(command_callable): - 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): - self.pre_execute_callback() - try: - result = command_callable(*args, **kw) - return self.process_result(result) - except Exception as err: - self.exception_execute_callback(err) - raise - finally: - self.post_execute_callback() - - return cls_super.command(name=name, **kwargs)(command_callable_decorator) - - return custom_command - - @staticmethod - def pre_execute_callback(): - """ - 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_execute_callback(exception: Exception): - """ - Callback executed on command execution error. - """ - log.debug("Executing command exception callback") - - @staticmethod - def post_execute_callback(): - """ - 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() - - class SnowCliMainTyper(typer.Typer): """ Top-level SnowCLI Typer. diff --git a/src/snowflake/cli/app/telemetry.py b/src/snowflake/cli/app/telemetry.py index bc87ce8b1e..4719a63736 100644 --- a/src/snowflake/cli/app/telemetry.py +++ b/src/snowflake/cli/app/telemetry.py @@ -6,7 +6,6 @@ from typing import Any, Dict, Union import click -from rich import print_json from snowflake.cli.__about__ import VERSION from snowflake.cli.api.cli_global_context import cli_context from snowflake.cli.api.output.formats import OutputFormat @@ -99,7 +98,6 @@ def send(self, payload: TelemetryDict): telemetry_data = TelemetryData.from_telemetry_data_dict( from_dict=message, timestamp=get_time_millis() ) - print_json(data=telemetry_data.to_dict()) self._telemetry.try_add_log_to_batch(telemetry_data) def flush(self): diff --git a/src/snowflake/cli/plugins/connection/commands.py b/src/snowflake/cli/plugins/connection/commands.py index 01aa75c484..f868df79bc 100644 --- a/src/snowflake/cli/plugins/connection/commands.py +++ b/src/snowflake/cli/plugins/connection/commands.py @@ -6,6 +6,7 @@ from click import ClickException from click.types import StringParamType 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, @@ -17,7 +18,6 @@ MessageResult, ObjectResult, ) -from snowflake.cli.app.main_typer import SnowTyper from snowflake.connector.config_manager import CONFIG_MANAGER app = SnowTyper( diff --git a/src/snowflake/cli/plugins/sql/commands.py b/src/snowflake/cli/plugins/sql/commands.py index 3ea1f19c2c..92f47ce296 100644 --- a/src/snowflake/cli/plugins/sql/commands.py +++ b/src/snowflake/cli/plugins/sql/commands.py @@ -2,8 +2,8 @@ from typing import Optional import typer +from snowflake.cli.api.commands.snow_typer import SnowTyper from snowflake.cli.api.output.types import CommandResult, MultipleResults, QueryResult -from snowflake.cli.app.main_typer import SnowTyper 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 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/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_typer.py b/tests/test_snow_typer.py new file mode 100644 index 0000000000..dec2bbc5b5 --- /dev/null +++ b/tests/test_snow_typer.py @@ -0,0 +1,140 @@ +from functools import partial +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 app_factory( + pre_execute=None, + result_handler=None, + exception_handler=None, + post_execute=None, +): + class _CustomTyper(SnowTyper): + @staticmethod + def pre_execute_callback(): + if pre_execute: + pre_execute() + + @staticmethod + def post_execute_callback(): + if post_execute: + post_execute() + + @staticmethod + def process_result(result): + if result_handler: + result_handler(result) + + @staticmethod + def exception_execute_callback(err): + if exception_handler: + exception_handler(err) + + app = _CustomTyper(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())(["simple_cmd", "Norman"]) + assert result.exit_code == 0, result.output + assert result.output == "" + + +def test_result_callbacks(cli): + result = cli(app_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( + 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( + 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())(["simple_cmd", "--help"]) + assert result.output == snapshot + + +def test_command_with_global_options(cli, snapshot): + result = cli(app_factory())(["cmd_with_global_options", "--help"]) + assert result.output == snapshot + + +def test_command_with_connection_options(cli, snapshot): + result = cli(app_factory())(["cmd_with_connection_options", "--help"]) + assert result.output == snapshot From abc339dbe1ffc36cd321ac6776bf42542b4d7b58 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Tue, 6 Feb 2024 10:12:36 +0100 Subject: [PATCH 3/5] fixup! Add tests --- src/snowflake/cli/api/commands/snow_typer.py | 5 +- src/snowflake/cli/api/utils/error_handling.py | 2 +- src/snowflake/cli/app/main_typer.py | 3 - src/snowflake/cli/app/telemetry.py | 6 +- tests/api/__init__.py | 0 tests/api/commands/__init__.py | 0 tests/{ => api/commands}/test_snow_typer.py | 59 ++++++++++++++----- tests/app/test_telemetry.py | 41 +++++++++++++ 8 files changed, 93 insertions(+), 23 deletions(-) create mode 100644 tests/api/__init__.py create mode 100644 tests/api/commands/__init__.py rename tests/{ => api/commands}/test_snow_typer.py (63%) create mode 100644 tests/app/test_telemetry.py diff --git a/src/snowflake/cli/api/commands/snow_typer.py b/src/snowflake/cli/api/commands/snow_typer.py index 678e923d12..80f7aec81c 100644 --- a/src/snowflake/cli/api/commands/snow_typer.py +++ b/src/snowflake/cli/api/commands/snow_typer.py @@ -39,7 +39,6 @@ def command( logic before and after execution as well as process the result and act on possible errors. """ - cls_super = super() def custom_command(command_callable): """Custom command wrapper similar to Typer.command.""" @@ -61,7 +60,9 @@ def command_callable_decorator(*args, **kw): finally: self.post_execute_callback() - return cls_super.command(name=name, **kwargs)(command_callable_decorator) + return super(SnowTyper, self).command(name=name, **kwargs)( + command_callable_decorator + ) return custom_command diff --git a/src/snowflake/cli/api/utils/error_handling.py b/src/snowflake/cli/api/utils/error_handling.py index ebeb91e33d..9f1e14b8ff 100644 --- a/src/snowflake/cli/api/utils/error_handling.py +++ b/src/snowflake/cli/api/utils/error_handling.py @@ -2,7 +2,7 @@ @contextmanager -def safe(): +def ignore_exceptions(): try: yield except: diff --git a/src/snowflake/cli/app/main_typer.py b/src/snowflake/cli/app/main_typer.py index 301011e514..eabcb0574a 100644 --- a/src/snowflake/cli/app/main_typer.py +++ b/src/snowflake/cli/app/main_typer.py @@ -1,6 +1,5 @@ from __future__ import annotations -import logging import sys import typer @@ -8,8 +7,6 @@ from snowflake.cli.api.cli_global_context import cli_context from snowflake.cli.api.commands.flags import DEFAULT_CONTEXT_SETTINGS, DebugOption -log = logging.getLogger(__name__) - def _handle_exception(exception: Exception): if cli_context.enable_tracebacks: diff --git a/src/snowflake/cli/app/telemetry.py b/src/snowflake/cli/app/telemetry.py index 4719a63736..6c9e8de2c1 100644 --- a/src/snowflake/cli/app/telemetry.py +++ b/src/snowflake/cli/app/telemetry.py @@ -9,7 +9,7 @@ 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 safe +from snowflake.cli.api.utils.error_handling import ignore_exceptions from snowflake.connector.telemetry import ( TelemetryData, TelemetryField, @@ -107,11 +107,11 @@ def flush(self): _telemetry = CLITelemetryClient(ctx=cli_context) -@safe() +@ignore_exceptions() def log_command_usage(): _telemetry.send({TelemetryField.KEY_TYPE: TelemetryEvent.CMD_EXECUTION.value}) -@safe() +@ignore_exceptions() def flush_telemetry(): _telemetry.flush() 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/test_snow_typer.py b/tests/api/commands/test_snow_typer.py similarity index 63% rename from tests/test_snow_typer.py rename to tests/api/commands/test_snow_typer.py index dec2bbc5b5..e87e23eac9 100644 --- a/tests/test_snow_typer.py +++ b/tests/api/commands/test_snow_typer.py @@ -1,4 +1,5 @@ from functools import partial +from unittest import mock from unittest.mock import MagicMock import pytest @@ -10,7 +11,7 @@ from typer.testing import CliRunner -def app_factory( +def class_factory( pre_execute=None, result_handler=None, exception_handler=None, @@ -37,7 +38,11 @@ def exception_execute_callback(err): if exception_handler: exception_handler(err) - app = _CustomTyper(name="snow") + 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()): @@ -71,13 +76,13 @@ def mock_cli(app): def test_no_callbacks(cli): - result = cli(app_factory())(["simple_cmd", "Norman"]) + 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(result_handler=lambda x: print(x.message)))( + result = cli(app_factory(class_factory(result_handler=lambda x: print(x.message))))( ["simple_cmd", "Norman"] ) assert result.exit_code == 0, result.output @@ -90,9 +95,11 @@ def test_pre_callback_green_path(cli): exception_callback = MagicMock() result = cli( app_factory( - pre_execute=pre_execute, - post_execute=post_execute, - exception_handler=exception_callback, + class_factory( + pre_execute=pre_execute, + post_execute=post_execute, + exception_handler=exception_callback, + ) ) )(["simple_cmd", "Norman"]) assert result.exit_code == 0, result.output @@ -110,10 +117,12 @@ def test_pre_callback_error_path(cli): result = cli( app_factory( - pre_execute=pre_execute, - post_execute=post_execute, - exception_handler=exception_callback, - result_handler=result_handler, + 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 @@ -126,15 +135,37 @@ def test_pre_callback_error_path(cli): def test_command_without_any_options(cli, snapshot): - result = cli(app_factory())(["simple_cmd", "--help"]) + result = cli(app_factory(SnowTyper))(["simple_cmd", "--help"]) assert result.output == snapshot def test_command_with_global_options(cli, snapshot): - result = cli(app_factory())(["cmd_with_global_options", "--help"]) + 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())(["cmd_with_connection_options", "--help"]) + 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", + } From f9ae29dd60a36441574cffca5350957546b16e10 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Tue, 6 Feb 2024 10:47:24 +0100 Subject: [PATCH 4/5] fixup! fixup! Add tests --- src/snowflake/cli/api/commands/snow_typer.py | 12 ++++++------ tests/api/commands/test_snow_typer.py | 6 +++--- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/snowflake/cli/api/commands/snow_typer.py b/src/snowflake/cli/api/commands/snow_typer.py index 80f7aec81c..edffab99d1 100644 --- a/src/snowflake/cli/api/commands/snow_typer.py +++ b/src/snowflake/cli/api/commands/snow_typer.py @@ -50,15 +50,15 @@ def custom_command(command_callable): @wraps(command_callable) def command_callable_decorator(*args, **kw): """Wrapper around command callable. This is what happens at "runtime".""" - self.pre_execute_callback() + self.pre_execute() try: result = command_callable(*args, **kw) return self.process_result(result) except Exception as err: - self.exception_execute_callback(err) + self.exception_handler(err) raise finally: - self.post_execute_callback() + self.post_execute() return super(SnowTyper, self).command(name=name, **kwargs)( command_callable_decorator @@ -67,7 +67,7 @@ def command_callable_decorator(*args, **kw): return custom_command @staticmethod - def pre_execute_callback(): + 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 @@ -84,14 +84,14 @@ def process_result(result): print_result(result) @staticmethod - def exception_execute_callback(exception: Exception): + def exception_handler(exception: Exception): """ Callback executed on command execution error. """ log.debug("Executing command exception callback") @staticmethod - def post_execute_callback(): + 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. diff --git a/tests/api/commands/test_snow_typer.py b/tests/api/commands/test_snow_typer.py index e87e23eac9..d3a22a1410 100644 --- a/tests/api/commands/test_snow_typer.py +++ b/tests/api/commands/test_snow_typer.py @@ -19,12 +19,12 @@ def class_factory( ): class _CustomTyper(SnowTyper): @staticmethod - def pre_execute_callback(): + def pre_execute(): if pre_execute: pre_execute() @staticmethod - def post_execute_callback(): + def post_execute(): if post_execute: post_execute() @@ -34,7 +34,7 @@ def process_result(result): result_handler(result) @staticmethod - def exception_execute_callback(err): + def exception_handler(err): if exception_handler: exception_handler(err) From 577d5758e20352dc65334dc3641ef3d3c10b6738 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Tue, 6 Feb 2024 14:33:22 +0100 Subject: [PATCH 5/5] Add missing snapshots --- .../__snapshots__/test_snow_typer.ambr | 100 ++++++++++++++++++ 1 file changed, 100 insertions(+) create mode 100644 tests/api/commands/__snapshots__/test_snow_typer.ambr 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. │ + ╰──────────────────────────────────────────────────────────────────────────────╯ + + + ''' +# ---