Migrate TaskInstance to UUID v7 primary key

[kaxil]

closes #43161 part of AIP-72.

As part of the ongoing work for AIP-72: Task Execution Interface, we are migrating the task_instance table to use a UUID primary key. This change is being made to simplify task instance identification, especially when communicating between the executor and workers.

Currently, the primary key of task_instance is a composite key consisting of dag_id, task_id, run_id, and map_index as shown below. This migration introduces a UUID v7 column (id) as the new primary key.

airflow/airflow/models/taskinstance.py

Lines 1815 to 1819 in b4269f3

	__tablename__ = "task_instance"
	task_id = Column(StringID(), primary_key=True, nullable=False)
	dag_id = Column(StringID(), primary_key=True, nullable=False)
	run_id = Column(StringID(), primary_key=True, nullable=False)
	map_index = Column(Integer, primary_key=True, nullable=False, server_default=text("-1"))

Why UUID v7?

The UUID v7 format was chosen because of its improved temporal sorting capabilities. For existing records, UUID v7 will be generated using either the queued_dttm, start_date, or the current timestamp.

(From this blog post.)

Changes

1. Migrated Primary Key to UUID v7

   • Replaced the composite primary key (dag_id, task_id, run_id, map_index) with a UUID v7 id field, ensuring temporal sorting and simplified task instance identification.

2. Database-Specific UUID v7 Functions for DB Migrations

   • Added UUID v7 functions for each database, just for DB migration since generating uuid's on Python and then storing it in DB can be slow:
     • PostgreSQL: Uses pgcrypto for generation with fallback.
     • MySQL: Custom deterministic UUID v7 function.
     • SQLite: Utilizes uuid6 Python package.

3. Updated Constraints and Indexes

   • Added UniqueConstraint on (dag_id, task_id, run_id, map_index) for compatibility.
   • Modified foreign key constraints for the new primary key, handling downgrades to restore previous constraints.

4. Model and API Adjustments

   • Updated TaskInstance model to use UUID v7 as the primary key via uuid6 library, that has uuid7 ! 😄 .
   • Adjusted REST API, views, and queries to support UUID-based lookups.
   • Modified tests for compatibility with the new primary key.

Issues identified

After updating the primary key on the TaskInstance model to a UUID7 id field (from the composite primary key of ["dag_id", "task_id", "run_id", "map_index"]), I have now run into issues in testing and session management:

1. session.merge() Compatibility

   The session.merge() function operates strictly on primary keys as mentioned in this SQLAlchemy docs, which means it no longer recognizes the unique constraint on ["dag_id", "task_id", "run_id", "map_index"] to identify existing TaskInstance records. This leads to issues in cases where session.merge() was previously used to either update or insert TaskInstance records, as it now fails to locate the intended record by the unique constraint.

   Example:

   airflow/tests/models/test_dag.py

   Lines 2013 to 2026 in b2ecb6c

   	dagrun_1 = dag.create_dagrun(
   	run_type=DagRunType.BACKFILL_JOB,
   	state=DagRunState.RUNNING,
   	start_date=DEFAULT_DATE,
   	execution_date=DEFAULT_DATE,
   	data_interval=(DEFAULT_DATE, DEFAULT_DATE),
   	**triggered_by_kwargs, # type: ignore
   	)
   	session.merge(dagrun_1)
   	task_instance_1 = TI(t_1, run_id=dagrun_1.run_id, state=ti_state_begin)
   	task_instance_1.job_id = 123
   	session.merge(task_instance_1)
   	session.commit()

   A possible workaround would be to follow session.merge() calls with refresh_from_db() to ensure the instance has the latest state:

   task_instance_1 = session.merge(TaskInstance(task=t_1, run_id=dagrun_1.run_id, state=ti_state_begin))
   task_instance_1.refresh_from_db()  # Update instance with current DB state

2. session.get() usage with TaskInstanceKey

   Similarly, session.get() and other primary key-based session operations now rely solely on the id UUID rather than the composite fields. This results in errors when accessing or sorting records based on the previous composite fields, as session.get() cannot match identifiers to TaskInstance records by the unique constraint.

   Current code in main:

   airflow/tests/models/test_cleartasks.py

   Line 742 in 78ff0a9

   session.get(TaskInstance, ti1.key.primary).try_number += 1

   I have been changing those references from TaskInstanceKey in Database calls with ti.id, example:

   airflow/tests/models/test_cleartasks.py

   Line 742 in cca7ce0

   session.get(TaskInstance, ti1.id).try_number += 1

