uranusjr commented on code in PR #76827:
URL: https://github.com/apache/airflow/pull/76827#discussion_r4546406754
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -2652,8 +2652,12 @@ def _asset_event_extras_from_aliases() ->
dict[tuple[SerializedAssetUniqueKey, s
)
if event is None:
ti.log.info("Dynamically creating AssetModel %s",
asset_key)
- session.add(AssetModel.from_serialized(asset))
- session.flush() # So event can set up its asset fk.
+ # Another task can create the same asset between our
lookup and this insert.
+ try:
+ with session.begin_nested():
+ session.add(AssetModel.from_serialized(asset))
+ except IntegrityError:
+ ti.log.info("AssetModel %s was created concurrently;
reusing it", asset_key)
asset_manager.register_asset_change(
Review Comment:
If this call returns None after an IntegrityError (no asset added), the
event is silently dropped, causing the exact issue. Either raise, or at least
log an error, in theiscase.
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -65,6 +65,7 @@
log = structlog.get_logger(__name__)
_APDR_GET_OR_CREATE_ATTEMPTS = 2
+_ALIAS_ASSOCIATION_ATTEMPTS = 3
Review Comment:
Why 3 here? What’s different from the APDR case, which only needed 2?
##########
airflow-core/src/airflow/models/taskinstance.py:
##########
@@ -2652,8 +2652,12 @@ def _asset_event_extras_from_aliases() ->
dict[tuple[SerializedAssetUniqueKey, s
)
if event is None:
ti.log.info("Dynamically creating AssetModel %s",
asset_key)
- session.add(AssetModel.from_serialized(asset))
- session.flush() # So event can set up its asset fk.
+ # Another task can create the same asset between our
lookup and this insert.
Review Comment:
```suggestion
```
Catching IntegrityError is straightfoward enough IMO.
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -112,16 +113,32 @@ def _add_asset_alias_association(
*,
session: Session,
) -> None:
- already_related = {m.name for m in asset_model.aliases}
- existing_aliases = {
- m.name: m
- for m in
session.scalars(select(AssetAliasModel).where(AssetAliasModel.name.in_(alias_names)))
- }
- asset_model.aliases.extend(
- existing_aliases.get(name, AssetAliasModel(name=name))
- for name in alias_names
- if name not in already_related
- )
+ # Another task can link the same alias between our read of
asset_model.aliases and the
+ # insert. The savepoint keeps that conflict from failing the whole
transaction; after a
+ # rollback the aliases are reloaded and only the still-missing links
are added.
Review Comment:
```suggestion
# Tolerate a concurrent writer linking the same alias.
```
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -112,16 +113,32 @@ def _add_asset_alias_association(
*,
session: Session,
) -> None:
- already_related = {m.name for m in asset_model.aliases}
- existing_aliases = {
- m.name: m
- for m in
session.scalars(select(AssetAliasModel).where(AssetAliasModel.name.in_(alias_names)))
- }
- asset_model.aliases.extend(
- existing_aliases.get(name, AssetAliasModel(name=name))
- for name in alias_names
- if name not in already_related
- )
+ # Another task can link the same alias between our read of
asset_model.aliases and the
+ # insert. The savepoint keeps that conflict from failing the whole
transaction; after a
+ # rollback the aliases are reloaded and only the still-missing links
are added.
+ for attempt in range(1, _ALIAS_ASSOCIATION_ATTEMPTS + 1):
+ already_related = {m.name for m in asset_model.aliases}
+ missing = [name for name in alias_names if name not in
already_related]
+ if not missing:
+ return
+ existing_aliases = {
+ m.name: m
+ for m in
session.scalars(select(AssetAliasModel).where(AssetAliasModel.name.in_(missing)))
+ }
+ try:
+ with session.begin_nested():
+ asset_model.aliases.extend(
+ existing_aliases.get(name, AssetAliasModel(name=name))
for name in missing
+ )
+ except exc.IntegrityError:
+ if attempt == _ALIAS_ASSOCIATION_ATTEMPTS:
+ raise
+ cls.logger().debug(
+ "Another process linked aliases %s to asset %s first;
retrying",
+ missing,
+ asset_model.uri,
+ )
+ session.expire(asset_model, ["aliases"])
Review Comment:
```suggestion
```
Rolling back the savepoint already expires objects modified inside it.
##########
airflow-core/src/airflow/assets/manager.py:
##########
@@ -112,16 +113,32 @@ def _add_asset_alias_association(
*,
session: Session,
) -> None:
- already_related = {m.name for m in asset_model.aliases}
- existing_aliases = {
- m.name: m
- for m in
session.scalars(select(AssetAliasModel).where(AssetAliasModel.name.in_(alias_names)))
- }
- asset_model.aliases.extend(
- existing_aliases.get(name, AssetAliasModel(name=name))
- for name in alias_names
- if name not in already_related
- )
+ # Another task can link the same alias between our read of
asset_model.aliases and the
+ # insert. The savepoint keeps that conflict from failing the whole
transaction; after a
+ # rollback the aliases are reloaded and only the still-missing links
are added.
+ for attempt in range(1, _ALIAS_ASSOCIATION_ATTEMPTS + 1):
+ already_related = {m.name for m in asset_model.aliases}
+ missing = [name for name in alias_names if name not in
already_related]
+ if not missing:
+ return
+ existing_aliases = {
+ m.name: m
+ for m in
session.scalars(select(AssetAliasModel).where(AssetAliasModel.name.in_(missing)))
+ }
+ try:
+ with session.begin_nested():
+ asset_model.aliases.extend(
+ existing_aliases.get(name, AssetAliasModel(name=name))
for name in missing
Review Comment:
Would it be possible for `AssetAliasModel(name=name)` to also conflict?
--
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]
