Skip to content

Reinstate tih.id as pkey and auto-increment on downgrade #58149

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Chais wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Reinstate tih.id as pkey and auto-increment on downgrade #58149

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0071913
Reinstate tih.id as pkey and auto-increment on downgrade
Nov 10, 2025
6cb636f
Merge branch 'main' into 7645189f3479-downgrade-fix
Chais Nov 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 1 addition & 27 deletions airflow-core/src/airflow/migrations/versions/0060_3_0_0_add_try_id_to_ti_and_tih.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,34 +127,8 @@ def upgrade():

def downgrade():
"""Unapply Add try_id to TI and TIH."""
dialect_name = op.get_bind().dialect.name
with op.batch_alter_table("task_instance_history", schema=None) as batch_op:
batch_op.drop_constraint(batch_op.f("task_instance_history_pkey"), type_="primary")
batch_op.add_column(sa.Column("id", sa.INTEGER, nullable=True))
batch_op.add_column(sa.Column("id", sa.INTEGER, primary_key=True, autoincrement=True))
Copy link
Contributor

@ephraimbuddy ephraimbuddy Dec 2, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This won't work with data as id would be null.
That's why we had to run the deleted statements to update the ID column before making it a primary key

Copy link
Author

@Chais Chais Jan 15, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not convinced this is the case. The downgrade SQL generated for psql looks like this:

BEGIN;
-- Running downgrade 7645189f3479 -> e00344393f31
ALTER TABLE task_instance_history DROP CONSTRAINT task_instance_history_pkey;
ALTER TABLE task_instance_history ADD COLUMN id SERIAL NOT NULL;
ALTER TABLE task_instance_history DROP COLUMN task_instance_id;
ALTER TABLE task_instance_history DROP COLUMN try_id;
UPDATE alembic_version SET version_num='e00344393f31' WHERE alembic_version.version_num = '7645189f3479';
COMMIT;

The relevant line is:

ALTER TABLE task_instance_history ADD COLUMN id SERIAL NOT NULL;

If I now create a test table and fill it with some data:

CREATE TABLE test ( column_a text, column_b integer );
INSERT INTO test (column_a, column_b) VALUES ('some text', 3);
INSERT INTO test (column_a, column_b) VALUES ('some other text', 17);

I see it as expected:

SELECT * FROM test;
    column_a     | column_b 
-----------------+----------
 some text       |        3
 some other text |       17
(2 rows)

If I now apply the relevant change to this table with its existing data like this:

ALTER TABLE test ADD COLUMN id SERIAL NOT NULL;

I can see the values being inserted:

SELECT * FROM test;
    column_a     | column_b | id 
-----------------+----------+----
 some text       |        3 |  1
 some other text |       17 |  2
(2 rows)

So this seems to work as intended on PostgreSQL. I tested it on MySQL when I wrote this PR, with a similar result.

Copy link
Contributor

@ephraimbuddy ephraimbuddy Jan 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you test with postgresql because according to alembic, autoincrement is only understood by mysql: https://alembic.sqlalchemy.org/en/latest/ops.html

Copy link
Author

@Chais Chais Jan 16, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I had hoped this would make it sufficiently clear that I did:

The downgrade SQL generated for psql looks like this:
[…]
So this seems to work as intended on PostgreSQL.

Technically, the generated code doesn't use AUTO_INCREMENT:

ALTER TABLE test ADD COLUMN id SERIAL NOT NULL;

But as we can see in the PostgreSQL documentation, the SERIAL type is an autoincrementing integer, which is just what we need: https://www.postgresql.org/docs/current/datatype-numeric.html
Maybe Alembic translates sa.INTEGER, autoincrement=True into SERIAL for Postgres.

Also, it's not like I invented this definition, either. task_instance_history was introduced in revision d482b7261ff9 (number 21, by the new count) with the id column defined like this:

sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),

So clearly, this worked before (and for all supported DBs, I might add) and I see no reason why it shouldn't work again.
And the fact remains, that the current code leaves the DB broken post-downgrade. While it restores the id column and refills it in what may be the most cumbersome way possible, it neglects the autoincrement the codebase expects at the relevant version, leading to queries failing. Please see the issue I linked in the original post for details.

batch_op.drop_column("task_instance_id")
if dialect_name == "postgresql":
op.execute(
"""
ALTER TABLE task_instance_history ADD COLUMN row_num SERIAL;
UPDATE task_instance_history SET id = row_num;
ALTER TABLE task_instance_history DROP COLUMN row_num;
"""
)
elif dialect_name == "mysql":
op.execute(
"""
SET @row_number = 0;
UPDATE task_instance_history
SET id = (@row_number := @row_number + 1)
ORDER BY try_id;
"""
)
else:
op.execute("""
UPDATE task_instance_history
SET id = (SELECT COUNT(*) FROM task_instance_history t2 WHERE t2.id <= task_instance_history.id);
""")
with op.batch_alter_table("task_instance_history", schema=None) as batch_op:
batch_op.alter_column("id", nullable=False, existing_type=sa.INTEGER)
batch_op.drop_column("try_id")
batch_op.create_primary_key("task_instance_history_pkey", ["id"])
Loading