Repository: kedro-org/kedro

Issue:
Our `after_dataset_saved` hook does not work with nodes saving more than one datasets
## Description

The code here has a bug:
https://github.com/kedro-org/kedro/blob/08c45b4ada1ec2160d3c2670d8207b4c98330a95/kedro/runner/runner.py#L445-L455
Our `after_dataset_saved` hook has been misbehaving and is always called with the same dataset name and data, regardless of how many other datasets the node will output.

## Context

This bug makes it impossible for users to attach working hooks for `after_dataset_saved` if their nodes have more than one outputs.




Discussion:

Resolving patch:
diff --git a/kedro/runner/runner.py b/kedro/runner/runner.py
--- a/kedro/runner/runner.py
+++ b/kedro/runner/runner.py
@@ -469,17 +469,16 @@ def _synchronous_dataset_load(dataset_name: str):
             node, catalog, inputs, is_async, hook_manager, session_id=session_id
         )
 
-        save_futures = set()
-
+        future_dataset_mapping = {}
         for name, data in outputs.items():
             hook_manager.hook.before_dataset_saved(dataset_name=name, data=data)
-            save_futures.add(pool.submit(catalog.save, name, data))
+            future = pool.submit(catalog.save, name, data)
+            future_dataset_mapping[future] = (name, data)
 
-        for future in as_completed(save_futures):
+        for future in as_completed(future_dataset_mapping):
             exception = future.exception()
             if exception:
                 raise exception
-            hook_manager.hook.after_dataset_saved(
-                dataset_name=name, data=data  # pylint: disable=undefined-loop-variable
-            )
+            name, data = future_dataset_mapping[future]
+            hook_manager.hook.after_dataset_saved(dataset_name=name, data=data)
     return node


Test patch:
diff --git a/tests/framework/session/test_session_extension_hooks.py b/tests/framework/session/test_session_extension_hooks.py
--- a/tests/framework/session/test_session_extension_hooks.py
+++ b/tests/framework/session/test_session_extension_hooks.py
@@ -9,8 +9,14 @@
 from dynaconf.validator import Validator
 
 from kedro.framework.context.context import _convert_paths_to_absolute_posix
-from kedro.framework.hooks import hook_impl
-from kedro.framework.project import _ProjectPipelines, _ProjectSettings, pipelines
+from kedro.framework.hooks import _create_hook_manager, hook_impl
+from kedro.framework.hooks.manager import _register_hooks, _register_hooks_setuptools
+from kedro.framework.project import (
+    _ProjectPipelines,
+    _ProjectSettings,
+    pipelines,
+    settings,
+)
 from kedro.framework.session import KedroSession
 from kedro.io import DataCatalog, MemoryDataSet
 from kedro.pipeline import node, pipeline
@@ -526,6 +532,20 @@ def wait_and_identity(x: Any):
     return node(wait_and_identity, inputs="ds1", outputs="ds2", name="test-node")
 
 
+@pytest.fixture
+def sample_node_multiple_outputs():
+    def wait_and_identity(x: Any, y: Any):
+        time.sleep(0.1)
+        return (x, y)
+
+    return node(
+        wait_and_identity,
+        inputs=["ds1", "ds2"],
+        outputs=["ds3", "ds4"],
+        name="test-node",
+    )
+
+
 class LogCatalog(DataCatalog):
     def load(self, name: str, version: str = None) -> Any:
         dataset = super().load(name=name, version=version)
@@ -537,7 +557,17 @@ def load(self, name: str, version: str = None) -> Any:
 def memory_catalog():
     ds1 = MemoryDataSet({"data": 42})
     ds2 = MemoryDataSet({"data": 42})
-    return LogCatalog({"ds1": ds1, "ds2": ds2})
+    ds3 = MemoryDataSet({"data": 42})
+    ds4 = MemoryDataSet({"data": 42})
+    return LogCatalog({"ds1": ds1, "ds2": ds2, "ds3": ds3, "ds4": ds4})
+
+
+@pytest.fixture
+def hook_manager():
+    hook_manager = _create_hook_manager()
+    _register_hooks(hook_manager, settings.HOOKS)
+    _register_hooks_setuptools(hook_manager, settings.DISABLE_HOOKS_FOR_PLUGINS)
+    return hook_manager
 
 
 class TestAsyncNodeDatasetHooks:
@@ -562,6 +592,32 @@ def test_after_dataset_load_hook_async(
             ["Before dataset loaded", "Catalog load", "After dataset loaded"]
         ).strip("[]") in str(hooks_log_messages).strip("[]")
 
+    def test_after_dataset_load_hook_async_multiple_outputs(
+        self,
+        mocker,
+        memory_catalog,
+        hook_manager,
+        sample_node_multiple_outputs,
+    ):
+        after_dataset_saved_mock = mocker.patch.object(
+            hook_manager.hook, "after_dataset_saved"
+        )
+
+        _run_node_async(
+            node=sample_node_multiple_outputs,
+            catalog=memory_catalog,
+            hook_manager=hook_manager,
+        )
+
+        after_dataset_saved_mock.assert_has_calls(
+            [
+                mocker.call(dataset_name="ds3", data={"data": 42}),
+                mocker.call(dataset_name="ds4", data={"data": 42}),
+            ],
+            any_order=True,
+        )
+        assert after_dataset_saved_mock.call_count == 2
+
 
 class TestKedroContextSpecsHook:
     """Test the behavior of `after_context_created` when updating node inputs."""