Repository: apache/airflow

Issue:
dynamic dataset ref breaks when viewed in UI or when triggered (dagbag.py:_add_dag_from_db)
### Apache Airflow version

2.4.0b1

### What happened

Here's a file which defines three dags.  "source" uses `Operator.partial` to reference either "sink".  I'm not sure if it's supported to do so, but airlflow should at least fail more gracefully than it does.

```python3
from datetime import datetime, timedelta
from time import sleep

from airflow import Dataset
from airflow.decorators import dag
from airflow.operators.dummy import DummyOperator
from airflow.operators.python import PythonOperator

ps = Dataset("partial_static")
p1 = Dataset("partial_dynamic_1")
p2 = Dataset("partial_dynamic_2")
p3 = Dataset("partial_dynamic_3")

def sleep_n(n):
    sleep(n)

@dag(start_date=datetime(1970, 1, 1), schedule=timedelta(days=365 * 30))
def two_kinds_dynamic_source():

    # dataset ref is not dynamic
    PythonOperator.partial(
        task_id="partial_static", python_callable=sleep_n, outlets=[ps]
    ).expand(op_args=[[1], [20], [40]])

    # dataset ref is dynamic
    PythonOperator.partial(
        task_id="partial_dynamic", python_callable=sleep_n
    ).expand_kwargs(
        [
            {"op_args": [1], "outlets": [p1]},
            {"op_args": [20], "outlets": [p2]},
            {"op_args": [40], "outlets": [p3]},
        ]
    )

two_kinds_dynamic_source()


@dag(schedule=[ps], start_date=datetime(1970, 1, 1))
def two_kinds_static_sink():
    DummyOperator(task_id="dummy")

two_kinds_static_sink()


@dag(schedule=[p1, p2, p3], start_date=datetime(1970, 1, 1))
def two_kinds_dynamic_sink():
    DummyOperator(task_id="dummy")

two_kinds_dynamic_sink()

```

Tried to trigger the dag in the browser, saw this traceback instead:

```
Python version: 3.9.13
Airflow version: 2.4.0.dev1640+astro.1
Node: airflow-webserver-6b969cbd87-4q5kh
-------------------------------------------------------------------------------
Traceback (most recent call last):
  File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 2525, in wsgi_app
    response = self.full_dispatch_request()
  File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 1822, in full_dispatch_request
    rv = self.handle_user_exception(e)
  File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 1820, in full_dispatch_request
    rv = self.dispatch_request()
  File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 1796, in dispatch_request
    return self.ensure_sync(self.view_functions[rule.endpoint])(**view_args)
  File "/usr/local/lib/python3.9/site-packages/airflow/www/auth.py", line 46, in decorated
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/airflow/www/decorators.py", line 117, in view_func
    return f(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/airflow/www/decorators.py", line 80, in wrapper
    return f(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/airflow/utils/session.py", line 73, in wrapper
    return func(*args, session=session, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/airflow/www/views.py", line 2532, in grid
    dag = get_airflow_app().dag_bag.get_dag(dag_id, session=session)
  File "/usr/local/lib/python3.9/site-packages/airflow/utils/session.py", line 70, in wrapper
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.9/site-packages/airflow/models/dagbag.py", line 176, in get_dag
    self._add_dag_from_db(dag_id=dag_id, session=session)
  File "/usr/local/lib/python3.9/site-packages/airflow/models/dagbag.py", line 251, in _add_dag_from_db
    dag = row.dag
  File "/usr/local/lib/python3.9/site-packages/airflow/models/serialized_dag.py", line 223, in dag
    dag = SerializedDAG.from_dict(self.data)  # type: Any
  File "/usr/local/lib/python3.9/site-packages/airflow/serialization/serialized_objects.py", line 1220, in from_dict
    return cls.deserialize_dag(serialized_obj['dag'])
  File "/usr/local/lib/python3.9/site-packages/airflow/serialization/serialized_objects.py", line 1197, in deserialize_dag
    setattr(task, k, kwargs_ref.deref(dag))
  File "/usr/local/lib/python3.9/site-packages/airflow/serialization/serialized_objects.py", line 224, in deref
    value = {k: v.deref(dag) if isinstance(v, _XComRef) else v for k, v in self.value.items()}
AttributeError: 'list' object has no attribute 'items'
```