3. Airflow REST API calls which filters TIs

   Similar to (2) above, some of our REST API call filters TI records based on current primary key as below:
   

   airflow/airflow/api_connexion/endpoints/task_instance_endpoint.py

   Lines 536 to 537 in 78ff0a9

   	if run_id and not session.get(
   	TI, {"task_id": task_id, "dag_id": dag_id, "run_id": run_id, "map_index": -1}

   I am changing that to session.scalars(select(TI).where(...))..one_or_none() queries. For example, the above code is changed to the following in this PR.
   

   airflow/airflow/api_connexion/endpoints/task_instance_endpoint.py

   Lines 536 to 540 in 3c61ba7

   	select_stmt = select(TI).where(
   	TI.dag_id == dag_id, TI.task_id == task_id, TI.run_id == run_id, TI.map_index == -1
   	)
   	if run_id and not session.scalars(select_stmt).one_or_none():

TODOs

• Check if the DB user has permissions to add CREATE extension for Postgres as we install pgcrypto
• Convert the raw SQL for dropping the primary key (CASCADE) into an Alembic operation. This will allow capturing foreign key constraints and managing them in a downgrade scenario. MySQL does not support dropping foreign keys using CASCADE with primary keys anyway.
• Ensure foreign key constraints are re-added during downgrades after the primary key has been changed back to the composite key.
• SQLite Support: Add SQLite support for UUID v7 generation using the uuid6 Python package.
• Add changes to TaskInstance model file to add id

Future Work

• Evaluate if the task_instance_history table requires similar updates for UUID v7 compatibility. (Evaluate if the TI History table requires similar updates (TI pk) for UUID v7 compatibility #43437)
• Batch Updates: Investigate batch processing for updating UUIDs to avoid long locks on large datasets. Maybe run with batches of 10000 rows??? -- TBD. We can also run some benchmarks once we do that. (Investigate batch processing for updating TI UUIDs to avoid long locks on large datasets #43438)
• Determine if including try_number in the UniqueConstraint improves index efficiency. (Determine if including try_number in the UniqueConstraint improves index efficiency #43439)

[ashb]

A neat trick to avoid the need for pgcrypto extension: https://brandur.org/fragments/secure-bytes-without-pgcrypto

Though uuid_send() fn it uses doesn't exist on Pg10, it only comes from Pg11 onwards. However I do think it's reasonable to say 12 or 13 is the min required Pg version for Airflow 3 (Postgres only supports v13 onwards).

Oh, Pg 12 is officially all we support anyway, so I think using the uuid_send() function approach could work fine.

[kaxil]

Issues identified

After updating the primary key on the TaskInstance model to a UUID7 id field (from the composite primary key of ["dag_id", "task_id", "run_id", "map_index"]), I have now run into issues in testing and session management:

1. session.merge() Compatibility

   The session.merge() function operates strictly on primary keys as mentioned in this SQLAlchemy docs, which means it no longer recognizes the unique constraint on ["dag_id", "task_id", "run_id", "map_index"] to identify existing TaskInstance records. This leads to issues in cases where session.merge() was previously used to either update or insert TaskInstance records, as it now fails to locate the intended record by the unique constraint.

   Example:

   airflow/tests/models/test_dag.py

   Lines 2013 to 2026 in b2ecb6c

   	dagrun_1 = dag.create_dagrun(
   	run_type=DagRunType.BACKFILL_JOB,
   	state=DagRunState.RUNNING,
   	start_date=DEFAULT_DATE,
   	execution_date=DEFAULT_DATE,
   	data_interval=(DEFAULT_DATE, DEFAULT_DATE),
   	**triggered_by_kwargs, # type: ignore
   	)
   	session.merge(dagrun_1)
   	task_instance_1 = TI(t_1, run_id=dagrun_1.run_id, state=ti_state_begin)
   	task_instance_1.job_id = 123
   	session.merge(task_instance_1)
   	session.commit()

   A possible workaround would be to follow session.merge() calls with refresh_from_db() to ensure the instance has the latest state:

   task_instance_1 = session.merge(TaskInstance(task=t_1, run_id=dagrun_1.run_id, state=ti_state_begin))
   task_instance_1.refresh_from_db()  # Update instance with current DB state

2. session.get() usage with TaskInstanceKey

   Similarly, session.get() and other primary key-based session operations now rely solely on the id UUID rather than the composite fields. This results in errors when accessing or sorting records based on the previous composite fields, as session.get() cannot match identifiers to TaskInstance records by the unique constraint.

   Current code in main:

   airflow/tests/models/test_cleartasks.py

   Line 742 in 78ff0a9

   session.get(TaskInstance, ti1.key.primary).try_number += 1

   I have been changing those references from TaskInstanceKey in Database calls with ti.id, example:

   airflow/tests/models/test_cleartasks.py

   Line 742 in cca7ce0

   session.get(TaskInstance, ti1.id).try_number += 1

3. Airflow REST API calls which filters TIs

   Similar to (2) above, some of our REST API call filters TI records based on current primary key as below:
   

   airflow/airflow/api_connexion/endpoints/task_instance_endpoint.py

   Lines 536 to 537 in 78ff0a9

   	if run_id and not session.get(
   	TI, {"task_id": task_id, "dag_id": dag_id, "run_id": run_id, "map_index": -1}

   I am changing that to session.scalars(select(TI).where(...))..one_or_none() queries. For example, the above code is changed to the following in this PR.
   

   airflow/airflow/api_connexion/endpoints/task_instance_endpoint.py

   Lines 536 to 540 in 3c61ba7

   	select_stmt = select(TI).where(
   	TI.dag_id == dag_id, TI.task_id == task_id, TI.run_id == run_id, TI.map_index == -1
   	)
   	if run_id and not session.scalars(select_stmt).one_or_none():

[ashb]

I do think long-term/by release of 3.0 TaskInstanceHistory should also have a uuid.

[ashb]

Yeah, our use of session.merge() in Airflow is v. bad :(

[potiuk]

Yeah, our use of session.merge() in Airflow is v. bad :(

Agree. This is one of the reasons I wanted to avoid all those db-schema changes for the initial proposal of multi-team and proposed the team-id prefix, becaue I was afraid we are going to have similar ripple-effects that will result in rewriting big part of our database code. Which seem that it was a pretty justified worry.

[potiuk]

BTW. One of the ways it could be helped - we could potentially generate the unique id by hashing the remaining fields - with appropriate hashing algorithm, we could have very low probability of collision (and maybe we could implement a mechanism to detect collisions and implement handling of those collision in similar ways they are handled in hash maps) - that would be another way how we could approach it (and there the merge code could be simpler and still use session.merge()

But this one also comes with it's own set of difficulties and collision handling is likely going to be complex. But I thought it's worth menioning it here as an option.

[potiuk]

One problem with generating the unique hash with other fields, is that it will be even more difficult to add new fields, so probably getting rid of session.merge() everywhere and replacing it with our own way of merging objects and database entries by retrieving the unique id based on the unique fields from the DB before merging is likely a long-term better approach.

[kaxil]

@potiuk Yeah, I considered hashing but it is bad for Databases for indexing and, as a result, querying since it won't have temporal properties apart from Collision Handling complexity.

UUID v7 is explicitly designed to support distributed databases with high insert rates due to its temporal ordering.

BTW. One of the ways it could be helped - we could potentially generate the unique id by hashing the remaining fields - with appropriate hashing algorithm, we could have very low probability of collision (and maybe we could implement a mechanism to detect collisions and implement handling of those collision in similar ways they are handled in hash maps) - that would be another way how we could approach it (and there the merge code could be simpler and still use session.merge()

But this one also comes with it's own set of difficulties and collision handling is likely going to be complex. But I thought it's worth menioning it here as an option.

[kaxil]

Yeah, our use of session.merge() in Airflow is v. bad :(

Indeed

[potiuk]

UUID v7 is explicitly designed to support distributed databases with high insert rates due to its temporal ordering.

Yep. that settles it as well. We would have to implement our hashing in similar way (if possible at all)