kedro-telemetry: Improve performance by switching to after_command_run

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]:

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: