Copilot commented on code in PR #74726:
URL: https://github.com/apache/airflow/pull/74726#discussion_r4501530425
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py:
##########
@@ -32,19 +39,314 @@
from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
from airflow.api_fastapi.execution_api.datamodels.asset import AssetResponse
from airflow.api_fastapi.execution_api.datamodels.asset_event import (
+ AssetEventReceipt,
AssetEventResponse,
AssetEventsResponse,
+ EmitAssetEventBody,
)
-from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel
+from airflow.api_fastapi.execution_api.datamodels.token import TIToken
+from airflow.api_fastapi.execution_api.security import (
+ _IN_PROCESS_NON_TI_CALLER,
+ CurrentTIToken,
+ ExecutionAPIRoute,
+ skip_auto_ti_attempt_live,
+)
+from airflow.assets.manager import AssetManager
+from airflow.configuration import conf
+from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel,
TaskOutletAssetReference
+from airflow.models.callback import Callback, ExecutorCallback
+from airflow.models.dagrun import DagRun
+from airflow.models.deadline import Deadline
+from airflow.models.taskinstance import TaskInstance
from airflow.utils.sqlalchemy import JsonContains
+from airflow.utils.state import CallbackState
router = APIRouter(
+ route_class=ExecutionAPIRoute,
responses={
status.HTTP_404_NOT_FOUND: {"description": "Asset not found"},
status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
},
)
+EMISSION_ORIGIN = "emission"
+# Prototype defaults; the AIP leaves the final values open.
+ARCHIVED_RETRY_EMISSION_WINDOW = timedelta(minutes=10)
+MAX_EMISSION_EXTRA_BYTES = 64 * 1024
+MAX_EMISSIONS_PER_SOURCE = 10_000
+
+
+class _EmissionSource(NamedTuple):
+ task_instance: TaskInstance | None
+ dag_run: DagRun
+ # Retries keep these values, so a repeated key stays in one scope; a clear
changes clear_number.
+ scope: tuple[str | int | None, ...]
+
+
+def _get_dag_run(dag_id: str | None, run_id: str | None, *, session: Session)
-> DagRun:
+ dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id,
DagRun.run_id == run_id))
+ if dag_run is None:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND, detail={"reason": "not_found",
"message": "Source Dag run not found"}
+ )
+ return dag_run
+
+
+def _build_dag_run_source(dag_run: DagRun) -> _EmissionSource:
+ return _EmissionSource(
+ task_instance=None,
+ dag_run=dag_run,
+ scope=("dag_run", dag_run.dag_id, dag_run.run_id,
dag_run.clear_number),
+ )
+
+
+def _build_task_source(ti: TaskInstance, *, session: Session) ->
_EmissionSource:
+ dag_run = _get_dag_run(ti.dag_id, ti.run_id, session=session)
+ return _EmissionSource(
+ task_instance=ti,
+ dag_run=dag_run,
+ scope=("task", ti.dag_id, ti.run_id, ti.task_id, ti.map_index,
dag_run.clear_number),
+ )
+
+
+def _require_emitting_attempt(ti: TaskInstance) -> None:
+ if ti.working_set:
+ return
+ # The task runner reports up_for_retry before it runs on_retry_callback,
and that report
+ # archives the attempt, so the callback can only emit through this bounded
exception.
+ if (
+ ti.archived_reason == "retry"
+ and ti.end_date is not None
+ and timezone.utcnow() - ti.end_date <= ARCHIVED_RETRY_EMISSION_WINDOW
+ ):
+ return
+ raise HTTPException(
+ status.HTTP_410_GONE,
+ detail={"reason": "attempt_archived", "message": "This task attempt
can no longer emit events"},
+ )
+
+
+def _resolve_emission_source(
+ body: EmitAssetEventBody, request: Request, token: TIToken, *, session:
Session
+) -> _EmissionSource:
+ if body.source_dag_id is not None or body.source_run_id is not None:
+ # Only the trusted, in-process Dag processor transport names a
callback's Dag run.
+ if not request.scope.get(_IN_PROCESS_NON_TI_CALLER):
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ detail={"reason": "forbidden", "message": "Callback provenance
requires the Dag processor"},
+ )
+ return _build_dag_run_source(_get_dag_run(body.source_dag_id,
body.source_run_id, session=session))
+
+ ti = session.scalar(
+ select(TaskInstance).where(TaskInstance.id ==
token.id).execution_options(include_all_attempts=True)
+ )
+ if ti is not None:
+ _require_emitting_attempt(ti)
+ return _build_task_source(ti, session=session)
+
+ callback = session.get(Callback, token.id)
+ if isinstance(callback, ExecutorCallback):
+ if callback.state != CallbackState.RUNNING:
+ raise HTTPException(
+ status.HTTP_410_GONE,
+ detail={"reason": "callback_finished", "message": "This
callback can no longer emit events"},
+ )
+ dagrun_id =
session.scalar(select(Deadline.dagrun_id).where(Deadline.callback_id ==
callback.id))
+ dag_run = session.get(DagRun, dagrun_id) if dagrun_id is not None else
None
+ if dag_run is None:
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ detail={"reason": "forbidden", "message": "This callback has
no Dag run to emit for"},
+ )
+ return _build_dag_run_source(dag_run)
+
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ detail={"reason": "forbidden", "message": "Asset emission requires a
task or a callback"},
+ )
+
+
+def _get_emission_asset(body: EmitAssetEventBody, *, session: Session) ->
AssetModel:
+ query = select(AssetModel).options(joinedload(AssetModel.active))
+ if body.name is not None:
+ query = query.where(AssetModel.name == body.name)
+ if body.uri is not None:
+ query = query.where(AssetModel.uri == body.uri)
+ assets = session.scalars(query).unique().all()
+ active = [asset for asset in assets if asset.active is not None]
+ if len(active) == 1:
+ return active[0]
+ if len(active) > 1:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ detail={"reason": "asset_ambiguous", "message": "The reference
matches more than one asset"},
+ )
+ if assets:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ detail={"reason": "asset_inactive", "message": "No active Dag
references this asset"},
+ )
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ detail={"reason": "asset_not_found", "message": "The asset is not in
the Airflow database"},
+ )
+
+
+def _compute_hash(value: object) -> str:
+ """Hash a value as canonical JSON; a NaN or infinite float raises
ValueError."""
+ payload = json.dumps(value, sort_keys=True, separators=(",", ":"),
allow_nan=False)
+ return hashlib.sha256(payload.encode()).hexdigest()
+
+
+def _compute_payload_hash(asset: AssetModel, body: EmitAssetEventBody) -> str:
+ try:
+ extra_size = len(json.dumps(body.extra, separators=(",", ":"),
allow_nan=False).encode())
+ payload_hash = _compute_hash(
+ {"asset_id": asset.id, "extra": body.extra, "partition_key":
body.partition_key}
+ )
+ except ValueError:
+ raise HTTPException(
+ status.HTTP_422_UNPROCESSABLE_CONTENT,
+ detail={
+ "reason": "invalid_json",
+ "message": "extra must be strict JSON, without NaN or
Infinity",
+ },
+ )
+ if extra_size > MAX_EMISSION_EXTRA_BYTES:
+ raise HTTPException(
+ status.HTTP_413_CONTENT_TOO_LARGE,
+ detail={
+ "reason": "extra_too_large",
+ "message": f"extra is larger than {MAX_EMISSION_EXTRA_BYTES}
bytes",
+ },
+ )
+ return payload_hash
+
+
+def _require_emission_list_entry(source: _EmissionSource, asset: AssetModel,
*, session: Session) -> None:
+ ti = source.task_instance
+ if ti is None or not conf.getboolean("core", "multi_team"):
+ return
+ declared = session.scalar(
+ select(TaskOutletAssetReference.asset_id).where(
+ TaskOutletAssetReference.dag_id == ti.dag_id,
+ TaskOutletAssetReference.task_id == ti.task_id,
+ TaskOutletAssetReference.asset_id == asset.id,
+ )
+ )
+ if declared is None:
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ detail={
+ "reason": "asset_not_in_emission_list",
+ "message": "The task does not declare this asset",
+ },
+ )
+
+
+def _require_emission_budget(source: _EmissionSource, *, session: Session) ->
None:
+ query =
select(func.count()).select_from(AssetEvent).where(AssetEvent.origin ==
EMISSION_ORIGIN)
+ if source.task_instance is not None:
+ query = query.where(AssetEvent.source_task_instance_id ==
source.task_instance.id)
+ else:
+ query = query.where(
+ AssetEvent.source_dag_id == source.dag_run.dag_id,
+ AssetEvent.source_run_id == source.dag_run.run_id,
+ AssetEvent.source_task_id.is_(None),
+ )
+ if (session.scalar(query) or 0) >= MAX_EMISSIONS_PER_SOURCE:
Review Comment:
The budget check and event insertion are not serialized. With one slot
remaining, two requests for the same source and different keys can both observe
a count below the maximum and both commit; the unique scope-hash index only
prevents requests using the same key. This also occurs with concurrent
callbacks sharing a Dag-run source. Reserve the budget atomically, or lock the
source before counting and retain that lock through insertion, and test
concurrent distinct-key requests at the limit.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py:
##########
@@ -32,19 +39,314 @@
from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
from airflow.api_fastapi.execution_api.datamodels.asset import AssetResponse
from airflow.api_fastapi.execution_api.datamodels.asset_event import (
+ AssetEventReceipt,
AssetEventResponse,
AssetEventsResponse,
+ EmitAssetEventBody,
)
-from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel
+from airflow.api_fastapi.execution_api.datamodels.token import TIToken
+from airflow.api_fastapi.execution_api.security import (
+ _IN_PROCESS_NON_TI_CALLER,
+ CurrentTIToken,
+ ExecutionAPIRoute,
+ skip_auto_ti_attempt_live,
+)
+from airflow.assets.manager import AssetManager
+from airflow.configuration import conf
+from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel,
TaskOutletAssetReference
+from airflow.models.callback import Callback, ExecutorCallback
+from airflow.models.dagrun import DagRun
+from airflow.models.deadline import Deadline
+from airflow.models.taskinstance import TaskInstance
from airflow.utils.sqlalchemy import JsonContains
+from airflow.utils.state import CallbackState
router = APIRouter(
+ route_class=ExecutionAPIRoute,
responses={
status.HTTP_404_NOT_FOUND: {"description": "Asset not found"},
status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
},
)
+EMISSION_ORIGIN = "emission"
+# Prototype defaults; the AIP leaves the final values open.
+ARCHIVED_RETRY_EMISSION_WINDOW = timedelta(minutes=10)
+MAX_EMISSION_EXTRA_BYTES = 64 * 1024
+MAX_EMISSIONS_PER_SOURCE = 10_000
+
+
+class _EmissionSource(NamedTuple):
+ task_instance: TaskInstance | None
+ dag_run: DagRun
+ # Retries keep these values, so a repeated key stays in one scope; a clear
changes clear_number.
+ scope: tuple[str | int | None, ...]
+
+
+def _get_dag_run(dag_id: str | None, run_id: str | None, *, session: Session)
-> DagRun:
+ dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id,
DagRun.run_id == run_id))
+ if dag_run is None:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND, detail={"reason": "not_found",
"message": "Source Dag run not found"}
+ )
+ return dag_run
+
+
+def _build_dag_run_source(dag_run: DagRun) -> _EmissionSource:
+ return _EmissionSource(
+ task_instance=None,
+ dag_run=dag_run,
+ scope=("dag_run", dag_run.dag_id, dag_run.run_id,
dag_run.clear_number),
+ )
+
+
+def _build_task_source(ti: TaskInstance, *, session: Session) ->
_EmissionSource:
+ dag_run = _get_dag_run(ti.dag_id, ti.run_id, session=session)
+ return _EmissionSource(
+ task_instance=ti,
+ dag_run=dag_run,
+ scope=("task", ti.dag_id, ti.run_id, ti.task_id, ti.map_index,
dag_run.clear_number),
Review Comment:
`TaskInstance.map_index` is now a synonym for `region_index`, which is not
unique across regional executions. A mapped task expanded inside two loop
iterations can have the same Dag/run/task/index but different `region_id`s, so
reusing a key incorrectly replays the other instance's event or returns an
idempotency conflict. Include the regional execution identity in the scope;
retries preserve `region_id`, so this can distinguish independent instances
without breaking retry deduplication.
##########
airflow-core/src/airflow/dag_processing/collection.py:
##########
@@ -1104,6 +1104,8 @@ def _get_task_ports(data: dict, inlets: bool, outlets:
bool) -> Iterable[str]:
yield from data.get("inlets") or ()
if outlets:
yield from data.get("outlets") or ()
+ # An emission list declares a producer like an outlet does, but
success emits nothing for it.
+ yield from data.get("emits") or ()
Review Comment:
When the same asset appears in both `outlets` and `emits` with different
access controls, this ordering makes `add_task_asset_references()` overwrite
the outlet's consumer policy with the `emits` policy for the shared row. Adding
an unrestricted emission declaration can thus cause existing success-only
outlet events to trigger teams previously excluded. Preserve the outlet policy
or reject conflicting declarations rather than silently replacing it.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py:
##########
@@ -32,19 +39,314 @@
from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
from airflow.api_fastapi.execution_api.datamodels.asset import AssetResponse
from airflow.api_fastapi.execution_api.datamodels.asset_event import (
+ AssetEventReceipt,
AssetEventResponse,
AssetEventsResponse,
+ EmitAssetEventBody,
)
-from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel
+from airflow.api_fastapi.execution_api.datamodels.token import TIToken
+from airflow.api_fastapi.execution_api.security import (
+ _IN_PROCESS_NON_TI_CALLER,
+ CurrentTIToken,
+ ExecutionAPIRoute,
+ skip_auto_ti_attempt_live,
+)
+from airflow.assets.manager import AssetManager
+from airflow.configuration import conf
+from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel,
TaskOutletAssetReference
+from airflow.models.callback import Callback, ExecutorCallback
+from airflow.models.dagrun import DagRun
+from airflow.models.deadline import Deadline
+from airflow.models.taskinstance import TaskInstance
from airflow.utils.sqlalchemy import JsonContains
+from airflow.utils.state import CallbackState
router = APIRouter(
+ route_class=ExecutionAPIRoute,
responses={
status.HTTP_404_NOT_FOUND: {"description": "Asset not found"},
status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
},
)
+EMISSION_ORIGIN = "emission"
+# Prototype defaults; the AIP leaves the final values open.
+ARCHIVED_RETRY_EMISSION_WINDOW = timedelta(minutes=10)
+MAX_EMISSION_EXTRA_BYTES = 64 * 1024
+MAX_EMISSIONS_PER_SOURCE = 10_000
+
+
+class _EmissionSource(NamedTuple):
+ task_instance: TaskInstance | None
+ dag_run: DagRun
+ # Retries keep these values, so a repeated key stays in one scope; a clear
changes clear_number.
+ scope: tuple[str | int | None, ...]
+
+
+def _get_dag_run(dag_id: str | None, run_id: str | None, *, session: Session)
-> DagRun:
+ dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id,
DagRun.run_id == run_id))
+ if dag_run is None:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND, detail={"reason": "not_found",
"message": "Source Dag run not found"}
+ )
+ return dag_run
+
+
+def _build_dag_run_source(dag_run: DagRun) -> _EmissionSource:
+ return _EmissionSource(
+ task_instance=None,
+ dag_run=dag_run,
+ scope=("dag_run", dag_run.dag_id, dag_run.run_id,
dag_run.clear_number),
+ )
+
+
+def _build_task_source(ti: TaskInstance, *, session: Session) ->
_EmissionSource:
+ dag_run = _get_dag_run(ti.dag_id, ti.run_id, session=session)
+ return _EmissionSource(
+ task_instance=ti,
+ dag_run=dag_run,
+ scope=("task", ti.dag_id, ti.run_id, ti.task_id, ti.map_index,
dag_run.clear_number),
+ )
+
+
+def _require_emitting_attempt(ti: TaskInstance) -> None:
+ if ti.working_set:
+ return
+ # The task runner reports up_for_retry before it runs on_retry_callback,
and that report
+ # archives the attempt, so the callback can only emit through this bounded
exception.
+ if (
+ ti.archived_reason == "retry"
+ and ti.end_date is not None
+ and timezone.utcnow() - ti.end_date <= ARCHIVED_RETRY_EMISSION_WINDOW
+ ):
+ return
+ raise HTTPException(
+ status.HTTP_410_GONE,
+ detail={"reason": "attempt_archived", "message": "This task attempt
can no longer emit events"},
+ )
Review Comment:
The reason this comment gives for the exception doesn't match the
supervisor. `ActivitySubprocess._handle_task_state` (registered for both
`TaskState` and `RetryTask`) only stores the message in
`_pending_terminal_state_msg`. `update_task_state_if_needed()` sends it to the
API after the child process exits, and `InProcessTestSupervisor` does the same
after `finalize()`. So `on_retry_callback` runs while the attempt is still in
the working set, and the normal `ti.working_set` check already lets it emit.
The exception is also wider than intended. `archived_reason == "retry"` is
set by every `prepare_db_for_next_try()`, including scheduler-driven retries:
`handle_failure` for heartbeat timeouts and executor failures, and the adoption
reset in `scheduler_job_runner`. With this exception, an attempt the server
already declared dead and replaced can keep writing asset events for 10 minutes
while its successor runs, and those events are attributed to the dead attempt.
Unless a concrete runtime path archives the attempt before its callbacks run,
please remove the exception (the `test_archived_attempt` "retry-callback" case
would then expect 410).
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py:
##########
@@ -32,19 +39,314 @@
from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
from airflow.api_fastapi.execution_api.datamodels.asset import AssetResponse
from airflow.api_fastapi.execution_api.datamodels.asset_event import (
+ AssetEventReceipt,
AssetEventResponse,
AssetEventsResponse,
+ EmitAssetEventBody,
)
-from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel
+from airflow.api_fastapi.execution_api.datamodels.token import TIToken
+from airflow.api_fastapi.execution_api.security import (
+ _IN_PROCESS_NON_TI_CALLER,
+ CurrentTIToken,
+ ExecutionAPIRoute,
+ skip_auto_ti_attempt_live,
+)
+from airflow.assets.manager import AssetManager
+from airflow.configuration import conf
+from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel,
TaskOutletAssetReference
+from airflow.models.callback import Callback, ExecutorCallback
+from airflow.models.dagrun import DagRun
+from airflow.models.deadline import Deadline
+from airflow.models.taskinstance import TaskInstance
from airflow.utils.sqlalchemy import JsonContains
+from airflow.utils.state import CallbackState
router = APIRouter(
+ route_class=ExecutionAPIRoute,
responses={
status.HTTP_404_NOT_FOUND: {"description": "Asset not found"},
status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
},
)
+EMISSION_ORIGIN = "emission"
+# Prototype defaults; the AIP leaves the final values open.
+ARCHIVED_RETRY_EMISSION_WINDOW = timedelta(minutes=10)
+MAX_EMISSION_EXTRA_BYTES = 64 * 1024
+MAX_EMISSIONS_PER_SOURCE = 10_000
+
+
+class _EmissionSource(NamedTuple):
+ task_instance: TaskInstance | None
+ dag_run: DagRun
+ # Retries keep these values, so a repeated key stays in one scope; a clear
changes clear_number.
+ scope: tuple[str | int | None, ...]
+
+
+def _get_dag_run(dag_id: str | None, run_id: str | None, *, session: Session)
-> DagRun:
+ dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id,
DagRun.run_id == run_id))
+ if dag_run is None:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND, detail={"reason": "not_found",
"message": "Source Dag run not found"}
+ )
+ return dag_run
+
+
+def _build_dag_run_source(dag_run: DagRun) -> _EmissionSource:
+ return _EmissionSource(
+ task_instance=None,
+ dag_run=dag_run,
+ scope=("dag_run", dag_run.dag_id, dag_run.run_id,
dag_run.clear_number),
+ )
+
+
+def _build_task_source(ti: TaskInstance, *, session: Session) ->
_EmissionSource:
+ dag_run = _get_dag_run(ti.dag_id, ti.run_id, session=session)
+ return _EmissionSource(
+ task_instance=ti,
+ dag_run=dag_run,
+ scope=("task", ti.dag_id, ti.run_id, ti.task_id, ti.map_index,
dag_run.clear_number),
+ )
+
+
+def _require_emitting_attempt(ti: TaskInstance) -> None:
+ if ti.working_set:
+ return
+ # The task runner reports up_for_retry before it runs on_retry_callback,
and that report
+ # archives the attempt, so the callback can only emit through this bounded
exception.
+ if (
+ ti.archived_reason == "retry"
+ and ti.end_date is not None
+ and timezone.utcnow() - ti.end_date <= ARCHIVED_RETRY_EMISSION_WINDOW
+ ):
+ return
+ raise HTTPException(
+ status.HTTP_410_GONE,
+ detail={"reason": "attempt_archived", "message": "This task attempt
can no longer emit events"},
+ )
+
+
+def _resolve_emission_source(
+ body: EmitAssetEventBody, request: Request, token: TIToken, *, session:
Session
+) -> _EmissionSource:
+ if body.source_dag_id is not None or body.source_run_id is not None:
+ # Only the trusted, in-process Dag processor transport names a
callback's Dag run.
+ if not request.scope.get(_IN_PROCESS_NON_TI_CALLER):
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ detail={"reason": "forbidden", "message": "Callback provenance
requires the Dag processor"},
+ )
+ return _build_dag_run_source(_get_dag_run(body.source_dag_id,
body.source_run_id, session=session))
+
+ ti = session.scalar(
+ select(TaskInstance).where(TaskInstance.id ==
token.id).execution_options(include_all_attempts=True)
+ )
+ if ti is not None:
+ _require_emitting_attempt(ti)
+ return _build_task_source(ti, session=session)
+
+ callback = session.get(Callback, token.id)
+ if isinstance(callback, ExecutorCallback):
+ if callback.state != CallbackState.RUNNING:
+ raise HTTPException(
+ status.HTTP_410_GONE,
+ detail={"reason": "callback_finished", "message": "This
callback can no longer emit events"},
+ )
+ dagrun_id =
session.scalar(select(Deadline.dagrun_id).where(Deadline.callback_id ==
callback.id))
+ dag_run = session.get(DagRun, dagrun_id) if dagrun_id is not None else
None
+ if dag_run is None:
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ detail={"reason": "forbidden", "message": "This callback has
no Dag run to emit for"},
+ )
+ return _build_dag_run_source(dag_run)
+
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ detail={"reason": "forbidden", "message": "Asset emission requires a
task or a callback"},
+ )
+
+
+def _get_emission_asset(body: EmitAssetEventBody, *, session: Session) ->
AssetModel:
+ query = select(AssetModel).options(joinedload(AssetModel.active))
+ if body.name is not None:
+ query = query.where(AssetModel.name == body.name)
+ if body.uri is not None:
+ query = query.where(AssetModel.uri == body.uri)
+ assets = session.scalars(query).unique().all()
+ active = [asset for asset in assets if asset.active is not None]
+ if len(active) == 1:
+ return active[0]
+ if len(active) > 1:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ detail={"reason": "asset_ambiguous", "message": "The reference
matches more than one asset"},
+ )
+ if assets:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ detail={"reason": "asset_inactive", "message": "No active Dag
references this asset"},
+ )
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ detail={"reason": "asset_not_found", "message": "The asset is not in
the Airflow database"},
+ )
+
+
+def _compute_hash(value: object) -> str:
+ """Hash a value as canonical JSON; a NaN or infinite float raises
ValueError."""
+ payload = json.dumps(value, sort_keys=True, separators=(",", ":"),
allow_nan=False)
+ return hashlib.sha256(payload.encode()).hexdigest()
+
+
+def _compute_payload_hash(asset: AssetModel, body: EmitAssetEventBody) -> str:
+ try:
+ extra_size = len(json.dumps(body.extra, separators=(",", ":"),
allow_nan=False).encode())
+ payload_hash = _compute_hash(
+ {"asset_id": asset.id, "extra": body.extra, "partition_key":
body.partition_key}
+ )
+ except ValueError:
+ raise HTTPException(
+ status.HTTP_422_UNPROCESSABLE_CONTENT,
+ detail={
+ "reason": "invalid_json",
+ "message": "extra must be strict JSON, without NaN or
Infinity",
+ },
+ )
+ if extra_size > MAX_EMISSION_EXTRA_BYTES:
+ raise HTTPException(
+ status.HTTP_413_CONTENT_TOO_LARGE,
+ detail={
+ "reason": "extra_too_large",
+ "message": f"extra is larger than {MAX_EMISSION_EXTRA_BYTES}
bytes",
+ },
+ )
+ return payload_hash
+
+
+def _require_emission_list_entry(source: _EmissionSource, asset: AssetModel,
*, session: Session) -> None:
+ ti = source.task_instance
+ if ti is None or not conf.getboolean("core", "multi_team"):
+ return
+ declared = session.scalar(
+ select(TaskOutletAssetReference.asset_id).where(
+ TaskOutletAssetReference.dag_id == ti.dag_id,
+ TaskOutletAssetReference.task_id == ti.task_id,
+ TaskOutletAssetReference.asset_id == asset.id,
+ )
+ )
+ if declared is None:
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ detail={
+ "reason": "asset_not_in_emission_list",
+ "message": "The task does not declare this asset",
+ },
+ )
+
+
+def _require_emission_budget(source: _EmissionSource, *, session: Session) ->
None:
+ query =
select(func.count()).select_from(AssetEvent).where(AssetEvent.origin ==
EMISSION_ORIGIN)
+ if source.task_instance is not None:
+ query = query.where(AssetEvent.source_task_instance_id ==
source.task_instance.id)
+ else:
+ query = query.where(
+ AssetEvent.source_dag_id == source.dag_run.dag_id,
+ AssetEvent.source_run_id == source.dag_run.run_id,
+ AssetEvent.source_task_id.is_(None),
+ )
+ if (session.scalar(query) or 0) >= MAX_EMISSIONS_PER_SOURCE:
+ raise HTTPException(
+ status.HTTP_429_TOO_MANY_REQUESTS,
+ detail={"reason": "emission_limit", "message": f"More than
{MAX_EMISSIONS_PER_SOURCE} emissions"},
+ )
+
+
+def _replay(event: AssetEvent, payload_hash: str, response: Response) ->
AssetEventReceipt:
+ if event.payload_hash != payload_hash:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ detail={
+ "reason": "idempotency_conflict",
+ "message": "The key was already used with a different payload",
+ "event_id": event.id,
+ },
+ )
+ response.status_code = status.HTTP_200_OK
+ return AssetEventReceipt(id=event.id, timestamp=event.timestamp,
replayed=True)
+
+
[email protected](
+ "/emit",
+ status_code=status.HTTP_201_CREATED,
+ responses=create_openapi_http_exception_doc(
+ [
+ status.HTTP_403_FORBIDDEN,
+ status.HTTP_404_NOT_FOUND,
+ status.HTTP_409_CONFLICT,
+ status.HTTP_410_GONE,
+ status.HTTP_413_CONTENT_TOO_LARGE,
+ status.HTTP_422_UNPROCESSABLE_CONTENT,
+ status.HTTP_429_TOO_MANY_REQUESTS,
+ ]
+ ),
+)
+@skip_auto_ti_attempt_live
+def emit_asset_event(
+ *,
+ body: EmitAssetEventBody,
+ request: Request,
+ response: Response,
+ session: SessionDep,
+ token: TIToken = CurrentTIToken,
+) -> AssetEventReceipt:
+ """Write an asset event now, independent of the outcome of the task or
callback that emits it."""
+ source = _resolve_emission_source(body, request, token, session=session)
+ asset = _get_emission_asset(body, session=session)
+ payload_hash = _compute_payload_hash(asset, body)
+ _require_emission_list_entry(source, asset, session=session)
+ scope_hash = _compute_hash([EMISSION_ORIGIN, asset.id, *source.scope,
body.idempotency_key])
+ if (
+ existing :=
session.scalar(select(AssetEvent).where(AssetEvent.scope_hash == scope_hash))
+ ) is not None:
+ return _replay(existing, payload_hash, response)
+ _require_emission_budget(source, session=session)
+
+ callbacks: list[Callable[[], None]] = []
+ try:
+ event = AssetManager.register_asset_change(
+ asset=asset,
+ task_instance=source.task_instance,
+ dag_run=None if source.task_instance is not None else
source.dag_run,
+ extra=body.extra,
+ partition_key=body.partition_key,
+ origin=EMISSION_ORIGIN,
+ idempotency_key=body.idempotency_key,
+ scope_hash=scope_hash,
+ payload_hash=payload_hash,
+ callback_sink=callbacks,
+ session=session,
+ )
+ except IntegrityError:
+ # Another request with the same key committed first. This request
wrote nothing before
+ # the event, so rolling back the whole transaction loses nothing.
+ session.rollback()
+ existing =
session.scalar(select(AssetEvent).where(AssetEvent.scope_hash == scope_hash))
+ if existing is None:
+ raise
+ return _replay(existing, payload_hash, response)
+ if event is None:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ detail={"reason": "asset_not_found", "message": "The asset is not
in the Airflow database"},
+ )
+ receipt = AssetEventReceipt(id=event.id, timestamp=event.timestamp)
+ # Listeners must not see an event that the database does not commit.
+ session.commit()
Review Comment:
This commits a caller-owned `SessionDep` inside the route. The session
dependency already owns the transaction through `_get_session()` and
`create_session()` (`common/db/common.py:47–52`), so this also commits any
pending caller/dependency work and prevents that work from being rolled back on
a subsequent failure. Keep transaction completion with the session owner and
arrange listener execution after its successful commit, or explicitly own a
separate transaction for emission instead of committing the supplied session.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py:
##########
@@ -32,19 +39,314 @@
from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
from airflow.api_fastapi.execution_api.datamodels.asset import AssetResponse
from airflow.api_fastapi.execution_api.datamodels.asset_event import (
+ AssetEventReceipt,
AssetEventResponse,
AssetEventsResponse,
+ EmitAssetEventBody,
)
-from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel
+from airflow.api_fastapi.execution_api.datamodels.token import TIToken
+from airflow.api_fastapi.execution_api.security import (
+ _IN_PROCESS_NON_TI_CALLER,
+ CurrentTIToken,
+ ExecutionAPIRoute,
+ skip_auto_ti_attempt_live,
+)
+from airflow.assets.manager import AssetManager
+from airflow.configuration import conf
+from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel,
TaskOutletAssetReference
+from airflow.models.callback import Callback, ExecutorCallback
+from airflow.models.dagrun import DagRun
+from airflow.models.deadline import Deadline
+from airflow.models.taskinstance import TaskInstance
from airflow.utils.sqlalchemy import JsonContains
+from airflow.utils.state import CallbackState
router = APIRouter(
+ route_class=ExecutionAPIRoute,
responses={
status.HTTP_404_NOT_FOUND: {"description": "Asset not found"},
status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"},
},
)
+EMISSION_ORIGIN = "emission"
+# Prototype defaults; the AIP leaves the final values open.
+ARCHIVED_RETRY_EMISSION_WINDOW = timedelta(minutes=10)
+MAX_EMISSION_EXTRA_BYTES = 64 * 1024
+MAX_EMISSIONS_PER_SOURCE = 10_000
+
+
+class _EmissionSource(NamedTuple):
+ task_instance: TaskInstance | None
+ dag_run: DagRun
+ # Retries keep these values, so a repeated key stays in one scope; a clear
changes clear_number.
+ scope: tuple[str | int | None, ...]
+
+
+def _get_dag_run(dag_id: str | None, run_id: str | None, *, session: Session)
-> DagRun:
+ dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id,
DagRun.run_id == run_id))
+ if dag_run is None:
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND, detail={"reason": "not_found",
"message": "Source Dag run not found"}
+ )
+ return dag_run
+
+
+def _build_dag_run_source(dag_run: DagRun) -> _EmissionSource:
+ return _EmissionSource(
+ task_instance=None,
+ dag_run=dag_run,
+ scope=("dag_run", dag_run.dag_id, dag_run.run_id,
dag_run.clear_number),
+ )
+
+
+def _build_task_source(ti: TaskInstance, *, session: Session) ->
_EmissionSource:
+ dag_run = _get_dag_run(ti.dag_id, ti.run_id, session=session)
+ return _EmissionSource(
+ task_instance=ti,
+ dag_run=dag_run,
+ scope=("task", ti.dag_id, ti.run_id, ti.task_id, ti.map_index,
dag_run.clear_number),
+ )
+
+
+def _require_emitting_attempt(ti: TaskInstance) -> None:
+ if ti.working_set:
+ return
+ # The task runner reports up_for_retry before it runs on_retry_callback,
and that report
+ # archives the attempt, so the callback can only emit through this bounded
exception.
+ if (
+ ti.archived_reason == "retry"
+ and ti.end_date is not None
+ and timezone.utcnow() - ti.end_date <= ARCHIVED_RETRY_EMISSION_WINDOW
+ ):
+ return
+ raise HTTPException(
+ status.HTTP_410_GONE,
+ detail={"reason": "attempt_archived", "message": "This task attempt
can no longer emit events"},
+ )
+
+
+def _resolve_emission_source(
+ body: EmitAssetEventBody, request: Request, token: TIToken, *, session:
Session
+) -> _EmissionSource:
+ if body.source_dag_id is not None or body.source_run_id is not None:
+ # Only the trusted, in-process Dag processor transport names a
callback's Dag run.
+ if not request.scope.get(_IN_PROCESS_NON_TI_CALLER):
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ detail={"reason": "forbidden", "message": "Callback provenance
requires the Dag processor"},
+ )
+ return _build_dag_run_source(_get_dag_run(body.source_dag_id,
body.source_run_id, session=session))
+
+ ti = session.scalar(
+ select(TaskInstance).where(TaskInstance.id ==
token.id).execution_options(include_all_attempts=True)
+ )
+ if ti is not None:
+ _require_emitting_attempt(ti)
+ return _build_task_source(ti, session=session)
+
+ callback = session.get(Callback, token.id)
+ if isinstance(callback, ExecutorCallback):
+ if callback.state != CallbackState.RUNNING:
+ raise HTTPException(
+ status.HTTP_410_GONE,
+ detail={"reason": "callback_finished", "message": "This
callback can no longer emit events"},
+ )
+ dagrun_id =
session.scalar(select(Deadline.dagrun_id).where(Deadline.callback_id ==
callback.id))
+ dag_run = session.get(DagRun, dagrun_id) if dagrun_id is not None else
None
+ if dag_run is None:
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ detail={"reason": "forbidden", "message": "This callback has
no Dag run to emit for"},
+ )
+ return _build_dag_run_source(dag_run)
+
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ detail={"reason": "forbidden", "message": "Asset emission requires a
task or a callback"},
+ )
+
+
+def _get_emission_asset(body: EmitAssetEventBody, *, session: Session) ->
AssetModel:
+ query = select(AssetModel).options(joinedload(AssetModel.active))
+ if body.name is not None:
+ query = query.where(AssetModel.name == body.name)
+ if body.uri is not None:
+ query = query.where(AssetModel.uri == body.uri)
+ assets = session.scalars(query).unique().all()
+ active = [asset for asset in assets if asset.active is not None]
+ if len(active) == 1:
+ return active[0]
+ if len(active) > 1:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ detail={"reason": "asset_ambiguous", "message": "The reference
matches more than one asset"},
+ )
+ if assets:
+ raise HTTPException(
+ status.HTTP_409_CONFLICT,
+ detail={"reason": "asset_inactive", "message": "No active Dag
references this asset"},
+ )
+ raise HTTPException(
+ status.HTTP_404_NOT_FOUND,
+ detail={"reason": "asset_not_found", "message": "The asset is not in
the Airflow database"},
+ )
+
+
+def _compute_hash(value: object) -> str:
+ """Hash a value as canonical JSON; a NaN or infinite float raises
ValueError."""
+ payload = json.dumps(value, sort_keys=True, separators=(",", ":"),
allow_nan=False)
+ return hashlib.sha256(payload.encode()).hexdigest()
+
+
+def _compute_payload_hash(asset: AssetModel, body: EmitAssetEventBody) -> str:
+ try:
+ extra_size = len(json.dumps(body.extra, separators=(",", ":"),
allow_nan=False).encode())
+ payload_hash = _compute_hash(
+ {"asset_id": asset.id, "extra": body.extra, "partition_key":
body.partition_key}
+ )
+ except ValueError:
+ raise HTTPException(
+ status.HTTP_422_UNPROCESSABLE_CONTENT,
+ detail={
+ "reason": "invalid_json",
+ "message": "extra must be strict JSON, without NaN or
Infinity",
+ },
+ )
+ if extra_size > MAX_EMISSION_EXTRA_BYTES:
+ raise HTTPException(
+ status.HTTP_413_CONTENT_TOO_LARGE,
+ detail={
+ "reason": "extra_too_large",
+ "message": f"extra is larger than {MAX_EMISSION_EXTRA_BYTES}
bytes",
+ },
+ )
+ return payload_hash
+
+
+def _require_emission_list_entry(source: _EmissionSource, asset: AssetModel,
*, session: Session) -> None:
+ ti = source.task_instance
+ if ti is None or not conf.getboolean("core", "multi_team"):
+ return
+ declared = session.scalar(
+ select(TaskOutletAssetReference.asset_id).where(
+ TaskOutletAssetReference.dag_id == ti.dag_id,
+ TaskOutletAssetReference.task_id == ti.task_id,
+ TaskOutletAssetReference.asset_id == asset.id,
+ )
+ )
+ if declared is None:
+ raise HTTPException(
+ status.HTTP_403_FORBIDDEN,
+ detail={
+ "reason": "asset_not_in_emission_list",
+ "message": "The task does not declare this asset",
+ },
+ )
+
+
+def _require_emission_budget(source: _EmissionSource, *, session: Session) ->
None:
+ query =
select(func.count()).select_from(AssetEvent).where(AssetEvent.origin ==
EMISSION_ORIGIN)
+ if source.task_instance is not None:
+ query = query.where(AssetEvent.source_task_instance_id ==
source.task_instance.id)
+ else:
+ query = query.where(
+ AssetEvent.source_dag_id == source.dag_run.dag_id,
+ AssetEvent.source_run_id == source.dag_run.run_id,
+ AssetEvent.source_task_id.is_(None),
Review Comment:
The Dag-source budget query has no supporting index on `asset_event`: the
existing indexes start with asset ID, task-instance ID, or scope hash, none of
which this branch filters on. Every Dag or executor callback emission therefore
counts through the historical event table, making emission latency grow with
all retained events rather than this source's events. Add a matching composite
source index in the model and migration, or use an indexed per-source counter.
##########
airflow-core/src/airflow/dag_processing/collection.py:
##########
@@ -1104,6 +1104,8 @@ def _get_task_ports(data: dict, inlets: bool, outlets:
bool) -> Iterable[str]:
yield from data.get("inlets") or ()
if outlets:
yield from data.get("outlets") or ()
+ # An emission list declares a producer like an outlet does, but
success emits nothing for it.
+ yield from data.get("emits") or ()
Review Comment:
This adds producer reference rows, but not the serialized scheduling
dependency. `_DependencyDetector.detect_task_dependencies()` in
`serialized_objects.py:948` still visits only `task.outlets`, while
`get_scheduling_dependencies()` consumes
`SerializedDagModel.get_dag_dependencies()`. An emission-only producer
therefore remains absent from the scheduling graph despite its new reference
row. Include `emits` in dependency detection for ordinary and mapped tasks, and
assert the producer edge in a serialization regression test.
##########
task-sdk/src/airflow/sdk/execution_time/asset_events.py:
##########
@@ -0,0 +1,150 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import json
+import uuid
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING
+
+import attrs
+
+from airflow.sdk.definitions.asset import Asset, AssetEventReceipt,
AssetNameRef, AssetUriRef
+from airflow.sdk.exceptions import (
+ AirflowRuntimeError,
+ AssetEmissionError,
+ AssetEmissionUnavailableError,
+ ErrorType,
+)
+from airflow.sdk.execution_time.comms import AssetEventEmissionResult,
EmitAssetEvent
+
+if TYPE_CHECKING:
+ from collections.abc import Iterator
+
+ from pydantic.types import JsonValue
+
+ from airflow.sdk.definitions.asset import AssetRef
+
+
[email protected]
+class _EmissionScope:
+ callback_index: int | None = None
+ """The callback's position in the trusted parent's batch, for a Dag
processor callback."""
+
+
+_scope: ContextVar[_EmissionScope | None] = ContextVar("asset_emission_scope",
default=None)
+
+
+@contextmanager
+def set_asset_emission_scope(*, callback_index: int | None = None) ->
Iterator[None]:
+ """
+ Let the code inside emit asset events.
+
+ The runtime that starts user code sets this: the task runner around a task
and its callbacks,
+ the callback supervisor around an executor callback, and the Dag processor
around a callback.
+ A Dag file parse has supervisor comms too, but no scope, so it cannot emit.
+ """
+ token = _scope.set(_EmissionScope(callback_index=callback_index))
+ try:
+ yield
+ finally:
+ _scope.reset(token)
+
+
+def _raise_emission_error(error: AirflowRuntimeError) -> None:
+ detail = error.error.detail if isinstance(error.error.detail, dict) else {}
+ body = detail.get("detail")
+ server_detail = body.get("detail") if isinstance(body, dict) else None
+ status_code = detail.get("status_code")
+ if error.error.error == ErrorType.PERMISSION_DENIED:
+ raise AssetEmissionUnavailableError(
+ "This context cannot emit asset events", reason="forbidden"
+ ) from error
+ if not isinstance(server_detail, dict):
+ # A server without the endpoint answers with FastAPI's plain "Not
Found".
+ if status_code == 404:
+ raise AssetEmissionUnavailableError(
+ "The API server cannot emit asset events",
reason="unsupported_server", status_code=404
+ ) from error
+ raise error
Review Comment:
FastAPI request-validation errors have a list-valued `detail`, so invalid
partition keys or an overlong idempotency key reach this branch and escape as
`AirflowRuntimeError`. Callers catching the public `AssetEmissionError` cannot
handle these ordinary emission failures, unlike the endpoint's structured
errors. Translate validation responses to `AssetEmissionError` while preserving
status 422, and add a test with a list-valued server detail.
##########
task-sdk/src/airflow/sdk/execution_time/asset_events.py:
##########
@@ -0,0 +1,150 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import json
+import uuid
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING
+
+import attrs
+
+from airflow.sdk.definitions.asset import Asset, AssetEventReceipt,
AssetNameRef, AssetUriRef
+from airflow.sdk.exceptions import (
+ AirflowRuntimeError,
+ AssetEmissionError,
+ AssetEmissionUnavailableError,
+ ErrorType,
+)
+from airflow.sdk.execution_time.comms import AssetEventEmissionResult,
EmitAssetEvent
+
+if TYPE_CHECKING:
+ from collections.abc import Iterator
+
+ from pydantic.types import JsonValue
+
+ from airflow.sdk.definitions.asset import AssetRef
+
+
[email protected]
+class _EmissionScope:
+ callback_index: int | None = None
+ """The callback's position in the trusted parent's batch, for a Dag
processor callback."""
+
+
+_scope: ContextVar[_EmissionScope | None] = ContextVar("asset_emission_scope",
default=None)
+
+
+@contextmanager
+def set_asset_emission_scope(*, callback_index: int | None = None) ->
Iterator[None]:
+ """
+ Let the code inside emit asset events.
+
+ The runtime that starts user code sets this: the task runner around a task
and its callbacks,
+ the callback supervisor around an executor callback, and the Dag processor
around a callback.
+ A Dag file parse has supervisor comms too, but no scope, so it cannot emit.
+ """
+ token = _scope.set(_EmissionScope(callback_index=callback_index))
+ try:
+ yield
+ finally:
+ _scope.reset(token)
+
+
+def _raise_emission_error(error: AirflowRuntimeError) -> None:
+ detail = error.error.detail if isinstance(error.error.detail, dict) else {}
+ body = detail.get("detail")
+ server_detail = body.get("detail") if isinstance(body, dict) else None
+ status_code = detail.get("status_code")
+ if error.error.error == ErrorType.PERMISSION_DENIED:
+ raise AssetEmissionUnavailableError(
+ "This context cannot emit asset events", reason="forbidden"
+ ) from error
+ if not isinstance(server_detail, dict):
+ # A server without the endpoint answers with FastAPI's plain "Not
Found".
+ if status_code == 404:
+ raise AssetEmissionUnavailableError(
+ "The API server cannot emit asset events",
reason="unsupported_server", status_code=404
+ ) from error
+ raise error
+ raise AssetEmissionError(
+ server_detail.get("message", str(error)),
+ reason=server_detail.get("reason"),
+ status_code=status_code,
+ ) from error
+
+
+def emit_asset_event(
+ asset: Asset | AssetRef,
+ *,
+ extra: dict[str, JsonValue] | None = None,
+ partition_key: str | None = None,
+ idempotency_key: str | None = None,
+) -> AssetEventReceipt:
+ """
+ Write an asset event now, from a running task or from a callback.
+
+ Unlike an outlet, the event does not wait for the task to succeed, and a
later failure does not
+ remove it. The asset must already be in the Airflow database. An omitted
partition key means no
+ partition; the event does not inherit the partition of the current run.
+
+ Give a stable ``idempotency_key`` when a retry must not write a second
event. Without one, the
+ SDK makes a key for this call, so only its own automatic HTTP retries
share it.
+ """
+ if isinstance(asset, Asset):
+ name, uri = asset.name, asset.uri
+ elif isinstance(asset, AssetNameRef):
+ name, uri = asset.name, None
+ elif isinstance(asset, AssetUriRef):
+ name, uri = None, asset.uri
+ else:
+ raise TypeError(f"Expected an Asset or an AssetRef, got
{type(asset).__name__}")
+
+ try:
+ json.dumps(extra or {}, allow_nan=False)
+ except (TypeError, ValueError) as error:
+ raise AssetEmissionError(
+ f"extra must be strict JSON, without NaN or Infinity: {error}",
reason="invalid_json"
+ ) from error
+
+ scope = _scope.get()
+ if scope is None:
+ raise AssetEmissionUnavailableError(
+ "Asset events can be emitted only from a running task or a
callback", reason="no_runtime"
+ )
Review Comment:
The emission scope is a `ContextVar`, and threads started with
`threading.Thread`, `ThreadPoolExecutor.submit` or `loop.run_in_executor` begin
with an empty context. So `emit_asset_event()` called from a worker thread
inside a running task raises
`AssetEmissionUnavailableError(reason="no_runtime")`, even though
`SUPERVISOR_COMMS` is a module global with its own thread lock. Other SDK calls
such as `Variable.get` and XCom work from those threads.
This breaks the batch-ingestion pattern the PR uses as motivation, when the
commit loop runs in a pool. Two options:
- Let the runtime set a process-level flag. Each task or callback process
runs one scope at a time, and the Dag processor sets the scope only around
callbacks, after parsing.
- At least document that emission must happen on the thread that runs the
task or callback, or that `contextvars.copy_context()` has to be propagated by
hand.
##########
task-sdk/src/airflow/sdk/execution_time/asset_events.py:
##########
@@ -0,0 +1,150 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import json
+import uuid
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING
+
+import attrs
+
+from airflow.sdk.definitions.asset import Asset, AssetEventReceipt,
AssetNameRef, AssetUriRef
+from airflow.sdk.exceptions import (
+ AirflowRuntimeError,
+ AssetEmissionError,
+ AssetEmissionUnavailableError,
+ ErrorType,
+)
+from airflow.sdk.execution_time.comms import AssetEventEmissionResult,
EmitAssetEvent
+
+if TYPE_CHECKING:
+ from collections.abc import Iterator
+
+ from pydantic.types import JsonValue
+
+ from airflow.sdk.definitions.asset import AssetRef
+
+
[email protected]
+class _EmissionScope:
+ callback_index: int | None = None
+ """The callback's position in the trusted parent's batch, for a Dag
processor callback."""
+
+
+_scope: ContextVar[_EmissionScope | None] = ContextVar("asset_emission_scope",
default=None)
+
+
+@contextmanager
+def set_asset_emission_scope(*, callback_index: int | None = None) ->
Iterator[None]:
+ """
+ Let the code inside emit asset events.
+
+ The runtime that starts user code sets this: the task runner around a task
and its callbacks,
+ the callback supervisor around an executor callback, and the Dag processor
around a callback.
+ A Dag file parse has supervisor comms too, but no scope, so it cannot emit.
+ """
+ token = _scope.set(_EmissionScope(callback_index=callback_index))
+ try:
+ yield
+ finally:
+ _scope.reset(token)
+
+
+def _raise_emission_error(error: AirflowRuntimeError) -> None:
+ detail = error.error.detail if isinstance(error.error.detail, dict) else {}
+ body = detail.get("detail")
+ server_detail = body.get("detail") if isinstance(body, dict) else None
+ status_code = detail.get("status_code")
+ if error.error.error == ErrorType.PERMISSION_DENIED:
+ raise AssetEmissionUnavailableError(
+ "This context cannot emit asset events", reason="forbidden"
+ ) from error
+ if not isinstance(server_detail, dict):
+ # A server without the endpoint answers with FastAPI's plain "Not
Found".
+ if status_code == 404:
+ raise AssetEmissionUnavailableError(
+ "The API server cannot emit asset events",
reason="unsupported_server", status_code=404
+ ) from error
+ raise error
+ raise AssetEmissionError(
+ server_detail.get("message", str(error)),
+ reason=server_detail.get("reason"),
+ status_code=status_code,
+ ) from error
+
+
+def emit_asset_event(
+ asset: Asset | AssetRef,
+ *,
+ extra: dict[str, JsonValue] | None = None,
+ partition_key: str | None = None,
+ idempotency_key: str | None = None,
+) -> AssetEventReceipt:
+ """
+ Write an asset event now, from a running task or from a callback.
+
+ Unlike an outlet, the event does not wait for the task to succeed, and a
later failure does not
+ remove it. The asset must already be in the Airflow database. An omitted
partition key means no
+ partition; the event does not inherit the partition of the current run.
+
+ Give a stable ``idempotency_key`` when a retry must not write a second
event. Without one, the
+ SDK makes a key for this call, so only its own automatic HTTP retries
share it.
+ """
+ if isinstance(asset, Asset):
+ name, uri = asset.name, asset.uri
+ elif isinstance(asset, AssetNameRef):
+ name, uri = asset.name, None
+ elif isinstance(asset, AssetUriRef):
+ name, uri = None, asset.uri
+ else:
+ raise TypeError(f"Expected an Asset or an AssetRef, got
{type(asset).__name__}")
+
+ try:
+ json.dumps(extra or {}, allow_nan=False)
+ except (TypeError, ValueError) as error:
+ raise AssetEmissionError(
+ f"extra must be strict JSON, without NaN or Infinity: {error}",
reason="invalid_json"
+ ) from error
+
+ scope = _scope.get()
+ if scope is None:
+ raise AssetEmissionUnavailableError(
+ "Asset events can be emitted only from a running task or a
callback", reason="no_runtime"
+ )
+
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+ try:
+ result = SUPERVISOR_COMMS.send(
+ EmitAssetEvent(
+ name=name,
+ uri=uri,
+ extra=extra or {},
+ partition_key=partition_key,
+ idempotency_key=idempotency_key or str(uuid.uuid4()),
Review Comment:
An explicitly supplied empty key is treated as though the caller omitted it,
generating a different UUID on every call. That silently publishes duplicate
events instead of rejecting the invalid key as the endpoint does. Only generate
a key for `None`; forward an empty string unchanged so the server rejects it,
or validate it locally.
##########
task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py:
##########
@@ -306,14 +308,15 @@ def _target():
)
_add_xcom_pull_to_context(callback_kwargs)
- success, error_msg = execute_callback(
- callback_path=callback_path,
- callback_kwargs=callback_kwargs,
- dag_rel_path=dag_rel_path,
- bundle_path=bundle_path,
- log=_log,
- callback_kwargs_encoding=callback_kwargs_encoding,
- )
+ with set_asset_emission_scope():
+ success, error_msg = execute_callback(
Review Comment:
This opens the emission scope before `execute_callback()` loads the
callable. That function executes the Dag module via `spec.loader.exec_module()`
(or imports the callback module) before invoking the callback, so import-time
code can now publish events with the deadline callback's identity. This
contradicts the parsing rejection enforced elsewhere and can produce events
during module loading rather than callback execution. Open the scope only
around the actual callback/notifier invocations, leaving module loading outside
it.
--
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]
