🐛 Coerce workgraph node attributes to JSON-safe values on save

src/aiida_workgraph/utils/__init__.py

@@ -18,6 +18,7 @@
 from node_graph.socket_spec import SocketSpec
 from aiida.orm.utils.serialize import serialize
 from aiida_workgraph.orm.utils import deserialize_safe
+import enum
 from copy import deepcopy
 
 LOGGER = logging.getLogger(__name__)
@@ -268,6 +269,58 @@ def clean_pickled_task_executor(tdata: Dict[str, Any]) -> None:
             tdata['error_handlers'][name] = RuntimeExecutor.from_callable(UnavailableExecutor).to_dict()
 
 
+def _ensure_json_safe_key(key: Any) -> Union[str, int, float, bool, None]:
+    """Coerce a dict key into a valid JSON object key.
+
+    JSON object keys must be ``str``/``int``/``float``/``bool``/``None``, and
+    aiida-core's ``clean_value`` does not inspect dict keys, so an invalid key
+    would only fail once it reaches the database.  Enum keys are unwrapped to
+    their ``.value``; anything else that is not already a valid key type is
+    stringified.
+    """
+    if key is None or isinstance(key, (bool, int, float, str)):
+        return key
+    if isinstance(key, enum.Enum):
+        return _ensure_json_safe_key(key.value)
+    return str(key)
+
+
+def _ensure_json_safe(value: Any) -> Any:
+    """Recursively coerce a nested structure into values that can be stored as
+    node attributes.
+
+    Node attributes must survive aiida-core's ``clean_value``, but workgraph
+    data can carry Python objects it rejects (e.g. ``enum.Enum`` members in
+    error-handler kwargs).  Enums are unwrapped to their ``.value``; mapping
+    keys that are not valid JSON object keys are coerced the same way (see
+    :func:`_ensure_json_safe_key`); one-shot iterators are materialized to
+    lists (``clean_value`` would otherwise exhaust them as a side effect of
+    validation); and any remaining value that ``clean_value`` rejects is
+    replaced by its ``str()`` representation.  Values that storage coerces
+    itself (sets to lists, numpy scalars to Python scalars, ``BaseType`` to its
+    value) pass through untouched, so the ``str()`` fallback — lossy by design
+    — only catches what would otherwise fail on store.
+    """
+    from collections.abc import Iterator, Mapping
+
+    from aiida.common.exceptions import ValidationError
+    from aiida.orm.implementation.utils import clean_value
+
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    if isinstance(value, enum.Enum):
+        return _ensure_json_safe(value.value)
+    if isinstance(value, Mapping):
+        return {_ensure_json_safe_key(k): _ensure_json_safe(v) for k, v in value.items()}
+    if isinstance(value, (list, tuple, Iterator)):
+        return [_ensure_json_safe(v) for v in value]
+    try:
+        clean_value(value)
+        return value
+    except ValidationError:
+        return str(value)
+
+
 def save_workgraph_data(node: Union[int, orm.Node], inputs: Dict[str, Any]) -> None:
     from aiida_workgraph.engine.workgraph import WorkGraphSpec
 
@@ -286,9 +339,9 @@ def save_workgraph_data(node: Union[int, orm.Node], inputs: Dict[str, Any]) -> N
     node.task_states = task_states
     node.task_processes = task_processes
     node.task_actions = task_actions
-    node.workgraph_data = wgdata
-    node.workgraph_data_short = short_wgdata
-    node.workgraph_error_handlers = wgdata.pop('error_handlers', {})
+    node.workgraph_data = _ensure_json_safe(wgdata)
+    node.workgraph_data_short = _ensure_json_safe(short_wgdata)
+    node.workgraph_error_handlers = _ensure_json_safe(wgdata.pop('error_handlers', {}))
     graph_inputs = dict(inputs.pop('graph_inputs', {}))
     tasks = dict(inputs.pop('tasks', {}))
     tasks['graph_inputs'] = graph_inputs

tests/test_utils.py

@@ -1,4 +1,226 @@
+import enum
+import json
+
+import numpy as np
+import pytest
 from aiida import orm