I can also summon a similar traceback by just trying to view the dag in the grid view, or when running `airflow dags trigger`

### What you think should happen instead

If there's something invalid about this dag, it should fail to parse--rather than successfully parsing and then breaking the UI.

I'm a bit uncertain about what should happen in the dag dependency graph when the source dag runs.  The dynamic outlets are not known until runtime, so it's reasonable that they don't show up in the graph.  But what about after the dag runs?
- do they still trigger the "sink" dag even though we didn't know about the dependency up front?
- do we update the dependency graph now that we know about the dynamic dependency?

Because of this error, we don't get far enough to find out.

### How to reproduce

Include the dag above, try to display it in the grid view.

### Operating System

kubernetes-in-docker on MacOS via helm

### Versions of Apache Airflow Providers

n/a

### Deployment

Other 3rd-party Helm chart

### Deployment details

Deployed using [the astronomer helm chart ](https://github.com/astronomer/airflow-chart)and these values:
```yaml
airflow:
  airflowHome: /usr/local/airflow
  airflowVersion: $VERSION
  defaultAirflowRepository: img
  defaultAirflowTag: $TAG
  executor: KubernetesExecutor
  gid: 50000
  images:
    airflow:
      repository: img
  logs:
    persistence:
      enabled: true
      size: 2Gi
```

### Anything else

_No response_

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:

Resolving patch:
diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py
--- a/airflow/serialization/serialized_objects.py
+++ b/airflow/serialization/serialized_objects.py
@@ -17,6 +17,7 @@
 """Serialized DAG and BaseOperator"""
 from __future__ import annotations
 
+import collections.abc
 import datetime
 import enum
 import logging
@@ -24,7 +25,7 @@
 import weakref
 from dataclasses import dataclass
 from inspect import Parameter, signature
-from typing import TYPE_CHECKING, Any, Iterable, NamedTuple, Type
+from typing import TYPE_CHECKING, Any, Collection, Iterable, Mapping, NamedTuple, Type, Union
 
 import cattr
 import lazy_object_proxy
@@ -207,6 +208,26 @@ def deref(self, dag: DAG) -> XComArg:
         return deserialize_xcom_arg(self.data, dag)
 
 
+# These two should be kept in sync. Note that these are intentionally not using
+# the type declarations in expandinput.py so we always remember to update
+# serialization logic when adding new ExpandInput variants. If you add things to
+# the unions, be sure to update _ExpandInputRef to match.
+_ExpandInputOriginalValue = Union[
+    # For .expand(**kwargs).
+    Mapping[str, Any],
+    # For expand_kwargs(arg).
+    XComArg,
+    Collection[Union[XComArg, Mapping[str, Any]]],
+]
+_ExpandInputSerializedValue = Union[
+    # For .expand(**kwargs).
+    Mapping[str, Any],
+    # For expand_kwargs(arg).
+    _XComRef,
+    Collection[Union[_XComRef, Mapping[str, Any]]],
+]
+
+
 class _ExpandInputRef(NamedTuple):
     """Used to store info needed to create a mapped operator's expand input.
 
@@ -215,13 +236,29 @@ class _ExpandInputRef(NamedTuple):
     """
 
     key: str
-    value: _XComRef | dict[str, Any]
+    value: _ExpandInputSerializedValue
+
+    @classmethod
+    def validate_expand_input_value(cls, value: _ExpandInputOriginalValue) -> None:
+        """Validate we've covered all ``ExpandInput.value`` types.
+
+        This function does not actually do anything, but is called during
+        serialization so Mypy will *statically* check we have handled all
+        possible ExpandInput cases.
+        """
 
     def deref(self, dag: DAG) -> ExpandInput:
+        """De-reference into a concrete ExpandInput object.
+
+        If you add more cases here, be sure to update _ExpandInputOriginalValue
+        and _ExpandInputSerializedValue to match the logic.
+        """
         if isinstance(self.value, _XComRef):
             value: Any = self.value.deref(dag)
-        else:
+        elif isinstance(self.value, collections.abc.Mapping):
             value = {k: v.deref(dag) if isinstance(v, _XComRef) else v for k, v in self.value.items()}
+        else:
+            value = [v.deref(dag) if isinstance(v, _XComRef) else v for v in self.value]
         return create_expand_input(self.key, value)
 
 
@@ -663,6 +700,8 @@ def serialize_mapped_operator(cls, op: MappedOperator) -> dict[str, Any]:
         serialized_op = cls._serialize_node(op, include_deps=op.deps != MappedOperator.deps_for(BaseOperator))
         # Handle expand_input and op_kwargs_expand_input.
         expansion_kwargs = op._get_specified_expand_input()
+        if TYPE_CHECKING:  # Let Mypy check the input type for us!
+            _ExpandInputRef.validate_expand_input_value(expansion_kwargs.value)
         serialized_op[op._expand_input_attr] = {
             "type": get_map_type_key(expansion_kwargs),
             "value": cls.serialize(expansion_kwargs.value),


Test patch:
diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py
--- a/tests/serialization/test_dag_serialization.py
+++ b/tests/serialization/test_dag_serialization.py
@@ -1961,7 +1961,62 @@ def test_operator_expand_xcomarg_serde():
 
 
 @pytest.mark.parametrize("strict", [True, False])
-def test_operator_expand_kwargs_serde(strict):
+def test_operator_expand_kwargs_literal_serde(strict):
+    from airflow.models.xcom_arg import PlainXComArg, XComArg
+    from airflow.serialization.serialized_objects import _XComRef
+
+    with DAG("test-dag", start_date=datetime(2020, 1, 1)) as dag:
+        task1 = BaseOperator(task_id="op1")
+        mapped = MockOperator.partial(task_id='task_2').expand_kwargs(
+            [{"a": "x"}, {"a": XComArg(task1)}],
+            strict=strict,
+        )
+
+    serialized = SerializedBaseOperator.serialize(mapped)
+    assert serialized == {
+        '_is_empty': False,
+        '_is_mapped': True,
+        '_task_module': 'tests.test_utils.mock_operators',
+        '_task_type': 'MockOperator',
+        'downstream_task_ids': [],
+        'expand_input': {
+            "type": "list-of-dicts",
+            "value": [
+                {"__type": "dict", "__var": {"a": "x"}},
+                {
+                    "__type": "dict",
+                    "__var": {"a": {'__type': 'xcomref', '__var': {'task_id': 'op1', 'key': 'return_value'}}},
+                },
+            ],
+        },
+        'partial_kwargs': {},
+        'task_id': 'task_2',
+        'template_fields': ['arg1', 'arg2'],
+        'template_ext': [],
+        'template_fields_renderers': {},
+        'operator_extra_links': [],
+        'ui_color': '#fff',
+        'ui_fgcolor': '#000',
+        "_disallow_kwargs_override": strict,
+        '_expand_input_attr': 'expand_input',
+    }
+
+    op = SerializedBaseOperator.deserialize_operator(serialized)
+    assert op.deps is MappedOperator.deps_for(BaseOperator)
+    assert op._disallow_kwargs_override == strict
+
+    # The XComArg can't be deserialized before the DAG is.
+    expand_value = op.expand_input.value
+    assert expand_value == [{"a": "x"}, {"a": _XComRef({"task_id": "op1", "key": XCOM_RETURN_KEY})}]
+
+    serialized_dag: DAG = SerializedDAG.from_dict(SerializedDAG.to_dict(dag))
+
+    resolved_expand_value = serialized_dag.task_dict['task_2'].expand_input.value
+    resolved_expand_value == [{"a": "x"}, {"a": PlainXComArg(serialized_dag.task_dict['op1'])}]
+
+
+@pytest.mark.parametrize("strict", [True, False])
+def test_operator_expand_kwargs_xcomarg_serde(strict):
     from airflow.models.xcom_arg import PlainXComArg, XComArg
     from airflow.serialization.serialized_objects import _XComRef