From 78594be9535978c815fb8d5950871177576c1823 Mon Sep 17 00:00:00 2001 From: "L. R. Couto" <57910428+lrcouto@users.noreply.github.com> Date: Thu, 1 Aug 2024 12:20:12 -0300 Subject: [PATCH] `kedro-telemetry`: Improve performance by switching to after_command_run (#4014) * Tentative fix for the hook issue Signed-off-by: Laura Couto * Add variable to track if hook was called already Signed-off-by: Laura Couto * Properly set exit code when there is an exception Signed-off-by: Laura Couto * Add test coverage for exception on after_command_hook Signed-off-by: Laura Couto * Add test for the finally block Signed-off-by: Laura Couto * Remove redundant logger configuration on cli.py Signed-off-by: Laura Couto * Add minimal required logger config to sned messages to stderr Signed-off-by: Laura Couto * Call sys.exit only once Signed-off-by: Laura Couto * Lint Signed-off-by: Laura Couto * Remove comment Signed-off-by: Laura Couto * Move exit into exception block Signed-off-by: Laura Couto * Change test Signed-off-by: Laura Couto --------- Signed-off-by: Laura Couto --- kedro/framework/cli/cli.py | 23 ++++++++++++++++++++++- tests/framework/cli/test_cli.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/kedro/framework/cli/cli.py b/kedro/framework/cli/cli.py index 87bd5c4505..0892deed35 100644 --- a/kedro/framework/cli/cli.py +++ b/kedro/framework/cli/cli.py @@ -5,6 +5,7 @@ from __future__ import annotations import importlib +import logging import sys import traceback from collections import defaultdict @@ -38,6 +39,9 @@ v{version} """ +logger = logging.getLogger(__name__) +logger.addHandler(logging.StreamHandler(sys.stderr)) + @click.group(context_settings=CONTEXT_SETTINGS, name="Kedro") @click.version_option(version, "--version", "-V", help="Show version and exit") @@ -154,6 +158,9 @@ def main( self._cli_hook_manager.hook.before_command_run( project_metadata=self._metadata, command_args=args ) + + hook_called = False + try: super().main( args=args, @@ -169,6 +176,8 @@ def main( self._cli_hook_manager.hook.after_command_run( project_metadata=self._metadata, command_args=args, exit_code=exc.code ) + hook_called = True + # When CLI is run outside of a project, project_groups are not registered catch_exception = "click.exceptions.UsageError: No such command" # click convert exception handles to error message @@ -198,7 +207,19 @@ def main( ) click.echo(message) click.echo(hint) - sys.exit(exc.code) + sys.exit(exc.code) + except Exception as error: + logger.error(f"An error has occurred: {error}") + self._cli_hook_manager.hook.after_command_run( + project_metadata=self._metadata, command_args=args, exit_code=1 + ) + hook_called = True + sys.exit(1) + finally: + if not hook_called: + self._cli_hook_manager.hook.after_command_run( + project_metadata=self._metadata, command_args=args, exit_code=0 + ) @property def global_groups(self) -> Sequence[click.MultiCommand]: diff --git a/tests/framework/cli/test_cli.py b/tests/framework/cli/test_cli.py index 317a596924..b9144ebc1b 100644 --- a/tests/framework/cli/test_cli.py +++ b/tests/framework/cli/test_cli.py @@ -2,6 +2,7 @@ from itertools import cycle from os import rename from pathlib import Path +from unittest.mock import MagicMock, patch import click from click.testing import CliRunner @@ -432,6 +433,37 @@ def test_kedro_cli_with_project(self, mocker, fake_metadata): assert "Global commands from Kedro" in result.output assert "Project specific commands from Kedro" in result.output + @patch("sys.exit") + def test_main_hook_exception_handling(self, fake_metadata): + kedro_cli = KedroCLI(fake_metadata.project_path) + kedro_cli._cli_hook_manager.hook.after_command_run = MagicMock() + + with patch.object( + click.CommandCollection, "main", side_effect=Exception("Test Exception") + ): + result = CliRunner().invoke(kedro_cli, []) + + kedro_cli._cli_hook_manager.hook.after_command_run.assert_called_once_with( + project_metadata=kedro_cli._metadata, command_args=[], exit_code=1 + ) + + assert "An error has occurred: Test Exception" in result.output + + @patch("sys.exit") + def test_main_hook_finally_block(self, fake_metadata): + kedro_cli = KedroCLI(fake_metadata.project_path) + kedro_cli._cli_hook_manager.hook.after_command_run = MagicMock() + + # No exception is raised, so it should go to the finally block and call the hook + with patch.object(click.CommandCollection, "main"): + result = CliRunner().invoke(kedro_cli, []) + + kedro_cli._cli_hook_manager.hook.after_command_run.assert_called_once_with( + project_metadata=kedro_cli._metadata, command_args=[], exit_code=0 + ) + + assert "An error has occurred:" not in result.output + @mark.usefixtures("chdir_to_dummy_project") class TestRunCommand: