This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new bf022cc64f2 Fix duplicate pending Dag runs for the same asset
partition key (#71074)
bf022cc64f2 is described below
commit bf022cc64f2a0f621922b6ed2ddfba59c0c705d4
Author: Anmol Mishra <[email protected]>
AuthorDate: Thu Oct 22 13:11:26 2026 +0530
Fix duplicate pending Dag runs for the same asset partition key (#71074)
Co-authored-by: Kaxil Naik <[email protected]>
Co-authored-by: anmolxlight <[email protected]>
Co-authored-by: Wei Lee <[email protected]>
---
airflow-core/docs/migrations-ref.rst | 5 +-
airflow-core/src/airflow/assets/manager.py | 281 ++++++++++----------
...0148_3_4_0_add_pending_partition_key_to_apdr.py | 84 ++++++
..._job.py => 0149_3_5_0_add_session_id_to_job.py} | 4 +-
airflow-core/src/airflow/models/asset.py | 33 ++-
airflow-core/src/airflow/utils/db.py | 2 +-
airflow-core/tests/unit/assets/test_manager.py | 283 +++++++++++++++++++--
.../test_0148_add_pending_partition_key_to_apdr.py | 187 ++++++++++++++
8 files changed, 714 insertions(+), 165 deletions(-)
diff --git a/airflow-core/docs/migrations-ref.rst
b/airflow-core/docs/migrations-ref.rst
index 9a7f31318f7..922105789ae 100644
--- a/airflow-core/docs/migrations-ref.rst
+++ b/airflow-core/docs/migrations-ref.rst
@@ -39,7 +39,10 @@ Here's the list of all the Database Migrations that are
executed via when you ru
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| Revision ID | Revises ID | Airflow Version | Description
|
+=========================+==================+===================+==============================================================+
-| ``c71065a1b14f`` (head) | ``9cffe90fa0c0`` | ``3.5.0`` | Add
session_id and registration_id to Job. |
+| ``c71065a1b14f`` (head) | ``ee86eed19e24`` | ``3.5.0`` | Add
session_id and registration_id to Job. |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``ee86eed19e24`` | ``9cffe90fa0c0`` | ``3.4.0`` | Add
pending_partition_key to asset_partition_dag_run and |
+| | | | enforce
single pending row per key. |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| ``9cffe90fa0c0`` | ``a5d7b9c13e40`` | ``3.4.0`` | Add
commit_allows_update to connection_test_request. |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
diff --git a/airflow-core/src/airflow/assets/manager.py
b/airflow-core/src/airflow/assets/manager.py
index ad133881db7..eb8ea92d207 100644
--- a/airflow-core/src/airflow/assets/manager.py
+++ b/airflow-core/src/airflow/assets/manager.py
@@ -18,7 +18,6 @@
from __future__ import annotations
from collections.abc import Callable, Collection, Iterable
-from contextlib import contextmanager
from functools import partial
from typing import TYPE_CHECKING
@@ -48,7 +47,7 @@ from airflow.models.log import Log
from airflow.timetables.base import compute_rollup_fingerprint
from airflow.utils.helpers import is_container, prune_dict
from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.utils.sqlalchemy import get_dialect_name, with_row_locks
+from airflow.utils.sqlalchemy import get_dialect_name
if TYPE_CHECKING:
from datetime import datetime
@@ -66,59 +65,7 @@ if TYPE_CHECKING:
log = structlog.get_logger(__name__)
-
-@contextmanager
-def _lock_asset_model(
- *,
- session: Session,
- asset_id: int,
- max_retries: int = 10,
- retry_delay: float = 0.1,
-):
- """
- Context manager to acquire a lock for AssetPartitionDagRun creation.
-
- - SQLite: Use a no-op ORM update to trigger a write-transaction and
acquire SQLite's global writer lock.
- - Postgres/MySQL: uses row-level lock on AssetModel.
- """
- if get_dialect_name(session) == "sqlite":
- import time
-
- from sqlalchemy import update
-
- # no-op update
- # This is used to acquire SQLite's global writer lock.
- stmt = update(AssetModel).where(AssetModel.id ==
asset_id).values(id=AssetModel.id)
- for _ in range(max_retries):
- try:
- session.execute(stmt)
- session.flush()
- except exc.OperationalError as err:
- err_msg = str(err).lower()
- if "locked" in err_msg or "busy" in err_msg:
- session.rollback()
- time.sleep(retry_delay)
- continue
-
- # lock acquired
- yield
- return
-
- raise RuntimeError(f"Could not acquire SQLite AssetModel writer lock
for asset_id={asset_id}")
- else:
- # Postgres/MySQL row-level lock
- if (
- session.scalar(
- with_row_locks(
- query=select(AssetModel.id).where(AssetModel.id ==
asset_id),
- session=session,
- key_share=True,
- )
- )
- ) is None:
- raise RuntimeError(f"Asset {asset_id} does not exist – cannot
lock.")
-
- yield
+_APDR_GET_OR_CREATE_ATTEMPTS = 2
class AssetManager(LoggingMixin):
@@ -524,7 +471,12 @@ class AssetManager(LoggingMixin):
if not dags_to_queue:
return None
- partition_dags = [x for x in dags_to_queue if x.timetable_partitioned
is True]
+ # An APDR INSERT that collides with another transaction's uncommitted
row on the
+ # pending-key unique constraint waits for it. Sorting keeps two events
from inserting
+ # rows for the same pair of Dags in opposite orders and deadlocking.
+ partition_dags = sorted(
+ (x for x in dags_to_queue if x.timetable_partitioned is True),
key=lambda d: d.dag_id
+ )
cls._queue_partitioned_dags(
asset_id=asset_id,
partition_dags=partition_dags,
@@ -693,7 +645,6 @@ class AssetManager(LoggingMixin):
target_partition_date=target_partition_date,
target_dag=target_dag,
rollup_fingerprint=fingerprint,
- asset_id=asset_id,
allow_reuse=timetable.batch_asset_events,
session=session,
)
@@ -715,28 +666,143 @@ class AssetManager(LoggingMixin):
target_partition_date: datetime | None,
target_dag: DagModel,
rollup_fingerprint: dict,
- asset_id: int,
allow_reuse: bool,
session: Session,
) -> AssetPartitionDagRun:
"""
Get or create an APDR.
- If 2 processes invoke this method at the same time using the same
(target_key, target_dag) pair,
- they may both check the database and, finding no existing APDR, create
separate instances.
- This leads to the unintended outcome of having two APDRs created
instead of one.
- To resolve this, we add a mutex lock to AssetModel for PostgreSQL and
MySQL and use
- AssetPartitionDagRunMutexLock table for SQLite.
+ When ``allow_reuse=True``, an existing pending APDR for the same
+ ``(target_dag, partition_key)`` is reused — multiple events accumulate
on one
+ APDR. When ``allow_reuse=False`` (set when the timetable's
``batch_asset_events``
+ is ``False``), a new APDR is always created so each event gets its own
APDR
+ and the scheduler produces one DagRun per event.
+
+ With reuse, two processes handling events for the same (target_key,
target_dag) pair
+ may both find no pending APDR and both try to create one. The events
can come from
+ different producer assets, so nothing they share can be locked to
serialize this
+ short of the Dag row. Instead, a reusable APDR is created with
``pending_partition_key``
+ set, the unique constraint on (target_dag_id, pending_partition_key)
rejects the second
+ INSERT, and the loser reuses the winner's row.
``rollup_fingerprint`` is the serialized mapper / window definition
for all partitioned
assets in the timetable at creation time; the scheduler discards APDRs
whose stamp no
longer matches the current timetable's fingerprint (mapper / window
may have changed).
+ """
+ if allow_reuse and (
+ pending_apdr := cls._get_latest_pending_apdr(
+ target_key=target_key, target_dag_id=target_dag.dag_id,
session=session
+ )
+ ):
+ cls._reconcile_partition_date(
+ apdr=pending_apdr,
+ target_partition_date=target_partition_date,
+ target_dag_id=target_dag.dag_id,
+ target_key=target_key,
+ session=session,
+ )
+ cls.logger().debug(
+ "Existing APDR found for key %s dag_id %s",
+ target_key,
+ target_dag.dag_id,
+ exc_info=True,
+ )
+ return pending_apdr
+
+ build_apdr = partial(
+ AssetPartitionDagRun,
+ target_dag_id=target_dag.dag_id,
+ created_dag_run_id=None,
+ partition_key=target_key,
+ pending_partition_key=target_key if allow_reuse else None,
+ partition_date=target_partition_date,
+ rollup_fingerprint=rollup_fingerprint,
+ )
+ if not allow_reuse:
+ apdr = build_apdr()
+ session.add(apdr)
+ session.flush()
+ else:
+ # begin_nested() flushes the caller's pending rows before the
savepoint starts, so
+ # rolling it back on a conflict expunges only ``apdr``. The
re-lookup after a
+ # conflict sees the winner's row on MySQL only because it runs at
READ COMMITTED
+ # (settings.py), unlike its REPEATABLE READ default.
+ for attempt in range(1, _APDR_GET_OR_CREATE_ATTEMPTS + 1):
+ apdr = build_apdr()
+ try:
+ with session.begin_nested():
+ session.add(apdr)
+ except exc.IntegrityError:
+ winner = cls._get_latest_pending_apdr(
+ target_key=target_key,
target_dag_id=target_dag.dag_id, session=session
+ )
+ if winner is None:
+ # The winner was already turned into a Dag run, which
cleared its key and
+ # freed the constraint, so the INSERT can be run again.
+ if attempt == _APDR_GET_OR_CREATE_ATTEMPTS:
+ raise
+ continue
+ cls.logger().debug(
+ "Another process created the APDR for key %s dag_id %s
first; reusing it",
+ target_key,
+ target_dag.dag_id,
+ )
+ cls._reconcile_partition_date(
+ apdr=winner,
+ target_partition_date=target_partition_date,
+ target_dag_id=target_dag.dag_id,
+ target_key=target_key,
+ session=session,
+ )
+ return winner
+ break
+
+ cls.logger().debug(
+ "No existing APDR found. Create APDR for key %s dag_id %s",
+ target_key,
+ target_dag.dag_id,
+ exc_info=True,
+ )
+ return apdr
+
+ @staticmethod
+ def _get_latest_pending_apdr(
+ *, target_key: str, target_dag_id: str, session: Session
+ ) -> AssetPartitionDagRun | None:
+ """
+ Return the latest pending APDR for (target_dag_id, target_key), if any.
+
+ This also matches pending rows created while batching was disabled,
which carry no
+ ``pending_partition_key``, so a Dag that turns batching on keeps
accumulating onto
+ its most recent pending row.
+ """
+ return session.scalar(
+ select(AssetPartitionDagRun)
+ .where(
+ AssetPartitionDagRun.partition_key == target_key,
+ AssetPartitionDagRun.target_dag_id == target_dag_id,
+ AssetPartitionDagRun.created_dag_run_id.is_(None),
+ )
+ .order_by(AssetPartitionDagRun.id.desc())
+ .limit(1)
+ )
- Reconciling the carried ``partition_date`` on an existing pending APDR
is best-effort:
- a partitioned consumer's feeding assets are expected to agree on the
partition's
- datetime. The carry only matters for ``IdentityMapper`` (whose key the
scheduler
- cannot decode); temporal/composite feeds re-derive the date from the
key at run
- creation regardless of what is stored here. Within that contract:
+ @staticmethod
+ def _reconcile_partition_date(
+ *,
+ apdr: AssetPartitionDagRun,
+ target_partition_date: datetime | None,
+ target_dag_id: str,
+ target_key: str,
+ session: Session,
+ ) -> None:
+ """
+ Reconcile the carried ``partition_date`` on a reused pending APDR.
+
+ This is best-effort: a partitioned consumer's feeding assets are
expected to agree
+ on the partition's datetime. The carry only matters for
``IdentityMapper`` (whose key
+ the scheduler cannot decode); temporal/composite feeds re-derive the
date from the key
+ at run creation regardless of what is stored here. Within that
contract:
- If the APDR carries no date yet (``None`` — created by an event that
carried none),
adopt the incoming date when this event carries one. There is
nothing to conflict
@@ -745,73 +811,24 @@ class AssetManager(LoggingMixin):
one, the producing assets disagree; picking one would be
order-dependent, so the
carried date is suppressed to ``None`` (and re-adoptable by a later
event).
- Otherwise (the dates agree, or this event carries none) the existing
value is kept.
-
- When ``allow_reuse=True``, an existing pending APDR for the same
- ``(target_dag, partition_key)`` is reused — multiple events accumulate
on one
- APDR. When ``allow_reuse=False`` (set when the timetable's
``batch_asset_events``
- is ``False``), a new APDR is always created so each event gets its own
APDR
- and the scheduler produces one DagRun per event.
"""
- with _lock_asset_model(session=session, asset_id=asset_id):
- latest_apdr: AssetPartitionDagRun | None = None
- if allow_reuse:
- latest_apdr = session.scalar(
- select(AssetPartitionDagRun)
- .where(
- AssetPartitionDagRun.partition_key == target_key,
- AssetPartitionDagRun.target_dag_id ==
target_dag.dag_id,
- )
- .order_by(AssetPartitionDagRun.id.desc())
- .limit(1)
- )
- if latest_apdr and latest_apdr.created_dag_run_id is None:
- existing_partition_date = latest_apdr.partition_date
- if existing_partition_date is None:
- # No carried date yet; adopt the incoming one if present
(no conflict
- # to resolve). Keeps a later identity event's date from
being dropped.
- if target_partition_date is not None:
- latest_apdr.partition_date = target_partition_date
- session.flush()
- elif target_partition_date is not None and
existing_partition_date != target_partition_date:
- # Two contributing events carry conflicting
partition_dates for the same
- # (target_key, target_dag). Choosing one would be
order-dependent, so
- # suppress: the consumer DagRun gets partition_date=None
rather than a
- # wrong, unstable value.
- log.warning(
- "Conflicting partition_date carried for the same
target key; "
- "suppressing it so the consumer DagRun's
partition_date is None. "
- "The producing assets likely disagree on the
partition's datetime.",
- target_dag_id=target_dag.dag_id,
- target_key=target_key,
- existing_partition_date=existing_partition_date,
- incoming_partition_date=target_partition_date,
- )
- latest_apdr.partition_date = None
- session.flush()
- cls.logger().debug(
- "Existing APDR found for key %s dag_id %s",
- target_key,
- target_dag.dag_id,
- exc_info=True,
- )
- return latest_apdr
-
- apdr = AssetPartitionDagRun(
- target_dag_id=target_dag.dag_id,
- created_dag_run_id=None,
- partition_key=target_key,
- partition_date=target_partition_date,
- rollup_fingerprint=rollup_fingerprint,
+ existing_partition_date = apdr.partition_date
+ if existing_partition_date is None:
+ if target_partition_date is not None:
+ apdr.partition_date = target_partition_date
+ session.flush()
+ elif target_partition_date is not None and existing_partition_date !=
target_partition_date:
+ log.warning(
+ "Conflicting partition_date carried for the same target key; "
+ "suppressing it so the consumer DagRun's partition_date is
None. "
+ "The producing assets likely disagree on the partition's
datetime.",
+ target_dag_id=target_dag_id,
+ target_key=target_key,
+ existing_partition_date=existing_partition_date,
+ incoming_partition_date=target_partition_date,
)
- session.add(apdr)
+ apdr.partition_date = None
session.flush()
- cls.logger().debug(
- "No existing APDR found. Create APDR for key %s dag_id %s",
- target_key,
- target_dag.dag_id,
- exc_info=True,
- )
- return apdr
@classmethod
def _queue_dagruns_nonpartitioned(
diff --git
a/airflow-core/src/airflow/migrations/versions/0148_3_4_0_add_pending_partition_key_to_apdr.py
b/airflow-core/src/airflow/migrations/versions/0148_3_4_0_add_pending_partition_key_to_apdr.py
new file mode 100644
index 00000000000..0b70761ff50
--- /dev/null
+++
b/airflow-core/src/airflow/migrations/versions/0148_3_4_0_add_pending_partition_key_to_apdr.py
@@ -0,0 +1,84 @@
+#
+# 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.
+
+"""
+Add pending_partition_key to asset_partition_dag_run and enforce single
pending row per key.
+
+Asset events from two producer assets that map to the same downstream
partition key
+could each create a pending AssetPartitionDagRun (APDR), splitting the events
between
+two rows that each wait forever for the other's events.
``pending_partition_key``
+mirrors ``partition_key`` on a pending row, and a unique constraint on
+(target_dag_id, pending_partition_key) stops a second one from being created.
It stands
+in for a partial unique index, which MySQL does not support.
+
+Existing pending rows are neither merged nor deleted: a Dag without
``batch_asset_events``
+legitimately has several pending rows for one key, one per event. Only the
latest (highest id)
+pending row per (target_dag_id, partition_key) is marked; older duplicates
keep a NULL
+``pending_partition_key``, which the unique constraint treats as distinct, so
they do not
+conflict. A pending row split by the old race is not healed by this migration.
This is plain
+SQL, so it behaves the same in offline mode.
+
+Revision ID: ee86eed19e24
+Revises: 9cffe90fa0c0
+Create Date: 2026-08-04 00:00:00.000000
+
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+from airflow.migrations.db_types import StringID
+from airflow.migrations.utils import disable_sqlite_fkeys
+
+revision = "ee86eed19e24"
+down_revision = "9cffe90fa0c0"
+branch_labels = None
+depends_on = None
+airflow_version = "3.4.0"
+
+_TABLE = "asset_partition_dag_run"
+_UQ_NAME = "apdr_target_dag_id_pending_partition_key_uq"
+
+
+def upgrade():
+ """Add pending_partition_key to asset_partition_dag_run and enforce single
pending row per key."""
+ with disable_sqlite_fkeys(op):
+ with op.batch_alter_table(_TABLE, schema=None) as batch_op:
+ batch_op.add_column(sa.Column("pending_partition_key", StringID(),
nullable=True))
+
+ # The derived table lets MySQL read the table it updates.
+ op.execute(
+ f"UPDATE {_TABLE} SET pending_partition_key = partition_key WHERE
id IN ("
+ "SELECT latest_id FROM ("
+ f"SELECT MAX(id) AS latest_id FROM {_TABLE} WHERE
created_dag_run_id IS NULL "
+ "GROUP BY target_dag_id, partition_key"
+ ") latest_pending)"
+ )
+
+ with op.batch_alter_table(_TABLE, schema=None) as batch_op:
+ batch_op.create_unique_constraint(_UQ_NAME, ["target_dag_id",
"pending_partition_key"])
+
+
+def downgrade():
+ """Drop the pending-partition uniqueness guard and pending_partition_key
column."""
+ with disable_sqlite_fkeys(op):
+ with op.batch_alter_table(_TABLE, schema=None) as batch_op:
+ batch_op.drop_constraint(_UQ_NAME, type_="unique")
+ batch_op.drop_column("pending_partition_key")
diff --git
a/airflow-core/src/airflow/migrations/versions/0148_3_5_0_add_session_id_to_job.py
b/airflow-core/src/airflow/migrations/versions/0149_3_5_0_add_session_id_to_job.py
similarity index 97%
rename from
airflow-core/src/airflow/migrations/versions/0148_3_5_0_add_session_id_to_job.py
rename to
airflow-core/src/airflow/migrations/versions/0149_3_5_0_add_session_id_to_job.py
index 93fa94f46e2..e595a52b12a 100644
---
a/airflow-core/src/airflow/migrations/versions/0148_3_5_0_add_session_id_to_job.py
+++
b/airflow-core/src/airflow/migrations/versions/0149_3_5_0_add_session_id_to_job.py
@@ -20,7 +20,7 @@
Add session_id and registration_id to Job.
Revision ID: c71065a1b14f
-Revises: 9cffe90fa0c0
+Revises: ee86eed19e24
Create Date: 2026-10-05 08:40:20.000000
"""
@@ -33,7 +33,7 @@ from alembic import op
from airflow.migrations.utils import disable_sqlite_fkeys
revision = "c71065a1b14f"
-down_revision = "9cffe90fa0c0"
+down_revision = "ee86eed19e24"
branch_labels = None
depends_on = None
airflow_version = "3.5.0"
diff --git a/airflow-core/src/airflow/models/asset.py
b/airflow-core/src/airflow/models/asset.py
index ab780be2e77..df845b53a9f 100644
--- a/airflow-core/src/airflow/models/asset.py
+++ b/airflow-core/src/airflow/models/asset.py
@@ -33,11 +33,12 @@ from sqlalchemy import (
PrimaryKeyConstraint,
String,
Table,
+ UniqueConstraint,
delete,
select,
)
from sqlalchemy.ext.associationproxy import association_proxy
-from sqlalchemy.orm import Mapped, mapped_column, relationship
+from sqlalchemy.orm import Mapped, mapped_column, relationship, validates
from airflow._shared.timezones import timezone
from airflow.configuration import conf as airflow_conf
@@ -950,16 +951,26 @@ class AssetPartitionDagRun(Base):
gets created leaves its APDR record behind.
Where created_dag_run_id is null, the Dag run has not yet been created.
- With event batching enabled, `_lock_asset_model` protects reuse of the
- latest pending row for a target_dag_id / partition_key pair. With batching
- disabled, each event creates its own row, so multiple pending rows for
- the same pair are expected.
+ With batching disabled, each event creates its own row, so multiple pending
+ rows for the same pair are expected. With event batching enabled, events
+ accumulate on one pending row per target_dag_id / partition_key, and a
+ unique constraint on (target_dag_id, pending_partition_key) keeps
+ concurrent writers from creating a second one.
+
+ pending_partition_key mirrors partition_key on a pending row created with
+ batching enabled, and is null otherwise; the `_sync_pending_partition_key`
+ validator clears it once the Dag run is created. A unique index treats null
+ as distinct from every value on all three supported backends, so it only
+ compares batched pending rows. It stands in for a partial unique index on
+ (target_dag_id, partition_key) WHERE created_dag_run_id IS NULL, which
+ MySQL does not support.
"""
id: Mapped[int] = mapped_column(Integer, primary_key=True,
autoincrement=True)
target_dag_id: Mapped[str] = mapped_column(StringID(), nullable=False)
created_dag_run_id: Mapped[int | None] = mapped_column(Integer(),
nullable=True)
partition_key: Mapped[str] = mapped_column(StringID(), nullable=False)
+ pending_partition_key: Mapped[str | None] = mapped_column(StringID(),
nullable=True)
partition_date: Mapped[datetime | None] = mapped_column(UtcDateTime,
nullable=True)
# Serialized snapshot of the rollup definition (mapper + window for every
# partitioned asset in the timetable) at the time this APDR was created.
@@ -984,8 +995,20 @@ class AssetPartitionDagRun(Base):
name="apdr_created_dag_run_id_fkey",
ondelete="CASCADE",
),
+ UniqueConstraint(
+ "target_dag_id",
+ "pending_partition_key",
+ name="apdr_target_dag_id_pending_partition_key_uq",
+ ),
)
+ @validates("created_dag_run_id")
+ def _sync_pending_partition_key(self, key: str, value: int | None) -> int
| None:
+ """Clear pending_partition_key once a Dag run is created for this
row."""
+ if value is not None:
+ self.pending_partition_key = None
+ return value
+
class PartitionedAssetKeyLog(Base):
"""
diff --git a/airflow-core/src/airflow/utils/db.py
b/airflow-core/src/airflow/utils/db.py
index 6e078f3b0e9..266222bd070 100644
--- a/airflow-core/src/airflow/utils/db.py
+++ b/airflow-core/src/airflow/utils/db.py
@@ -117,7 +117,7 @@ _REVISION_HEADS_MAP: dict[str, str] = {
"3.1.8": "509b94a1042d",
"3.2.0": "1d6611b6ab7c",
"3.3.0": "d2f4e1b3c5a7",
- "3.4.0": "9cffe90fa0c0",
+ "3.4.0": "ee86eed19e24",
"3.5.0": "c71065a1b14f",
}
diff --git a/airflow-core/tests/unit/assets/test_manager.py
b/airflow-core/tests/unit/assets/test_manager.py
index 049df518569..c700605c24f 100644
--- a/airflow-core/tests/unit/assets/test_manager.py
+++ b/airflow-core/tests/unit/assets/test_manager.py
@@ -20,12 +20,13 @@ from __future__ import annotations
import concurrent.futures
import itertools
import logging
+import threading
from collections import Counter
from typing import TYPE_CHECKING
from unittest import mock
import pytest
-from sqlalchemy import func, select
+from sqlalchemy import event, exc, func, insert, select
from sqlalchemy.dialects import mysql, postgresql, sqlite
from sqlalchemy.orm import Session
@@ -43,6 +44,7 @@ from airflow.models.asset import (
DagScheduleAssetNameReference,
DagScheduleAssetReference,
DagScheduleAssetUriReference,
+ PartitionedAssetKeyLog,
)
from airflow.models.dag import DAG, DagModel
from airflow.models.dagbundle import DagBundleModel
@@ -54,6 +56,7 @@ from airflow.partition_mappers.window import WeekWindow
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.sdk.definitions.asset import Asset
from airflow.sdk.definitions.timetables.assets import PartitionedAssetTimetable
+from airflow.utils.session import create_session
from tests_common.test_utils.asserts import capture_orm_selects
from tests_common.test_utils.config import conf_vars
@@ -417,9 +420,8 @@ class TestAssetManager:
@pytest.mark.usefixtures("dag_maker", "testing_dag_bundle")
def test_get_or_create_apdr_race_condition(self, session, caplog):
- asm = AssetModel(uri="test://asset1/", name="partition_asset",
group="asset")
testing_dag = DagModel(dag_id="testing_dag", is_stale=False,
bundle_name="testing")
- session.add_all([asm, testing_dag])
+ session.add(testing_dag)
session.commit()
session.flush()
assert
session.scalar(select(func.count()).select_from(AssetPartitionDagRun)) == 0
@@ -439,7 +441,6 @@ class TestAssetManager:
target_partition_date=None,
target_dag=testing_dag,
rollup_fingerprint=rollup_fingerprint,
- asset_id=asm.id,
allow_reuse=True,
session=_session,
).id
@@ -452,14 +453,68 @@ class TestAssetManager:
with
concurrent.futures.ThreadPoolExecutor(max_workers=thread_count) as pool:
ids = pool.map(lambda _: _get_or_create_apdr(), [None] *
thread_count)
- assert Counter(r.msg for r in caplog.records) == {
- "Existing APDR found for key test_partition_key dag_id
testing_dag": thread_count - 1,
- "No existing APDR found. Create APDR for key test_partition_key
dag_id testing_dag": 1,
- }
+ # A thread either finds the row with its lookup, or loses the INSERT
to the unique
+ # constraint and reuses the winner's row; which one happens is a
timing detail.
+ msg_counts = Counter(r.msg for r in caplog.records)
+ create_msg = "No existing APDR found. Create APDR for key
test_partition_key dag_id testing_dag"
+ found_msg = "Existing APDR found for key test_partition_key dag_id
testing_dag"
+ lost_race_msg = (
+ "Another process created the APDR for key test_partition_key
dag_id testing_dag first; reusing it"
+ )
+ assert msg_counts[create_msg] == 1
+ assert msg_counts[found_msg] + msg_counts[lost_race_msg] ==
thread_count - 1
assert len(set(ids)) == 1
assert
session.scalar(select(func.count()).select_from(AssetPartitionDagRun)) == 1
+ @pytest.mark.usefixtures("testing_dag_bundle")
+ def
test_get_or_create_apdr_reuses_winner_after_losing_insert_and_keeps_unflushed_rows(self,
session):
+ """An APDR INSERT that loses to the unique constraint reuses the row
that won.
+
+ Rolling back the savepoint must not expunge rows the caller added
before and has not
+ flushed, such as the ``Log`` rows ``_queue_partitioned_dags`` adds
between APDRs.
+ """
+ testing_dag = DagModel(dag_id="testing_dag_lost_race", is_stale=False,
bundle_name="testing")
+ session.add(testing_dag)
+ session.commit()
+ log_row = Log(event="added before the savepoint")
+ session.add(log_row)
+ real_lookup = AssetManager._get_latest_pending_apdr
+
+ def lookup_misses_concurrent_insert(*, target_key, target_dag_id,
session):
+ # Another writer inserts the row right after this lookup, so the
INSERT that
+ # follows collides with it on the unique constraint.
+ lookup.side_effect = real_lookup
+ session.execute(
+ insert(AssetPartitionDagRun).values(
+ target_dag_id=target_dag_id, partition_key=target_key,
pending_partition_key=target_key
+ )
+ )
+ return None
+
+ with mock.patch.object(
+ AssetManager,
+ "_get_latest_pending_apdr",
+ autospec=True,
+ side_effect=lookup_misses_concurrent_insert,
+ ) as lookup:
+ apdr = AssetManager._get_or_create_apdr(
+ target_key="2026-05-20",
+ target_partition_date=None,
+ target_dag=testing_dag,
+ rollup_fingerprint={},
+ allow_reuse=True,
+ session=session,
+ )
+
+ assert lookup.call_count == 2
+ assert apdr.pending_partition_key == "2026-05-20"
+ assert session.scalars(
+
select(AssetPartitionDagRun.id).where(AssetPartitionDagRun.target_dag_id ==
testing_dag.dag_id)
+ ).all() == [apdr.id]
+ session.flush()
+ assert session.scalar(select(Log.event).where(Log.id == log_row.id))
== "added before the savepoint"
+
@pytest.mark.usefixtures("testing_dag_bundle")
def test_get_or_create_apdr_suppresses_conflicting_partition_date(self,
session):
"""Two events resolving the same target key to different dates →
suppress to None.
@@ -467,9 +522,8 @@ class TestAssetManager:
Rather than an order-dependent first-event-wins, conflicting carried
dates produce a
deterministic ``None`` so the consumer DagRun is not stamped with a
wrong, unstable date.
"""
- asm = AssetModel(uri="test://asset1/", name="partition_asset",
group="asset")
testing_dag = DagModel(dag_id="testing_dag_pd_conflict",
is_stale=False, bundle_name="testing")
- session.add_all([asm, testing_dag])
+ session.add(testing_dag)
session.commit()
fp = {"asset-1|test://asset1/": {"__type": "IdentityMapper", "__var":
{}}}
@@ -478,7 +532,6 @@ class TestAssetManager:
target_partition_date=timezone.parse("2026-05-20T00:00:00"),
target_dag=testing_dag,
rollup_fingerprint=fp,
- asset_id=asm.id,
allow_reuse=True,
session=session,
)
@@ -490,7 +543,6 @@ class TestAssetManager:
target_partition_date=timezone.parse("2026-05-21T00:00:00"),
target_dag=testing_dag,
rollup_fingerprint=fp,
- asset_id=asm.id,
allow_reuse=True,
session=session,
)
@@ -500,9 +552,8 @@ class TestAssetManager:
@pytest.mark.usefixtures("testing_dag_bundle")
def test_get_or_create_apdr_keeps_agreeing_partition_date(self, session):
"""A later event carrying the same (or no) date does not trip the
conflict suppression."""
- asm = AssetModel(uri="test://asset1/", name="partition_asset",
group="asset")
testing_dag = DagModel(dag_id="testing_dag_pd_agree", is_stale=False,
bundle_name="testing")
- session.add_all([asm, testing_dag])
+ session.add(testing_dag)
session.commit()
fp = {"asset-1|test://asset1/": {"__type": "IdentityMapper", "__var":
{}}}
source_date = timezone.parse("2026-05-20T00:00:00")
@@ -511,7 +562,6 @@ class TestAssetManager:
target_key="2026-05-20",
target_dag=testing_dag,
rollup_fingerprint=fp,
- asset_id=asm.id,
allow_reuse=True,
session=session,
)
@@ -529,9 +579,8 @@ class TestAssetManager:
@pytest.mark.usefixtures("testing_dag_bundle")
def test_get_or_create_apdr_adopts_date_when_existing_is_none(self,
session):
"""An APDR created with no date adopts a later event's carried date
(not dropped)."""
- asm = AssetModel(uri="test://asset1/", name="partition_asset",
group="asset")
testing_dag = DagModel(dag_id="testing_dag_pd_adopt", is_stale=False,
bundle_name="testing")
- session.add_all([asm, testing_dag])
+ session.add(testing_dag)
session.commit()
fp = {"asset-1|test://asset1/": {"__type": "IdentityMapper", "__var":
{}}}
source_date = timezone.parse("2026-05-20T00:00:00")
@@ -540,7 +589,6 @@ class TestAssetManager:
target_key="2026-05-20",
target_dag=testing_dag,
rollup_fingerprint=fp,
- asset_id=asm.id,
allow_reuse=True,
session=session,
)
@@ -555,9 +603,8 @@ class TestAssetManager:
@pytest.mark.usefixtures("testing_dag_bundle")
def test_get_or_create_apdr_recovers_after_conflict(self, session):
"""Once a conflict has suppressed the date to None, a later event
re-adopts a date."""
- asm = AssetModel(uri="test://asset1/", name="partition_asset",
group="asset")
testing_dag = DagModel(dag_id="testing_dag_pd_recover",
is_stale=False, bundle_name="testing")
- session.add_all([asm, testing_dag])
+ session.add(testing_dag)
session.commit()
fp = {"asset-1|test://asset1/": {"__type": "IdentityMapper", "__var":
{}}}
date_1 = timezone.parse("2026-05-20T00:00:00")
@@ -567,7 +614,6 @@ class TestAssetManager:
target_key="2026-05-20",
target_dag=testing_dag,
rollup_fingerprint=fp,
- asset_id=asm.id,
allow_reuse=True,
session=session,
)
@@ -628,9 +674,8 @@ class TestAssetManager:
def test_get_or_create_apdr_allow_reuse(self, session, allow_reuse):
clear_db_apdr()
clear_db_pakl()
- asset = AssetModel(uri="test://reuse/", name="reuse_asset",
group="asset")
dag = DagModel(dag_id="reuse_test_dag", is_stale=False,
bundle_name="testing")
- session.add_all([asset, dag])
+ session.add(dag)
session.flush()
with capture_orm_selects("asset_partition_dag_run") as lookups:
@@ -640,7 +685,6 @@ class TestAssetManager:
target_partition_date=None,
target_dag=dag,
rollup_fingerprint={},
- asset_id=asset.id,
allow_reuse=allow_reuse,
session=session,
)
@@ -652,6 +696,197 @@ class TestAssetManager:
assert len({apdr.id for apdr in apdrs}) == expected_count
assert
session.scalar(select(func.count()).select_from(AssetPartitionDagRun)) ==
expected_count
assert all(apdr.created_dag_run_id is None for apdr in apdrs)
+ # Only a reusable row takes part in the one-pending-row-per-key
constraint.
+ assert {apdr.pending_partition_key for apdr in apdrs} == {"key-1" if
allow_reuse else None}
+
+ @pytest.mark.backend("postgres", "mysql")
+ @pytest.mark.usefixtures("testing_dag_bundle")
+ def
test_concurrent_events_from_different_assets_share_one_pending_apdr(self,
dag_maker, session):
+ """Events from two producer assets racing on one partition key
accumulate on one APDR.
+
+ If the second writer looks the key up before the first writer commits
its APDR, it finds
+ nothing and inserts its own. Two pending APDRs holding one event each
can
+ never satisfy ``asset_a & asset_b``, and the consumer Dag never runs.
+ """
+ _clear_partition_db()
+ try:
+ asset_a = Asset(name="race-asset-a")
+ asset_b = Asset(name="race-asset-b")
+ with dag_maker(
+ dag_id="race-consumer",
+ schedule=PartitionedAssetTimetable(
+ assets=asset_a & asset_b,
default_partition_mapper=IdentityMapper()
+ ),
+ serialized=True,
+ ):
+ EmptyOperator(task_id="hi")
+ dag_maker.sync_dagbag_to_db()
+ session.commit()
+
+ test_thread = threading.get_ident()
+ second_writer_inserting_apdr = threading.Event()
+
+ def on_execute(conn, cursor, statement, parameters, context,
executemany):
+ if threading.get_ident() != test_thread and
statement.startswith(
+ "INSERT INTO asset_partition_dag_run "
+ ):
+ second_writer_inserting_apdr.set()
+
+ def run_second_writer():
+ with create_session(scoped=False) as second:
+ AssetManager.register_asset_change(asset=asset_b,
partition_key="day-1", session=second)
+
+ event.listen(settings.engine, "before_cursor_execute", on_execute)
+ try:
+ with create_session(scoped=False) as first:
+ AssetManager.register_asset_change(asset=asset_a,
partition_key="day-1", session=first)
+ first.flush()
+ with concurrent.futures.ThreadPoolExecutor(max_workers=1)
as pool:
+ second_done = pool.submit(run_second_writer)
+ assert second_writer_inserting_apdr.wait(timeout=30)
+ first.commit()
+ second_done.result(timeout=30)
+ finally:
+ event.remove(settings.engine, "before_cursor_execute",
on_execute)
+
+ pending = session.scalars(
+
select(AssetPartitionDagRun).where(AssetPartitionDagRun.target_dag_id ==
"race-consumer")
+ ).all()
+ assert len(pending) == 1
+ logged_events = session.scalar(
+ select(func.count())
+ .select_from(PartitionedAssetKeyLog)
+ .where(PartitionedAssetKeyLog.asset_partition_dag_run_id ==
pending[0].id)
+ )
+ assert logged_events == 2
+ finally:
+ _clear_partition_db()
+
+ @pytest.mark.usefixtures("testing_dag_bundle")
+ def test_get_or_create_apdr_inserts_again_when_winner_was_consumed(self,
session, dag_maker):
+ """A winner turned into a Dag run before the loser looks it up has
freed the key, so the loser inserts anew."""
+ with dag_maker(dag_id="consumed_winner_producer", serialized=True):
+ EmptyOperator(task_id="hi")
+ dag_run = dag_maker.create_dagrun()
+ testing_dag = DagModel(dag_id="testing_dag_consumed_winner",
is_stale=False, bundle_name="testing")
+ session.add(testing_dag)
+ winner = AssetPartitionDagRun(
+ target_dag_id=testing_dag.dag_id, partition_key="k",
pending_partition_key="k"
+ )
+ session.add(winner)
+ session.commit()
+ lookups = []
+
+ def lookup_misses_then_consumes_winner(*, target_key, target_dag_id,
session):
+ if lookups:
+ winner.created_dag_run_id = dag_run.id
+ session.flush()
+ lookups.append(target_key)
+ return None
+
+ with mock.patch.object(
+ AssetManager,
+ "_get_latest_pending_apdr",
+ autospec=True,
+ side_effect=lookup_misses_then_consumes_winner,
+ ):
+ apdr = AssetManager._get_or_create_apdr(
+ target_key="k",
+ target_partition_date=None,
+ target_dag=testing_dag,
+ rollup_fingerprint={},
+ allow_reuse=True,
+ session=session,
+ )
+ session.flush()
+
+ assert apdr.id != winner.id
+ assert apdr.pending_partition_key == "k"
+ assert dict(
+ session.execute(
+ select(AssetPartitionDagRun.id,
AssetPartitionDagRun.pending_partition_key).where(
+ AssetPartitionDagRun.target_dag_id == testing_dag.dag_id
+ )
+ ).all()
+ ) == {winner.id: None, apdr.id: "k"}
+
+ @pytest.mark.usefixtures("testing_dag_bundle")
+ def test_get_or_create_apdr_reraises_when_insert_keeps_conflicting(self,
session):
+ testing_dag = DagModel(dag_id="testing_dag_no_winner", is_stale=False,
bundle_name="testing")
+ session.add(testing_dag)
+ session.commit()
+ session.add(
+ AssetPartitionDagRun(
+ target_dag_id=testing_dag.dag_id, partition_key="k",
pending_partition_key="k"
+ )
+ )
+ session.flush()
+
+ with mock.patch.object(AssetManager, "_get_latest_pending_apdr",
autospec=True, return_value=None):
+ with pytest.raises(exc.IntegrityError):
+ AssetManager._get_or_create_apdr(
+ target_key="k",
+ target_partition_date=None,
+ target_dag=testing_dag,
+ rollup_fingerprint={},
+ allow_reuse=True,
+ session=session,
+ )
+
+ assert (
+ session.scalar(
+ select(func.count())
+ .select_from(AssetPartitionDagRun)
+ .where(AssetPartitionDagRun.target_dag_id ==
testing_dag.dag_id)
+ )
+ == 1
+ )
+
+ @pytest.mark.usefixtures("testing_dag_bundle")
+ def test_queue_dagruns_creates_apdrs_in_dag_id_order(self, session,
dag_maker):
+ """Partitioned consumers get their APDRs in dag_id order whatever
order the set iterates in."""
+ _clear_partition_db()
+ try:
+ asset = Asset(name="ordered-asset")
+ for dag_id in ("z-consumer", "a-consumer"):
+ with dag_maker(
+ dag_id=dag_id,
+ schedule=PartitionedAssetTimetable(
+ assets=asset, default_partition_mapper=IdentityMapper()
+ ),
+ serialized=True,
+ ):
+ EmptyOperator(task_id="hi")
+ dag_maker.sync_dagbag_to_db()
+ session.commit()
+
+ class ReverseOrderSet(set):
+ def __iter__(self):
+ return iter(sorted(super().__iter__(), key=lambda d:
d.dag_id, reverse=True))
+
+ real_queue_dagruns = AssetManager._queue_dagruns
+
+ def queue_dagruns_in_reverse_order(*, dags_to_queue, **kwargs):
+
real_queue_dagruns(dags_to_queue=ReverseOrderSet(dags_to_queue), **kwargs)
+
+ with mock.patch.object(
+ AssetManager, "_queue_dagruns", autospec=True,
side_effect=queue_dagruns_in_reverse_order
+ ):
+ AssetManager.register_asset_change(asset=asset,
partition_key="day-1", session=session)
+ session.flush()
+
+ assert session.scalars(
+
select(AssetPartitionDagRun.target_dag_id).order_by(AssetPartitionDagRun.id)
+ ).all() == ["a-consumer", "z-consumer"]
+ finally:
+ _clear_partition_db()
+
+ def test_assigning_created_dag_run_id_clears_pending_partition_key(self):
+ apdr = AssetPartitionDagRun(target_dag_id="d", partition_key="k",
pending_partition_key="k")
+
+ apdr.created_dag_run_id = 123
+
+ assert apdr.pending_partition_key is None
@pytest.mark.need_serialized_dag
@pytest.mark.usefixtures("testing_dag_bundle")
diff --git
a/airflow-core/tests/unit/migrations/test_0148_add_pending_partition_key_to_apdr.py
b/airflow-core/tests/unit/migrations/test_0148_add_pending_partition_key_to_apdr.py
new file mode 100644
index 00000000000..064fee68e10
--- /dev/null
+++
b/airflow-core/tests/unit/migrations/test_0148_add_pending_partition_key_to_apdr.py
@@ -0,0 +1,187 @@
+# 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.
+"""
+Tests for the pending partition key migration, written against the schema as
it stood at that revision.
+
+The tables are declared here instead of taken from ``airflow.models``, and
alembic moves the database
+to the revision, so later migrations and model changes cannot change what
these tests exercise.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+
+import pytest
+import sqlalchemy as sa
+from sqlalchemy import inspect
+
+from airflow import settings
+from airflow._shared.timezones import timezone
+from airflow.utils.db import downgrade, upgradedb
+from airflow.utils.sqlalchemy import UtcDateTime
+
+pytestmark = pytest.mark.db_test
+
+_migration =
import_module("airflow.migrations.versions.0148_3_4_0_add_pending_partition_key_to_apdr")
+REVISION = _migration.revision
+PREVIOUS_REVISION = _migration.down_revision
+DAG_ID = "migration_0148"
+OTHER_DAG_ID = "migration_0148_other"
+PARTITION_DATE = timezone.datetime(2026, 5, 20)
+
+metadata = sa.MetaData()
+dag_run = sa.Table(
+ "dag_run",
+ metadata,
+ sa.Column("id", sa.Integer, primary_key=True),
+ sa.Column("dag_id", sa.String(250)),
+ sa.Column("run_id", sa.String(250)),
+ sa.Column("run_type", sa.String(50)),
+ sa.Column("run_after", UtcDateTime),
+)
+apdr = sa.Table(
+ "asset_partition_dag_run",
+ metadata,
+ sa.Column("id", sa.Integer, primary_key=True),
+ sa.Column("target_dag_id", sa.String(250)),
+ sa.Column("created_dag_run_id", sa.Integer),
+ sa.Column("partition_key", sa.String(250)),
+ sa.Column("partition_date", UtcDateTime),
+ sa.Column("created_at", UtcDateTime),
+ sa.Column("updated_at", UtcDateTime),
+)
+key_log = sa.Table(
+ "partitioned_asset_key_log",
+ metadata,
+ sa.Column("id", sa.Integer, primary_key=True),
+ sa.Column("asset_id", sa.Integer),
+ sa.Column("asset_event_id", sa.Integer),
+ sa.Column("asset_partition_dag_run_id", sa.Integer),
+ sa.Column("source_partition_key", sa.String(250)),
+ sa.Column("target_dag_id", sa.String(250)),
+ sa.Column("target_partition_key", sa.String(250)),
+ sa.Column("created_at", UtcDateTime),
+)
+
+
+def _delete_rows(connection):
+ dag_ids = (DAG_ID, OTHER_DAG_ID)
+
connection.execute(key_log.delete().where(key_log.c.target_dag_id.in_(dag_ids)))
+ connection.execute(apdr.delete().where(apdr.c.target_dag_id.in_(dag_ids)))
+ connection.execute(dag_run.delete().where(dag_run.c.dag_id == DAG_ID))
+
+
[email protected]
+def database_before_revision():
+ """Move the database to the revision before the migration for one test,
then back to the current head."""
+ original_async_conn = settings.SQL_ALCHEMY_CONN_ASYNC
+ settings.SQL_ALCHEMY_CONN_ASYNC = ""
+ settings.dispose_orm(do_log=False)
+ settings.configure_orm(disable_connection_pool=True)
+ try:
+ downgrade(to_revision=PREVIOUS_REVISION)
+ with settings.engine.begin() as connection:
+ _delete_rows(connection)
+ yield settings.engine
+ finally:
+ try:
+ with settings.engine.begin() as connection:
+ _delete_rows(connection)
+ upgradedb()
+ finally:
+ settings.SQL_ALCHEMY_CONN_ASYNC = original_async_conn
+ settings.reconfigure_orm()
+
+
+def
test_upgrade_marks_only_the_latest_pending_row_per_key_and_keeps_every_row(database_before_revision):
+ now = timezone.utcnow()
+ with database_before_revision.begin() as connection:
+ dag_run_id = connection.execute(
+ dag_run.insert().values(dag_id=DAG_ID, run_id="run",
run_type="asset_triggered", run_after=now)
+ ).inserted_primary_key[0]
+
+ def add_apdr(partition_key, *, dag_id=DAG_ID, created_dag_run_id=None,
partition_date=None):
+ return connection.execute(
+ apdr.insert().values(
+ target_dag_id=dag_id,
+ partition_key=partition_key,
+ created_dag_run_id=created_dag_run_id,
+ partition_date=partition_date,
+ created_at=now,
+ updated_at=now,
+ )
+ ).inserted_primary_key[0]
+
+ completed = add_apdr("day-1", created_dag_run_id=dag_run_id)
+ older_pending = add_apdr("day-1", partition_date=PARTITION_DATE)
+ other_dag = add_apdr("day-1", dag_id=OTHER_DAG_ID)
+ latest_pending = add_apdr("day-1")
+ other_key = add_apdr("day-2")
+ apdr_ids = (completed, older_pending, other_dag, latest_pending,
other_key)
+ for asset_event_id, apdr_id in enumerate(apdr_ids):
+ connection.execute(
+ key_log.insert().values(
+ asset_id=asset_event_id,
+ asset_event_id=asset_event_id,
+ asset_partition_dag_run_id=apdr_id,
+ source_partition_key="source",
+ target_dag_id=OTHER_DAG_ID if apdr_id == other_dag else
DAG_ID,
+ target_partition_key="target",
+ created_at=now,
+ )
+ )
+
+ upgradedb(to_revision=REVISION)
+
+ upgraded_apdr = sa.Table(apdr.name, sa.MetaData(),
autoload_with=database_before_revision)
+ with database_before_revision.connect() as connection:
+ apdr_rows = {
+ row.id: (row.partition_key, row.created_dag_run_id,
row.pending_partition_key)
+ for row in connection.execute(
+
sa.select(upgraded_apdr).where(upgraded_apdr.c.target_dag_id.in_((DAG_ID,
OTHER_DAG_ID)))
+ )
+ }
+ partition_dates = dict(
+ connection.execute(
+ sa.select(apdr.c.id, apdr.c.partition_date).where(
+ apdr.c.target_dag_id.in_((DAG_ID, OTHER_DAG_ID))
+ )
+ ).all()
+ )
+ log_rows = dict(
+ connection.execute(
+ sa.select(key_log.c.asset_event_id,
key_log.c.asset_partition_dag_run_id).where(
+ key_log.c.target_dag_id.in_((DAG_ID, OTHER_DAG_ID))
+ )
+ ).all()
+ )
+
+ # A Dag without batch_asset_events has one pending row per event, so none
are merged or deleted:
+ # only the latest pending row of each key is marked, and a completed row
never is.
+ assert apdr_rows == {
+ completed: ("day-1", dag_run_id, None),
+ older_pending: ("day-1", None, None),
+ other_dag: ("day-1", None, "day-1"),
+ latest_pending: ("day-1", None, "day-1"),
+ other_key: ("day-2", None, "day-2"),
+ }
+ assert partition_dates[older_pending] == PARTITION_DATE
+ assert log_rows == dict(enumerate(apdr_ids))
+ assert _migration._UQ_NAME in {
+ constraint["name"]
+ for constraint in
inspect(database_before_revision).get_unique_constraints(apdr.name)
+ }
