kaxil commented on code in PR #74347:
URL: https://github.com/apache/airflow/pull/74347#discussion_r4222267215
##########
airflow-core/src/airflow/api_fastapi/core_api/services/public/event_logs.py:
##########
@@ -16,14 +16,29 @@
# under the License.
from __future__ import annotations
-from sqlalchemy import inspect
+from sqlalchemy import case, inspect, select
from sqlalchemy.orm.attributes import set_committed_value
from airflow.api_fastapi.core_api.datamodels.event_logs import EventLogResponse
from airflow.models import Log
+from airflow.models.task_coordinates import public_map_index_expression
+from airflow.models.taskinstance import TaskInstance
-def event_log_to_response(event_log: Log) -> EventLogResponse:
+def event_log_public_map_index():
+ """Resolve an attributed execution's public index without guessing once
its row has been purged."""
+ # The Table, unlike the ORM entity, is not narrowed to working_set rows,
so archived tries still resolve.
+ task_instances = TaskInstance.__table__
+ attributed = (
+ select(public_map_index_expression(task_instances))
+ .where(task_instances.c.id == Log.task_instance_id)
+ .correlate(Log)
+ .scalar_subquery()
+ )
+ return case((Log.task_instance_id.is_(None), Log.map_index),
else_=attributed).label("public_map_index")
Review Comment:
Asset events store the public map index when they're written
(`assets/manager.py:344-352`), but event logs keep the stored index and resolve
the public one here at read time. `db clean` purges `task_instance` by
`start_date` and `log` by `dttm`, so a log row can outlive its attempt and then
returns `map_index: null`, even for a plain unmapped task whose stored -1 was
already right. It also puts this CASE in the `map_index` filter, where it can't
use `idx_log_task_instance`. Would writing the public index onto `Log` when the
row is created, the way asset events do, work here?
##########
airflow-core/src/airflow/api_fastapi/core_api/routes/public/event_logs.py:
##########
@@ -130,8 +136,14 @@ def get_event_logs(
# Exact match filters (for backward compatibility)
dag_id: Annotated[FilterParam[str | None],
Depends(filter_param_factory(Log.dag_id, str | None))],
task_id: Annotated[FilterParam[str | None],
Depends(filter_param_factory(Log.task_id, str | None))],
+ task_instance_id: Annotated[
Review Comment:
The list endpoints filter on `task_instance_id` here and
`source_task_instance_id` in `assets.py`, but `EventLogResponse` and
`AssetEventResponse` still don't return either field. Two loop passes of an
unmapped task on try 1 come back with the same `task_id`, `map_index=-1` and
`try_number=1`, so a client listing events can't tell which pass wrote a row
without one request per attempt. Could both carry it as an optional field, like
`AssetStateStoreLastUpdatedBy` does?
##########
airflow-core/src/airflow/models/log.py:
##########
@@ -73,15 +73,16 @@ class Log(Base):
task_instance: Mapped[TaskInstance | None] = relationship(
"TaskInstance",
viewonly=True,
- foreign_keys=[dag_id, task_id, run_id, map_index],
- primaryjoin="and_(Log.dag_id == TaskInstance.dag_id, Log.task_id ==
TaskInstance.task_id, Log.run_id == TaskInstance.run_id, Log.map_index ==
TaskInstance.region_index, TaskInstance.working_set.is_(True))",
+ foreign_keys=[task_instance_id],
+ primaryjoin="Log.task_instance_id == TaskInstance.id",
Review Comment:
`task_display_name` still resolves only through `task_instance_id`, so any
row without it comes back with no display name, where 3.3 filled it from the
coordinates. That's not only rows written before 0142 (which adds the column
without a backfill). `BaseExecutor.log_task_event` sets the id only for
`TaskInstanceUuid` keys, and the released Celery and Kubernetes providers don't
set `supports_task_instance_uuid`, so their executor events keep landing
without it after the upgrade. Since the display name is task-level, a fallback
for rows where `task_instance_id` is null that looks it up by
`dag_id`/`task_id` would cover both cases.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -1466,6 +1495,14 @@ def is_dag_run_asset_triggered(
return dag_run.run_type == DagRunType.DATASET_TRIGGERED # type:
ignore[attr-defined] # This attr is available on AF2, but mypy can't see it
+def get_regional_task_instance_run_id(task_instance: TaskInstance |
RuntimeTaskInstance) -> str | None:
+ if AIRFLOW_V_3_4_PLUS:
+ region_id = getattr(task_instance, "region_id", None)
+ if isinstance(region_id, UUID) and region_id.int != 0:
Review Comment:
This still returns `str(task_instance.id)` for any non-sentinel region, and
`DagRun` creates a region for every task that needs expansion
(`dagrun.py:2198`), so plain `.expand()` tasks still move from the
deterministic `build_task_instance_ol_run_id` value to the TI UUID. Lineage
consumers that rebuild run ids from dag, task, try, logical date and map index
will stop matching those runs, and the provider changelog doesn't mention it.
If that isn't intended, a region whose `node_id` is the task itself and that
has no `parent_region_id` is a plain expansion and could keep the deterministic
id.
##########
providers/openlineage/src/airflow/providers/openlineage/utils/utils.py:
##########
@@ -664,15 +666,40 @@ def get_dag_documentation(dag: DAG | SerializedDAG |
None) -> tuple[str | None,
return None, None
+def get_task_instance_map_index(task_instance: TaskInstance |
RuntimeTaskInstance) -> int:
+ if AIRFLOW_V_3_4_PLUS and isinstance(task_instance, TaskInstance):
+ if task_instance.region_id.int == 0:
+ return task_instance.region_index
+
+ from sqlalchemy import select
+ from sqlalchemy.orm import object_session
+
+ from airflow.models.dynamic_region import DynamicRegion
Review Comment:
This reimplements core's public map index rule
(`public_map_index_expression` in `task_coordinates.py`) against
`DynamicRegion`, which isn't part of the public interface, so the two can drift
without either test suite noticing. It also runs a query for every regional ORM
TI. Could core expose the public map index on the task instance (or as a small
public helper) and have the provider call that?
##########
providers/openlineage/tests/unit/openlineage/utils/test_utils.py:
##########
@@ -3833,9 +3842,127 @@ def test_validate_uuid_invalid():
assert is_valid_uuid(uuid_str) is False
[email protected]_test
[email protected](not AIRFLOW_V_3_4_PLUS, reason="Region projection requires
Airflow 3.4")
[email protected]("mapped", [False, True])
+def test_orm_lineage_facets_project_public_index(dag_maker, session, mapped):
+ with dag_maker(serialized=True):
+ if mapped:
+ PythonOperator.partial(task_id="work",
python_callable=list).expand(op_kwargs=[{}, {}, {}])
+ else:
+ EmptyOperator(task_id="work")
+ run = dag_maker.create_dagrun()
+ if mapped:
+ ti = next(ti for ti in run.task_instances if ti.region_index == 2)
+ else:
+ region = DynamicRegion(dag_id=run.dag_id, run_id=run.run_id,
node_id="loop")
+ session.add(region)
+ session.flush()
+ ti = run.task_instances[0]
+ ti.region_id = region.id
+ ti.region_index = 2
+ session.flush()
+
+ with patch.object(DBDagBag, "get_dag", autospec=True) as get_dag:
+ info = dict(TaskInstanceInfo(ti))
+ facets = get_airflow_mapped_task_facet(ti)
+
+ get_dag.assert_not_called()
+ assert info["map_index"] == (2 if mapped else None)
+ if mapped:
+ assert facets["airflow_mappedTask"].mapIndex == 2
+ else:
+ assert facets == {}
+ assert ti in session
+
+
class TestExtractOlInfoFromAssetEvent:
"""Tests for _extract_ol_info_from_asset_event function."""
+ @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region attribution
requires Airflow3.4")
+ @pytest.mark.parametrize("map_index", [-1, 3])
+ def test_regional_source_run_identity_uses_execution_uuid(self, map_index):
+ region_id = uuid4()
+ logical_date = datetime.datetime(2024, 1, 1, tzinfo=datetime.UTC)
+ ti = SimpleNamespace(
+ id=uuid4(),
+ dag_id="source",
+ task_id="task",
+ try_number=1,
+ region_id=region_id,
+ region_index=3,
+ map_index=map_index,
+ )
+ event = SimpleNamespace(
+ source_task_instance_id=ti.id,
+ source_task_instance=ti,
+ source_dag_id="source",
+ source_task_id="task",
+ source_map_index=map_index,
+ source_dag_run=SimpleNamespace(logical_date=logical_date),
+ )
+ result = _extract_ol_info_from_asset_event(event)
+
+ expected = {"job_name": "source.task", "job_namespace": namespace(),
"run_id": str(ti.id)}
+ assert result == expected
+
+ @pytest.mark.db_test
+ @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Region attribution
requires Airflow3.4")
+ def test_live_regional_asset_sources_need_no_graph_lookup(self, dag_maker,
session):
+ with dag_maker():
+ EmptyOperator(task_id="work")
+ run = dag_maker.create_dagrun()
+ root = DynamicRegion(dag_id=run.dag_id, run_id=run.run_id,
node_id="mapped")
+ loop = DynamicRegion(dag_id=run.dag_id, run_id=run.run_id,
node_id="loop")
+ session.add_all([root, loop])
+ session.flush()
+ nested = DynamicRegion(
+ dag_id=run.dag_id,
+ run_id=run.run_id,
+ node_id="nested",
+ parent_region_id=loop.id,
+ parent_region_index=2,
+ )
+ session.add(nested)
+ session.flush()
+ events = []
+ for region, task_id in [(root, "mapped"), (loop, "work"), (nested,
"nested")]:
+ ti = SimpleNamespace(
+ id=uuid4(),
+ dag_id=run.dag_id,
+ task_id=task_id,
+ try_number=1,
+ region_id=region.id,
+ region_index=2,
+ )
+ for index in range(10):
+ events.append(
+ SimpleNamespace(
+ id=len(events),
+ asset_id=index,
+ uri=f"test://asset{index}",
+ extra={},
+ source_task_instance=ti,
+ source_dag_id=ti.dag_id,
+ source_task_id=ti.task_id,
+ source_run_id=run.run_id,
+ source_map_index=-1 if task_id == "work" else 2,
+
source_dag_run=SimpleNamespace(logical_date=run.run_after),
+ )
+ )
+ session.commit()
+
+ with assert_queries_count(0):
+ dependencies = _get_ol_job_dependencies_from_asset_events(events)
+
+ assert len(dependencies) == 3
+ assert {item["job_name"] for item in dependencies if "run_id" in item}
== {
Review Comment:
The code under test only checks `region_id.int != 0` on the
`SimpleNamespace` TIs, so the `DynamicRegion` rows aren't read, and this
assertion only covers job names. A wrong `run_id` (the deterministic one, say)
would still pass. Asserting the `(job_name, run_id)` pairs against each TI's
`id` would catch that.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