+from aiida.common.exceptions import ValidationError
+
+from aiida_workgraph.utils import _ensure_json_safe, _ensure_json_safe_key
+
+
+class _Color(enum.Enum):
+    """Plain Enum: NOT a str/int subclass -> not JSON-serializable."""
+
+    RED = 1
+    BLUE = 2
+
+
+class _SpinChannel(str, enum.Enum):
+    """str-Enum: IS a str instance, so already JSON-serializable."""
+
+    UP = 'up'
+    DOWN = 'down'
+
+
+def test_ensure_json_safe_plain_enum_unwraps_to_value():
+    """A plain Enum (the real bug class) is unwrapped to its ``.value``."""
+    assert _ensure_json_safe(_Color.RED) == 1
+    assert _ensure_json_safe(_Color.BLUE) == 2
+    # nested exactly as an enum default would sit inside wgdata
+    wgdata = {'tasks': {'t': {'inputs': {'color': {'value': _Color.RED}}}}}
+    assert _ensure_json_safe(wgdata) == {'tasks': {'t': {'inputs': {'color': {'value': 1}}}}}
+
+
+def test_ensure_json_safe_dict_enum_key_is_coerced():
+    """Enum dict keys are coerced (regression: keys were previously ignored)."""
+    result = _ensure_json_safe({_Color.RED: 'x', _Color.BLUE: 'y'})
+    assert result == {1: 'x', 2: 'y'}
+    # the whole thing must now be JSON-serializable (keys included)
+    assert json.dumps(result) == '{"1": "x", "2": "y"}'
+
+
+def test_ensure_json_safe_non_enum_key_stringified():
+    """A non-enum, non-primitive key is stringified so json.dumps succeeds."""
+
+    class _Key:
+        def __str__(self):
+            return 'k'
+
+    assert _ensure_json_safe_key(_Key()) == 'k'
+    assert _ensure_json_safe({_Key(): 1}) == {'k': 1}
+
+
+@pytest.mark.parametrize(
+    'factory',
+    [
+        pytest.param(lambda: {1, 2, 3}, id='set'),
+        pytest.param(lambda: frozenset({1, 2}), id='frozenset'),
+        pytest.param(lambda: np.int64(7), id='numpy-int'),
+        pytest.param(lambda: orm.Int(3), id='orm-int'),
+    ],
+)
+def test_ensure_json_safe_preserves_clean_value_coercible(factory):
+    """Values that ``clean_value`` coerces itself pass through untouched.
+
+    The helper must not pre-empt storage's own coercion: a set becomes a list
+    at store time, a numpy scalar a Python scalar, a ``BaseType`` its value —
+    none of them may be stringified by the helper.  (Factories, not values:
+    ``orm.Int`` needs a loaded profile, which is not available at collection
+    time when parametrize arguments are evaluated.)
+    """
+    value = factory()
+    assert _ensure_json_safe(value) is value
+
+
+def test_ensure_json_safe_coercible_values_store_end_to_end():
+    """Pass-through values are coerced by storage itself and round-trip."""
+    node = orm.WorkflowNode()
+    node.base.attributes.set('data', _ensure_json_safe({'tags': {1, 2, 3}, 'n': np.int64(7)}))
+    node.store()
+    stored = orm.load_node(node.pk).base.attributes.get('data')
+    assert sorted(stored['tags']) == [1, 2, 3]
+    assert stored['n'] == 7
+
+
+def test_ensure_json_safe_non_dict_mapping_is_coerced():
+    """A non-``dict`` ``Mapping`` takes the mapping branch, keys included.
+
+    ``clean_value`` does not inspect mapping keys, so a bad key inside e.g. a
+    ``MappingProxyType`` would otherwise pass the helper and only fail in the
+    database driver at store time.
+    """
+    from types import MappingProxyType
+
+    result = _ensure_json_safe(MappingProxyType({(1, 2): 'v', _Color.RED: 'w'}))
+    assert result == {'(1, 2)': 'v', 1: 'w'}
+
+    node = orm.WorkflowNode()
+    node.base.attributes.set('data', result)
+    node.store()
+    assert orm.load_node(node.pk).base.attributes.get('data') == {'(1, 2)': 'v', '1': 'w'}
+
+
+@pytest.mark.parametrize(
+    ('factory', 'expected'),
+    [
+        pytest.param(lambda: iter([1, 2]), [1, 2], id='iterator'),
+        pytest.param(lambda: (v for v in (_Color.RED, 3)), [1, 3], id='generator'),
+    ],
+)
+def test_ensure_json_safe_iterator_is_materialized(factory, expected):
+    """A one-shot iterator is materialized to a list, not passed to
+    ``clean_value`` (which would exhaust it as a side effect of validation and
+    leave an empty value to be stored).  Factories, so each run gets a fresh,
+    unconsumed iterator."""
+    assert _ensure_json_safe(factory()) == expected
+
+
+def test_ensure_json_safe_fallback_does_not_unwrap_arbitrary_value_attr():
+    """The fallback is narrowed to Enum: a plain object exposing ``.value`` is
+    NOT silently unwrapped (it is stringified instead)."""
+
+    class _Wrapper:
+        value = 42
+
+    result = _ensure_json_safe(_Wrapper())
+    assert result != 42
+    assert isinstance(result, str)
+    json.dumps(result)  # must not raise
+
+
+def test_ensure_json_safe_output_is_storable():
+    """Whatever the helper returns must always survive ``clean_value``."""
+    from aiida.orm.implementation.utils import clean_value
+
+    payload = {
+        'a_plain_enum': _Color.RED,
+        'a_str_enum': _SpinChannel.DOWN,
+        'a_list': [_Color.BLUE, (1, 2), {_Color.RED: 'nested'}],
+        'a_set': {1, 2},
+        _Color.RED: 'enum-key',
+    }
+    clean_value(_ensure_json_safe(payload))  # must not raise
+
+
+def test_enum_attribute_store_negative_control():
+    """End-to-end at the real storage gate (``clean_value`` + JSONB).
+
+    Negative control: a plain-Enum-bearing dict set as a node attribute raises
+    ``ValidationError`` when the node is stored (``clean_value`` runs at store
+    time); wrapping it in ``_ensure_json_safe`` first lets the store succeed and
+    the value round-trips through the database.
+    """
+    wgdata = {'tasks': {'t': {'inputs': {'color': {'value': _Color.RED}}}}}
+
+    raw_node = orm.WorkflowNode()
+    raw_node.base.attributes.set('workgraph_data', wgdata)
+    with pytest.raises(ValidationError):
+        raw_node.store()
+
+    safe_node = orm.WorkflowNode()
+    safe_node.base.attributes.set('workgraph_data', _ensure_json_safe(wgdata))
+    safe_node.store()
+    stored = orm.load_node(safe_node.pk).base.attributes.get('workgraph_data')
+    assert stored == {'tasks': {'t': {'inputs': {'color': {'value': 1}}}}}
+
+
+def test_save_coerces_graph_level_error_handler_kwargs():
+    """`WorkGraph.save()` must coerce graph-level error-handler ``kwargs``.
+
+    Error-handler ``kwargs`` are copied verbatim into the
+    ``workgraph_error_handlers`` attribute, so a plain-Enum value there reaches
+    the storage gate.  This exercises the wiring of ``_ensure_json_safe`` inside
+    ``save_workgraph_data``, not just the helper in isolation.
+    """
+    from node_graph.error_handler import normalize_error_handlers
+
+    from aiida_workgraph import WorkGraph, task
+
+    @task()
+    def add(x: int = 1, y: int = 2):
+        return x + y
+
+    def handle(task, **kwargs):  # never runs, only stored
+        return 'retrying'
+
+    wg = WorkGraph(
+        'graph_level_handler',
+        error_handlers=normalize_error_handlers(
+            {'h': {'executor': handle, 'exit_codes': [1], 'kwargs': {'color': _Color.RED}}}
+        ),
+    )
+    wg.add_task(add, name='add1')
+    wg.save()
+
+    stored = orm.load_node(wg.process.pk).base.attributes.get('workgraph_error_handlers')
+    assert stored['h']['kwargs']['color'] == 1
+
+
+def test_save_coerces_task_level_error_handler_kwargs():
+    """`WorkGraph.save()` must coerce task-level error-handler ``kwargs``.
+
+    A task-level handler is stored inside the ``workgraph_data`` attribute under
+    that task's spec; a plain-Enum value in its ``kwargs`` must be unwrapped
+    there too.
+    """
+    from aiida_workgraph import WorkGraph, task
+
+    @task()
+    def add(x: int = 1, y: int = 2):
+        return x + y
+
+    def handle(task, **kwargs):  # never runs, only stored
+        return 'retrying'
+
+    wg = WorkGraph('task_level_handler')
+    add1 = wg.add_task(add, name='add1')
+    add1.add_error_handler({'h': {'executor': handle, 'exit_codes': [1], 'kwargs': {'color': _Color.RED}}})
+    wg.save()
+
+    stored = orm.load_node(wg.process.pk).base.attributes.get('workgraph_data')
+    handlers = stored['tasks']['add1']['spec']['attached_error_handlers']
+    assert handlers['h']['kwargs']['color'] == 1
 
 
 def test_get_or_create_code(fixture_localhost):