You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "dstandish (via GitHub)" <gi...@apache.org> on 2025/11/14 01:43:58 UTC
[PR] Update asset manager to be partition-aware [airflow]
dstandish opened a new pull request, #58289:
URL: https://github.com/apache/airflow/pull/58289
In this PR we make it so that when an asset is updated in a dag run that has a partition key, the asset event also gets the partition key, _and_ records are created in the appropriate tables as needed for scheduling dag runs that listen for the asset.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2577460469
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -1317,6 +1317,8 @@ def register_asset_changes_in_db(
) -> None:
from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetNameRef, AssetUniqueKey, AssetUriRef
+ # TODO: AIP-76 should we provide an interface to override this?
Review Comment:
expanded on the comment in the todo
basically, task might do something different, and if it does, it should be able to tell the truth
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2536789080
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,13 +284,77 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
- return
-
if get_dialect_name(session) == "postgresql":
return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
+ @classmethod
+ def _queue_partitioned_dags(
+ cls,
+ asset_id: int,
+ all_dags: set[DagModel],
+ event: AssetEvent,
+ partition_key: str,
+ session: Session,
+ ) -> list[DagModel]:
+ # todo: AIP-76 there may be a better way to identify that timetable is partition-driven
+ partition_dags = [x for x in all_dags if x.timetable_summary == "Partitioned Asset"]
Review Comment:
I seel you access the timetable from serdag below anyway. So maybe there’s a way to just merge this logic somewhere below instead? It would be a bit less performant (need to deserialise unrelated dags) though. Alternatively, maybe we can just add a bool flag on DagModel to indicate if the dag has a PartitionedAssetTimetable?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2542432760
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -229,12 +242,39 @@ def notify_asset_alias_created(asset_assets: AssetAlias):
def notify_asset_changed(asset: Asset):
"""Run applicable notification actions when an asset is changed."""
try:
+ # todo: AIP-76 this will have to change. needs to know *what* happened to the asset (e.g. partition key)
Review Comment:
my editor (pycharm) doesn't care --- do some editors care?!?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2542913017
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
Review Comment:
so...
inverse_map is actually downstream_to_upstream
are you proposing these changes?
```
map --> upstream_to_downstream
inverse_map --> downstream_to_upstream
```
?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2565194192
##########
airflow-core/src/airflow/serialization/serialized_objects.py:
##########
@@ -458,6 +471,45 @@ def decode_timetable(var: dict[str, Any]) -> Timetable:
return timetable_class.deserialize(var[Encoding.VAR])
+def _load_partition_mapper(importable_string) -> PartitionMapper | None:
+ if importable_string.startswith("airflow.timetables."):
+ return import_string(importable_string)
+ else:
+ return None
+
+
+def encode_partition_mapper(var: PartitionMapper) -> dict[str, Any]:
+ """
+ Encode a PartitionMapper instance.
+
+ This delegates most of the serialization work to the type, so the behavior
+ can be completely controlled by a custom subclass.
+
+ :meta private:
+ """
+ partition_mapper_class = type(var)
+ importable_string = qualname(partition_mapper_class)
+ if _load_partition_mapper(importable_string) is None:
+ raise _PartitionMapperNotFound(importable_string)
Review Comment:
i copied the behavior for timetables; i think we might ultimately make this pluggable in the same way.
but meantime there's an odd error message that says something about sessions and such in the timetable case; i'm not sure what that's about and if it might also apply here.... so i just copied it; if you think it's not necessary, happy to remove it. lmk.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2525493833
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,13 +284,77 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
- return
-
if get_dialect_name(session) == "postgresql":
return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
+ @classmethod
+ def _queue_partitioned_dags(
+ cls,
+ asset_id: int,
+ all_dags: set[DagModel],
+ event: AssetEvent,
+ partition_key: str,
+ session: Session,
+ ) -> list[DagModel]:
+ # todo: AIP-76 there may be a better way to identify that timetable is partition-driven
+ partition_dags = [x for x in all_dags if x.timetable_summary == "Partitioned Asset"]
Review Comment:
@uranusjr @Lee-W do you guys see a better way to do this?
"This" means, identifying which dags are listening to for assets in a partition-aware way. Which, for the moment, means, using PartitionedAssetTimetable.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2577518168
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -229,12 +242,39 @@ def notify_asset_alias_created(asset_assets: AssetAlias):
def notify_asset_changed(asset: Asset):
"""Run applicable notification actions when an asset is changed."""
try:
+ # TODO: AIP-76 this will have to change. needs to know *what* happened to the asset (e.g. partition key)
Review Comment:
btw the issue for that is here https://github.com/orgs/apache/projects/542/views/1?pane=issue&itemId=139253269&issue=apache%7Cairflow%7C58290, and i guess here also https://github.com/orgs/apache/projects/542?pane=issue&itemId=140179031&issue=apache%7Cairflow%7C42254
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2526480811
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,13 +284,77 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
- return
-
if get_dialect_name(session) == "postgresql":
return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
+ @classmethod
+ def _queue_partitioned_dags(
+ cls,
+ asset_id: int,
+ all_dags: set[DagModel],
+ event: AssetEvent,
+ partition_key: str,
+ session: Session,
+ ) -> list[DagModel]:
+ # todo: AIP-76 there may be a better way to identify that timetable is partition-driven
+ partition_dags = [x for x in all_dags if x.timetable_summary == "Partitioned Asset"]
Review Comment:
instead of passing all_dags, should we just query the db and filter 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2539295918
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,13 +284,77 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
- return
-
if get_dialect_name(session) == "postgresql":
return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
+ @classmethod
+ def _queue_partitioned_dags(
+ cls,
+ asset_id: int,
+ all_dags: set[DagModel],
+ event: AssetEvent,
+ partition_key: str,
+ session: Session,
+ ) -> list[DagModel]:
+ # todo: AIP-76 there may be a better way to identify that timetable is partition-driven
+ partition_dags = [x for x in all_dags if x.timetable_summary == "Partitioned Asset"]
Review Comment:
Just to keep the PRs more focused, I propose we leave this as a to-do / task (created on here)
https://github.com/orgs/apache/projects/542/views/1?filterQuery=-status%3ADone&pane=issue&itemId=140081000&issue=apache%7Cairflow%7C58445
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2541556065
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
+
+
+class IdentityMapper(PartitionMapper):
+ """Partition mapper that does not change the key."""
+
+ def map(self, key: str) -> str:
+ return key
+
+ def inverse_map(self, key: str) -> Iterable[str]:
+ yield key
+
+
+class PartitionedAssetTimetable(AssetTriggeredTimetable):
+ """Asset-driven timetable that listens for partitioned assets."""
+
+ @property
+ def summary(self) -> str:
+ return "Partitioned Asset"
+
+ def __init__(self, assets: BaseAsset, partition_mapper: PartitionMapper) -> None:
+ super().__init__(assets=assets)
+ self.asset_condition = assets
+ self.partition_mapper = partition_mapper
+
+ def serialize(self) -> dict[str, Any]:
+ from airflow.serialization.serialized_objects import encode_asset_condition
+
+ return {
+ "asset_condition": encode_asset_condition(self.asset_condition),
+ "partition_mapper_cls": self.partition_mapper.__class__.__name__,
+ }
+
+ @classmethod
+ def deserialize(cls, data: dict[str, Any]) -> Timetable:
+ from airflow.serialization.serialized_objects import decode_asset_condition
+
+ ser_asset_condition = data["asset_condition"]
+ mapper_class_name = data.get("partition_mapper_cls", None)
+ if not mapper_class_name:
+ mapper_class = IdentityMapper
+ else:
+ # todo: AIP-76 what is the right way to do this?
+ if mapper_class_name in globals():
+ mapper_class = globals()[mapper_class_name]
+ else:
+ mapper_class = import_string(mapper_class_name)
Review Comment:
> we could add something like encode_timetable / decode_timetable for partition mappers.
I like this more
> do we need to add partition mappers as a distinct pluggable entity alongside timetables?
Next phase, probably? to keep it as simple as possible for the first phase
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -229,12 +242,39 @@ def notify_asset_alias_created(asset_assets: AssetAlias):
def notify_asset_changed(asset: Asset):
"""Run applicable notification actions when an asset is changed."""
try:
+ # todo: AIP-76 this will have to change. needs to know *what* happened to the asset (e.g. partition key)
Review Comment:
```suggestion
# TODO: AIP-76 this will have to change. needs to know *what* happened to the asset (e.g. partition key)
```
nit: AFAIK, many editors highlight upper case `TODO`
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
Review Comment:
IMO, it would be better 🤔
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on PR #58289:
URL: https://github.com/apache/airflow/pull/58289#issuecomment-3562939595
@uranusjr @Lee-W I have added a couple tests and done the requested renames
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2581529254
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -229,12 +242,44 @@ def notify_asset_alias_created(asset_assets: AssetAlias):
def notify_asset_changed(asset: Asset):
"""Run applicable notification actions when an asset is changed."""
try:
+ # TODO: AIP-76 this will have to change. needs to know *what* happened to the asset (e.g. partition key)
Review Comment:
nit: let's add the link or issue number in the comment here so that we can find where to modify easier
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,15 +288,91 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
+ if get_dialect_name(session) == "postgresql":
+ return cls._queue_dagruns_nonpartitioned_postgres(asset_id, non_partitioned_dags, session)
+ return cls._queue_dagruns_nonpartitioned_slow_path(asset_id, non_partitioned_dags, session)
+
+ @classmethod
+ def _queue_partitioned_dags(
+ cls,
+ asset_id: int,
+ partition_dags: Iterable[DagModel],
+ event: AssetEvent,
+ partition_key: str | None,
+ session: Session,
+ ) -> None:
+ if partition_dags and not partition_key:
+ # TODO: AIP-76 how to best ensure users can see this? Probably add Log record.
+ log.warning(
+ "Listening dags are partition-aware but run has no partition key",
+ listening_dags=[x.dag_id for x in partition_dags],
+ asset_id=asset_id,
+ run_id=event.source_run_id,
+ dag_id=event.source_dag_id,
+ task_id=event.source_task_id,
+ )
return
Review Comment:
I feel this logic should exist in line 269. but non-blocking
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,15 +288,91 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
+ if get_dialect_name(session) == "postgresql":
+ return cls._queue_dagruns_nonpartitioned_postgres(asset_id, non_partitioned_dags, session)
+ return cls._queue_dagruns_nonpartitioned_slow_path(asset_id, non_partitioned_dags, session)
+
+ @classmethod
+ def _queue_partitioned_dags(
+ cls,
+ asset_id: int,
+ partition_dags: Iterable[DagModel],
+ event: AssetEvent,
+ partition_key: str | None,
+ session: Session,
+ ) -> None:
+ if partition_dags and not partition_key:
+ # TODO: AIP-76 how to best ensure users can see this? Probably add Log record.
+ log.warning(
+ "Listening dags are partition-aware but run has no partition key",
+ listening_dags=[x.dag_id for x in partition_dags],
+ asset_id=asset_id,
+ run_id=event.source_run_id,
+ dag_id=event.source_dag_id,
+ task_id=event.source_task_id,
+ )
return
- if get_dialect_name(session) == "postgresql":
- return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
- return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
+ for target_dag in partition_dags:
+ if TYPE_CHECKING:
+ assert partition_key is not None
+ from airflow.models.serialized_dag import SerializedDagModel
+
+ serdag = SerializedDagModel.get(dag_id=target_dag.dag_id, session=session)
+ if not serdag:
Review Comment:
```suggestion
if not (serdag := SerializedDagModel.get(dag_id=target_dag.dag_id, session=session)):
```
##########
airflow-core/src/airflow/example_dags/example_outlet_event_extra.py:
##########
@@ -68,6 +68,9 @@ def asset_with_extra_by_context(*, outlet_events=None):
def _asset_with_extra_from_classic_operator_post_execute(context, result):
context["outlet_events"][asset].extra = {"hi": "bye"}
+ # TODO: AIP-76 probably we want to make it so this could be
Review Comment:
I feel like it add some complexity I'm not sure we what to handle 🤔 Is it needed for partition? Or would it be better to make it the next improvement?
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -1315,6 +1315,9 @@ def register_asset_changes_in_db(
) -> None:
from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetNameRef, AssetUniqueKey, AssetUriRef
+ # TODO: AIP-76 should we provide an interface to override this, so that the task can
+ # tell the truth if for some reason it touches a different partition?
Review Comment:
Why would the "for some reason" be?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2549751632
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
+
+
+class IdentityMapper(PartitionMapper):
+ """Partition mapper that does not change the key."""
+
+ def map(self, key: str) -> str:
+ return key
+
+ def inverse_map(self, key: str) -> Iterable[str]:
+ yield key
+
+
+class PartitionedAssetTimetable(AssetTriggeredTimetable):
+ """Asset-driven timetable that listens for partitioned assets."""
+
+ @property
+ def summary(self) -> str:
+ return "Partitioned Asset"
+
+ def __init__(self, assets: BaseAsset, partition_mapper: PartitionMapper) -> None:
+ super().__init__(assets=assets)
+ self.asset_condition = assets
+ self.partition_mapper = partition_mapper
+
+ def serialize(self) -> dict[str, Any]:
+ from airflow.serialization.serialized_objects import encode_asset_condition
+
+ return {
+ "asset_condition": encode_asset_condition(self.asset_condition),
+ "partition_mapper_cls": self.partition_mapper.__class__.__name__,
+ }
+
+ @classmethod
+ def deserialize(cls, data: dict[str, Any]) -> Timetable:
+ from airflow.serialization.serialized_objects import decode_asset_condition
+
+ ser_asset_condition = data["asset_condition"]
+ mapper_class_name = data.get("partition_mapper_cls", None)
+ if not mapper_class_name:
+ mapper_class = IdentityMapper
+ else:
+ # todo: AIP-76 what is the right way to do this?
+ if mapper_class_name in globals():
+ mapper_class = globals()[mapper_class_name]
+ else:
+ mapper_class = import_string(mapper_class_name)
Review Comment:
resolving this now
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2565283856
##########
airflow-core/src/airflow/serialization/serialized_objects.py:
##########
@@ -458,6 +471,45 @@ def decode_timetable(var: dict[str, Any]) -> Timetable:
return timetable_class.deserialize(var[Encoding.VAR])
+def _load_partition_mapper(importable_string) -> PartitionMapper | None:
+ if importable_string.startswith("airflow.timetables."):
+ return import_string(importable_string)
+ else:
+ return None
+
+
+def encode_partition_mapper(var: PartitionMapper) -> dict[str, Any]:
+ """
+ Encode a PartitionMapper instance.
+
+ This delegates most of the serialization work to the type, so the behavior
+ can be completely controlled by a custom subclass.
+
+ :meta private:
+ """
+ partition_mapper_class = type(var)
+ importable_string = qualname(partition_mapper_class)
+ if _load_partition_mapper(importable_string) is None:
+ raise _PartitionMapperNotFound(importable_string)
Review Comment:
well the other thing is, it restricts to this namespace. we won't always get an import error
importable_string.startswith("airflow.timetables.")
basically it's just banning custom mappers for now
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2530001844
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
+
+
+class IdentityMapper(PartitionMapper):
+ """Partition mapper that does not change the key."""
+
+ def map(self, key: str) -> str:
+ return key
+
+ def inverse_map(self, key: str) -> Iterable[str]:
+ yield key
+
+
+class PartitionedAssetTimetable(AssetTriggeredTimetable):
+ """Asset-driven timetable that listens for partitioned assets."""
+
+ @property
+ def summary(self) -> str:
+ return "Partitioned Asset"
+
+ def __init__(self, assets: BaseAsset, partition_mapper: PartitionMapper) -> None:
+ super().__init__(assets=assets)
+ self.asset_condition = assets
+ self.partition_mapper = partition_mapper
+
+ def serialize(self) -> dict[str, Any]:
+ from airflow.serialization.serialized_objects import encode_asset_condition
+
+ return {
+ "asset_condition": encode_asset_condition(self.asset_condition),
+ "partition_mapper_cls": self.partition_mapper.__class__.__name__,
+ }
+
+ @classmethod
+ def deserialize(cls, data: dict[str, Any]) -> Timetable:
+ from airflow.serialization.serialized_objects import decode_asset_condition
+
+ ser_asset_condition = data["asset_condition"]
+ mapper_class_name = data.get("partition_mapper_cls", None)
+ if not mapper_class_name:
+ mapper_class = IdentityMapper
+ else:
+ # todo: AIP-76 what is the right way to do this?
+ if mapper_class_name in globals():
+ mapper_class = globals()[mapper_class_name]
+ else:
+ mapper_class = import_string(mapper_class_name)
Review Comment:
Is this the best way to do this @uranusjr @Lee-W ?
When serializing / deserializing a timetable, we also need to ser / deser the mapper class. What would you do?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2536706045
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -229,12 +242,40 @@ def notify_asset_alias_created(asset_assets: AssetAlias):
def notify_asset_changed(asset: Asset):
"""Run applicable notification actions when an asset is changed."""
try:
+ # todo: AIP-76 this will have to change. needs to know *what* happened to the asset (e.g. partition key)
+ # maybe we should just add the event to the signature
+ # or add a new hook `on_asset_event`
get_listener_manager().hook.on_asset_changed(asset=asset)
except Exception:
log.exception("error calling listener")
@classmethod
- def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Session) -> None:
+ def _queue_dagruns(
+ cls,
+ asset_id: int,
+ dags_to_queue: set[DagModel],
+ partition_key: str | None,
+ event: AssetEvent,
+ session: Session,
+ ) -> None:
+ log.debug("dags to queue", dags_to_queue=dags_to_queue)
+
+ if not dags_to_queue:
+ return None
+
+ partition_dags = cls._queue_partitioned_dags(
+ asset_id=asset_id,
+ all_dags=dags_to_queue,
+ event=event,
+ partition_key=partition_key,
+ session=session,
+ )
+ for d in partition_dags:
+ dags_to_queue.remove(d) # don't double process
Review Comment:
```suggestion
dags_to_queue.difference_update(partition_dags) # don't double process
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2582280287
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -1315,6 +1315,9 @@ def register_asset_changes_in_db(
) -> None:
from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetNameRef, AssetUniqueKey, AssetUriRef
+ # TODO: AIP-76 should we provide an interface to override this, so that the task can
+ # tell the truth if for some reason it touches a different partition?
Review Comment:
Could be that the task has conditional logic
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2530000995
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,13 +284,77 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
- return
-
if get_dialect_name(session) == "postgresql":
return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
+ @classmethod
+ def _queue_partitioned_dags(
+ cls,
+ asset_id: int,
+ all_dags: set[DagModel],
+ event: AssetEvent,
+ partition_key: str,
+ session: Session,
+ ) -> list[DagModel]:
+ # todo: AIP-76 there may be a better way to identify that timetable is partition-driven
+ partition_dags = [x for x in all_dags if x.timetable_summary == "Partitioned Asset"]
Review Comment:
I realized that my question wasn't clear
What I am asking about is, how to identify which timetables are "partitioned". Since at this point we only have DagModel, we don't have the full timetable. To get that we need to load the serdag. So what I'm doing is using timetable summary to filter out the non-partitioned dags. Having done that, now with the partitioned dags in hand, for each of them, we load the serdag and do the partition mapping.
My question is, is using timetable summary robust enough? Should we instead just load the serdag and check inheritance? Or add some other identifier somehow?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2536706045
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -229,12 +242,40 @@ def notify_asset_alias_created(asset_assets: AssetAlias):
def notify_asset_changed(asset: Asset):
"""Run applicable notification actions when an asset is changed."""
try:
+ # todo: AIP-76 this will have to change. needs to know *what* happened to the asset (e.g. partition key)
+ # maybe we should just add the event to the signature
+ # or add a new hook `on_asset_event`
get_listener_manager().hook.on_asset_changed(asset=asset)
except Exception:
log.exception("error calling listener")
@classmethod
- def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Session) -> None:
+ def _queue_dagruns(
+ cls,
+ asset_id: int,
+ dags_to_queue: set[DagModel],
+ partition_key: str | None,
+ event: AssetEvent,
+ session: Session,
+ ) -> None:
+ log.debug("dags to queue", dags_to_queue=dags_to_queue)
+
+ if not dags_to_queue:
+ return None
+
+ partition_dags = cls._queue_partitioned_dags(
+ asset_id=asset_id,
+ all_dags=dags_to_queue,
+ event=event,
+ partition_key=partition_key,
+ session=session,
+ )
+ for d in partition_dags:
+ dags_to_queue.remove(d) # don't double process
Review Comment:
```suggestion
dags_to_queue = dags_to_queue.difference(partition_dags) # don't double process
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2577445731
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -229,12 +242,39 @@ def notify_asset_alias_created(asset_assets: AssetAlias):
def notify_asset_changed(asset: Asset):
"""Run applicable notification actions when an asset is changed."""
try:
+ # TODO: AIP-76 this will have to change. needs to know *what* happened to the asset (e.g. partition key)
Review Comment:
> on_asset_event
i agree. let's do this in followup.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on PR #58289:
URL: https://github.com/apache/airflow/pull/58289#issuecomment-3603314090
> a few nits and questions, but non of them is blocking this PR
Thank you
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2539269249
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,13 +284,77 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
- return
-
if get_dialect_name(session) == "postgresql":
return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
+ @classmethod
+ def _queue_partitioned_dags(
+ cls,
+ asset_id: int,
+ all_dags: set[DagModel],
+ event: AssetEvent,
+ partition_key: str,
+ session: Session,
+ ) -> list[DagModel]:
+ # todo: AIP-76 there may be a better way to identify that timetable is partition-driven
+ partition_dags = [x for x in all_dags if x.timetable_summary == "Partitioned Asset"]
Review Comment:
though if we get to a point where we can do partial deserialization (e.g. only the timetable) then that would change things
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2577515092
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,15 +283,94 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
- return
-
if get_dialect_name(session) == "postgresql":
- return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
- return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
+ return cls._queue_dagruns_nonpartitioned_postgres(asset_id, dags_to_queue, session)
+ return cls._queue_dagruns_nonpartitioned_slow_path(asset_id, dags_to_queue, session)
@classmethod
- def _slow_path_queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Session) -> None:
+ def _queue_partitioned_dags(
+ cls,
+ asset_id: int,
+ all_dags: set[DagModel],
+ event: AssetEvent,
+ partition_key: str | None,
+ session: Session,
+ ) -> list[DagModel]:
+ # TODO: AIP-76 there may be a better way to identify that timetable is partition-driven
+ partition_dags = [x for x in all_dags if x.timetable_summary == "Partitioned Asset"]
+ if partition_dags and not partition_key:
+ # TODO: AIP-76 how to best ensure users can see this? Probably add Log record.
+ log.warning(
+ "Listening dags are partition-aware but run has no partition key",
+ listening_dags=[x.dag_id for x in partition_dags],
+ asset_id=asset_id,
+ run_id=event.source_run_id,
+ dag_id=event.source_dag_id,
+ task_id=event.source_task_id,
+ )
Review Comment:
sure
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2577485097
##########
airflow-core/src/airflow/example_dags/example_outlet_event_extra.py:
##########
@@ -68,6 +68,9 @@ def asset_with_extra_by_context(*, outlet_events=None):
def _asset_with_extra_from_classic_operator_post_execute(context, result):
context["outlet_events"][asset].extra = {"hi": "bye"}
+ # TODO: AIP-76 probably we want to make it so this could be
Review Comment:
Added more context to this todo.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on PR #58289:
URL: https://github.com/apache/airflow/pull/58289#issuecomment-3597148902
OK @Lee-W @uranusjr i have dealt with all the comments PTAL
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2539936713
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
Review Comment:
🤷
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2555169916
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,13 +283,90 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
- return
-
if get_dialect_name(session) == "postgresql":
return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
Review Comment:
Let’s rename these to clarify they are for non-partitioned
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2565183082
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,13 +283,90 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
- return
-
if get_dialect_name(session) == "postgresql":
return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
Review Comment:
ok done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2555244934
##########
airflow-core/src/airflow/serialization/serialized_objects.py:
##########
@@ -458,6 +471,45 @@ def decode_timetable(var: dict[str, Any]) -> Timetable:
return timetable_class.deserialize(var[Encoding.VAR])
+def _load_partition_mapper(importable_string) -> PartitionMapper | None:
+ if importable_string.startswith("airflow.timetables."):
+ return import_string(importable_string)
+ else:
+ return None
+
+
+def encode_partition_mapper(var: PartitionMapper) -> dict[str, Any]:
+ """
+ Encode a PartitionMapper instance.
+
+ This delegates most of the serialization work to the type, so the behavior
+ can be completely controlled by a custom subclass.
+
+ :meta private:
+ """
+ partition_mapper_class = type(var)
+ importable_string = qualname(partition_mapper_class)
+ if _load_partition_mapper(importable_string) is None:
+ raise _PartitionMapperNotFound(importable_string)
Review Comment:
Why do we need to check this? Isn’t the mapper guarenteed to exist since we just accepted an instance `var`?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2528611179
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,13 +284,77 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
- return
-
if get_dialect_name(session) == "postgresql":
return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
+ @classmethod
+ def _queue_partitioned_dags(
+ cls,
+ asset_id: int,
+ all_dags: set[DagModel],
+ event: AssetEvent,
+ partition_key: str,
+ session: Session,
+ ) -> list[DagModel]:
+ # todo: AIP-76 there may be a better way to identify that timetable is partition-driven
+ partition_dags = [x for x in all_dags if x.timetable_summary == "Partitioned Asset"]
Review Comment:
it's not _all_ dags --- it's like, all dags that are listening to this asset
the important part here is, not just the ones that are partition driven
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2530004636
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
+
+
+class IdentityMapper(PartitionMapper):
+ """Partition mapper that does not change the key."""
+
+ def map(self, key: str) -> str:
+ return key
+
+ def inverse_map(self, key: str) -> Iterable[str]:
+ yield key
+
+
+class PartitionedAssetTimetable(AssetTriggeredTimetable):
+ """Asset-driven timetable that listens for partitioned assets."""
+
+ @property
+ def summary(self) -> str:
+ return "Partitioned Asset"
+
+ def __init__(self, assets: BaseAsset, partition_mapper: PartitionMapper) -> None:
+ super().__init__(assets=assets)
+ self.asset_condition = assets
+ self.partition_mapper = partition_mapper
+
+ def serialize(self) -> dict[str, Any]:
+ from airflow.serialization.serialized_objects import encode_asset_condition
+
+ return {
+ "asset_condition": encode_asset_condition(self.asset_condition),
+ "partition_mapper_cls": self.partition_mapper.__class__.__name__,
+ }
+
+ @classmethod
+ def deserialize(cls, data: dict[str, Any]) -> Timetable:
+ from airflow.serialization.serialized_objects import decode_asset_condition
+
+ ser_asset_condition = data["asset_condition"]
+ mapper_class_name = data.get("partition_mapper_cls", None)
+ if not mapper_class_name:
+ mapper_class = IdentityMapper
+ else:
+ # todo: AIP-76 what is the right way to do this?
+ if mapper_class_name in globals():
+ mapper_class = globals()[mapper_class_name]
+ else:
+ mapper_class = import_string(mapper_class_name)
Review Comment:
One thought occurs to me, we could add something like `encode_timetable` / `decode_timetable` for _partition mappers_. But then, that begs the question, do we need to add partition mappers as a distinct pluggable entity alongside timetables? Maybe so. Appreciate your thoughts on this.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2577518168
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -229,12 +242,39 @@ def notify_asset_alias_created(asset_assets: AssetAlias):
def notify_asset_changed(asset: Asset):
"""Run applicable notification actions when an asset is changed."""
try:
+ # TODO: AIP-76 this will have to change. needs to know *what* happened to the asset (e.g. partition key)
Review Comment:
btw the issue for that is here https://github.com/orgs/apache/projects/542?pane=issue&itemId=140179031&issue=apache%7Cairflow%7C42254
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2565296345
##########
airflow-core/src/airflow/serialization/serialized_objects.py:
##########
@@ -458,6 +471,45 @@ def decode_timetable(var: dict[str, Any]) -> Timetable:
return timetable_class.deserialize(var[Encoding.VAR])
+def _load_partition_mapper(importable_string) -> PartitionMapper | None:
+ if importable_string.startswith("airflow.timetables."):
+ return import_string(importable_string)
+ else:
+ return None
+
+
+def encode_partition_mapper(var: PartitionMapper) -> dict[str, Any]:
+ """
+ Encode a PartitionMapper instance.
+
+ This delegates most of the serialization work to the type, so the behavior
+ can be completely controlled by a custom subclass.
+
+ :meta private:
+ """
+ partition_mapper_class = type(var)
+ importable_string = qualname(partition_mapper_class)
+ if _load_partition_mapper(importable_string) is None:
+ raise _PartitionMapperNotFound(importable_string)
Review Comment:
lmk if concern or if ok---- i'm not wed to 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2539268030
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,13 +284,77 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
- return
-
if get_dialect_name(session) == "postgresql":
return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
+ @classmethod
+ def _queue_partitioned_dags(
+ cls,
+ asset_id: int,
+ all_dags: set[DagModel],
+ event: AssetEvent,
+ partition_key: str,
+ session: Session,
+ ) -> list[DagModel]:
+ # todo: AIP-76 there may be a better way to identify that timetable is partition-driven
+ partition_dags = [x for x in all_dags if x.timetable_summary == "Partitioned Asset"]
Review Comment:
> It would be a bit less performant (need to deserialise unrelated dags) though.
exactly. deserializing a dag is pretty expensive.
so for me i prefer a solution that avoids deserializing dags unnecessarily
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2549750970
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
+
+
+class IdentityMapper(PartitionMapper):
+ """Partition mapper that does not change the key."""
+
+ def map(self, key: str) -> str:
+ return key
+
+ def inverse_map(self, key: str) -> Iterable[str]:
+ yield key
+
+
+class PartitionedAssetTimetable(AssetTriggeredTimetable):
+ """Asset-driven timetable that listens for partitioned assets."""
+
+ @property
+ def summary(self) -> str:
+ return "Partitioned Asset"
+
+ def __init__(self, assets: BaseAsset, partition_mapper: PartitionMapper) -> None:
+ super().__init__(assets=assets)
+ self.asset_condition = assets
+ self.partition_mapper = partition_mapper
+
+ def serialize(self) -> dict[str, Any]:
+ from airflow.serialization.serialized_objects import encode_asset_condition
+
+ return {
+ "asset_condition": encode_asset_condition(self.asset_condition),
+ "partition_mapper_cls": self.partition_mapper.__class__.__name__,
+ }
+
+ @classmethod
+ def deserialize(cls, data: dict[str, Any]) -> Timetable:
+ from airflow.serialization.serialized_objects import decode_asset_condition
+
+ ser_asset_condition = data["asset_condition"]
+ mapper_class_name = data.get("partition_mapper_cls", None)
+ if not mapper_class_name:
+ mapper_class = IdentityMapper
+ else:
+ # todo: AIP-76 what is the right way to do this?
+ if mapper_class_name in globals():
+ mapper_class = globals()[mapper_class_name]
+ else:
+ mapper_class = import_string(mapper_class_name)
Review Comment:
Ok what I have done is update the serde logic to copy exactly what is done for timetables minus the pluggability. so we should be able to add pluggability in the future while maintaining backward compatibility. and i'll create a task for this.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2536795565
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
Review Comment:
I wonder if it’d be better to name these more explicitly like `upstream_to_downstream`
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2539272431
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
+
+
+class IdentityMapper(PartitionMapper):
+ """Partition mapper that does not change the key."""
+
+ def map(self, key: str) -> str:
+ return key
+
+ def inverse_map(self, key: str) -> Iterable[str]:
+ yield key
+
+
+class PartitionedAssetTimetable(AssetTriggeredTimetable):
+ """Asset-driven timetable that listens for partitioned assets."""
+
+ @property
+ def summary(self) -> str:
+ return "Partitioned Asset"
+
+ def __init__(self, assets: BaseAsset, partition_mapper: PartitionMapper) -> None:
+ super().__init__(assets=assets)
+ self.asset_condition = assets
+ self.partition_mapper = partition_mapper
+
+ def serialize(self) -> dict[str, Any]:
+ from airflow.serialization.serialized_objects import encode_asset_condition
+
+ return {
+ "asset_condition": encode_asset_condition(self.asset_condition),
+ "partition_mapper_cls": self.partition_mapper.__class__.__name__,
+ }
+
+ @classmethod
+ def deserialize(cls, data: dict[str, Any]) -> Timetable:
+ from airflow.serialization.serialized_objects import decode_asset_condition
+
+ ser_asset_condition = data["asset_condition"]
+ mapper_class_name = data.get("partition_mapper_cls", None)
+ if not mapper_class_name:
+ mapper_class = IdentityMapper
+ else:
+ # todo: AIP-76 what is the right way to do this?
+ if mapper_class_name in globals():
+ mapper_class = globals()[mapper_class_name]
+ else:
+ mapper_class = import_string(mapper_class_name)
Review Comment:
Any thoughts on this one @Lee-W @uranusjr ?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2547094609
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
Review Comment:
ok i renamed to `to_downstream` and `to_upstream`
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2546973757
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
Review Comment:
@uranusjr @Lee-W
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2542428447
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
Review Comment:
ok 2 against 1 wins
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2548259661
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,65 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
+ """
+ Base partition mapper class.
+
+ Maps keys from asset events to target dag run partitions.
+ """
+
+ def map(self, key: str) -> str: ...
+
+ def inverse_map(self, key: str) -> Iterable[str]: ...
Review Comment:
like 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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2577478846
##########
airflow-core/src/airflow/example_dags/example_outlet_event_extra.py:
##########
@@ -68,6 +68,9 @@ def asset_with_extra_by_context(*, outlet_events=None):
def _asset_with_extra_from_classic_operator_post_execute(context, result):
context["outlet_events"][asset].extra = {"hi": "bye"}
+ # TODO: AIP-76 probably we want to make it so this could be
Review Comment:
my point is to give the task the ability to truthfully record what it does. so it should be able to add or remove events i'd say. or alter the event. or suppress the event (i.e. don't make the event) and yeah i guess one of the bigger differences i'm suggesting is, instead of one and always one event, then it could potentially be more (e.g. multiple partitions affected in one task) or none (nothing done)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "dstandish (via GitHub)" <gi...@apache.org>.
dstandish commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2577450421
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -229,12 +242,39 @@ def notify_asset_alias_created(asset_assets: AssetAlias):
def notify_asset_changed(asset: Asset):
"""Run applicable notification actions when an asset is changed."""
try:
+ # TODO: AIP-76 this will have to change. needs to know *what* happened to the asset (e.g. partition key)
+ # maybe we should just add the event to the signature
+ # or add a new hook `on_asset_event`
get_listener_manager().hook.on_asset_changed(asset=asset)
except Exception:
log.exception("error calling listener")
@classmethod
- def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Session) -> None:
+ def _queue_dagruns(
+ cls,
+ asset_id: int,
Review Comment:
ok
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Update asset manager to be partition-aware [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #58289:
URL: https://github.com/apache/airflow/pull/58289#discussion_r2570609713
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -229,12 +242,39 @@ def notify_asset_alias_created(asset_assets: AssetAlias):
def notify_asset_changed(asset: Asset):
"""Run applicable notification actions when an asset is changed."""
try:
+ # TODO: AIP-76 this will have to change. needs to know *what* happened to the asset (e.g. partition key)
Review Comment:
`on_asset_changed` itself might make more sense for the cases that asset became inactive or active?
`on_asset_event` for these thing makes more sense to me
##########
airflow-core/src/airflow/serialization/serialized_objects.py:
##########
@@ -459,6 +472,45 @@ def decode_timetable(var: dict[str, Any]) -> Timetable:
return timetable_class.deserialize(var[Encoding.VAR])
+def _load_partition_mapper(importable_string) -> PartitionMapper | None:
+ if importable_string.startswith("airflow.timetables."):
+ return import_string(importable_string)
+ else:
+ return None
Review Comment:
```suggestion
if importable_string.startswith("airflow.timetables."):
return import_string(importable_string)
return None
```
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -229,12 +242,39 @@ def notify_asset_alias_created(asset_assets: AssetAlias):
def notify_asset_changed(asset: Asset):
"""Run applicable notification actions when an asset is changed."""
try:
+ # TODO: AIP-76 this will have to change. needs to know *what* happened to the asset (e.g. partition key)
+ # maybe we should just add the event to the signature
+ # or add a new hook `on_asset_event`
get_listener_manager().hook.on_asset_changed(asset=asset)
except Exception:
log.exception("error calling listener")
@classmethod
- def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Session) -> None:
+ def _queue_dagruns(
+ cls,
+ asset_id: int,
Review Comment:
```suggestion
cls,
*,
asset_id: int,
```
since we only call it with keyword argument, not a bad idea to make it keyword only
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -1317,6 +1317,8 @@ def register_asset_changes_in_db(
) -> None:
from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetNameRef, AssetUniqueKey, AssetUriRef
+ # TODO: AIP-76 should we provide an interface to override this?
Review Comment:
Why should we? What's the use case in your mind
##########
airflow-core/src/airflow/example_dags/example_outlet_event_extra.py:
##########
@@ -68,6 +68,9 @@ def asset_with_extra_by_context(*, outlet_events=None):
def _asset_with_extra_from_classic_operator_post_execute(context, result):
context["outlet_events"][asset].extra = {"hi": "bye"}
+ # TODO: AIP-76 probably we want to make it so this could be
Review Comment:
In the current syntax, `context["outlet_events"][asset]` works more like operating events of asset. Are you suggesting something like
`context["outlet_events"][asset].events = ...`
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,15 +283,94 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
- return
-
if get_dialect_name(session) == "postgresql":
- return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
- return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
+ return cls._queue_dagruns_nonpartitioned_postgres(asset_id, dags_to_queue, session)
+ return cls._queue_dagruns_nonpartitioned_slow_path(asset_id, dags_to_queue, session)
@classmethod
- def _slow_path_queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Session) -> None:
+ def _queue_partitioned_dags(
+ cls,
+ asset_id: int,
+ all_dags: set[DagModel],
+ event: AssetEvent,
+ partition_key: str | None,
+ session: Session,
+ ) -> list[DagModel]:
+ # TODO: AIP-76 there may be a better way to identify that timetable is partition-driven
+ partition_dags = [x for x in all_dags if x.timetable_summary == "Partitioned Asset"]
+ if partition_dags and not partition_key:
+ # TODO: AIP-76 how to best ensure users can see this? Probably add Log record.
+ log.warning(
+ "Listening dags are partition-aware but run has no partition key",
+ listening_dags=[x.dag_id for x in partition_dags],
+ asset_id=asset_id,
+ run_id=event.source_run_id,
+ dag_id=event.source_dag_id,
+ task_id=event.source_task_id,
+ )
+ return partition_dags
+
+ for target_dag in partition_dags:
+ if TYPE_CHECKING:
+ assert partition_key is not None
+ from airflow.models.serialized_dag import SerializedDagModel
+
+ serdag = SerializedDagModel.get(dag_id=target_dag.dag_id, session=session)
+ if not serdag:
+ raise RuntimeError(f"Could not find serialized dag for dag_id={target_dag.dag_id}")
+ timetable = serdag.dag.timetable
+ if TYPE_CHECKING:
+ assert isinstance(timetable, PartitionedAssetTimetable)
+ target_key = timetable.partition_mapper.to_downstream(partition_key)
+
+ apdr = cls._get_or_create_apdr(
+ target_key=target_key,
+ target_dag=target_dag,
+ session=session,
+ )
+ log_record = PartitionedAssetKeyLog(
+ asset_id=asset_id,
+ asset_event_id=event.id,
+ asset_partition_dag_run_id=apdr.id,
+ source_partition_key=partition_key,
+ target_dag_id=target_dag.dag_id,
+ target_partition_key=target_key,
+ )
+ session.add(log_record)
+ return partition_dags
+
+ @classmethod
+ def _get_or_create_apdr(
+ cls,
+ target_key: str,
+ target_dag: SerializedDagModel,
+ session: Session,
+ ):
Review Comment:
```suggestion
) -> AssetPartitionDagRun:
```
##########
airflow-core/src/airflow/timetables/simple.py:
##########
@@ -217,3 +218,66 @@ def next_dagrun_info(
restriction: TimeRestriction,
) -> DagRunInfo | None:
return None
+
+
+class PartitionMapper:
Review Comment:
```suggestion
class PartitionMapper(ABC):
```
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -243,15 +283,94 @@ def _queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Se
# "fallback" to running this in a nested transaction. This is needed
# so that the adding of these rows happens in the same transaction
# where `ti.state` is changed.
- if not dags_to_queue:
- return
-
if get_dialect_name(session) == "postgresql":
- return cls._postgres_queue_dagruns(asset_id, dags_to_queue, session)
- return cls._slow_path_queue_dagruns(asset_id, dags_to_queue, session)
+ return cls._queue_dagruns_nonpartitioned_postgres(asset_id, dags_to_queue, session)
+ return cls._queue_dagruns_nonpartitioned_slow_path(asset_id, dags_to_queue, session)
@classmethod
- def _slow_path_queue_dagruns(cls, asset_id: int, dags_to_queue: set[DagModel], session: Session) -> None:
+ def _queue_partitioned_dags(
+ cls,
+ asset_id: int,
+ all_dags: set[DagModel],
+ event: AssetEvent,
+ partition_key: str | None,
+ session: Session,
+ ) -> list[DagModel]:
+ # TODO: AIP-76 there may be a better way to identify that timetable is partition-driven
+ partition_dags = [x for x in all_dags if x.timetable_summary == "Partitioned Asset"]
+ if partition_dags and not partition_key:
+ # TODO: AIP-76 how to best ensure users can see this? Probably add Log record.
+ log.warning(
+ "Listening dags are partition-aware but run has no partition key",
+ listening_dags=[x.dag_id for x in partition_dags],
+ asset_id=asset_id,
+ run_id=event.source_run_id,
+ dag_id=event.source_dag_id,
+ task_id=event.source_task_id,
+ )
Review Comment:
Should we move these logic to `_queue_dagruns` and run the following partition Dag logic in `_queue_partitioned_dags` only?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org