kaxil commented on code in PR #66016:
URL: https://github.com/apache/airflow/pull/66016#discussion_r3164800574
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_modify_deadline_callback_schema.py:
##########
@@ -38,17 +41,113 @@
airflow_version = "3.1.0"
+_ASYNC_CALLBACK_CLASSNAME = "airflow.sdk.definitions.deadline.AsyncCallback"
+
+
def upgrade():
"""Replace deadline table's string callback and JSON callback_kwargs with
JSON callback."""
+ if context.is_offline_mode():
+ # In offline (--sql) mode there is no live connection so data cannot be
+ # migrated. Delete existing rows so the NOT NULL column can be added
+ # safely. This matches the offline behaviour of migration 0094.
+ print(
+ dedent("""
+ ------------
+ -- WARNING: Unable to migrate the data in the deadline table
+ -- while in offline mode! All rows in the deadline table will
+ -- be deleted.
+ ------------
+ """)
+ )
+ op.execute("DELETE FROM deadline")
+ with op.batch_alter_table("deadline", schema=None) as batch_op:
+ batch_op.drop_column("callback")
+ batch_op.drop_column("callback_kwargs")
+ batch_op.add_column(sa.Column("callback", sa.JSON(),
nullable=False))
+ return
+
+ conn = op.get_bind()
+
+ rows = conn.execute(sa.text("SELECT id, callback, callback_kwargs FROM
deadline")).fetchall()
+ callback_data: dict = {}
+ for row in rows:
+ path = row[1] or ""
+ kwargs = row[2]
+ if isinstance(kwargs, str):
+ kwargs = json.loads(kwargs) if kwargs else {}
+ elif kwargs is None:
+ kwargs = {}
+ callback_data[row[0]] = {
+ "__data__": {"path": path, "kwargs": kwargs},
+ "__classname__": _ASYNC_CALLBACK_CLASSNAME,
+ "__version__": 0,
+ }
+
with op.batch_alter_table("deadline", schema=None) as batch_op:
batch_op.drop_column("callback")
batch_op.drop_column("callback_kwargs")
- batch_op.add_column(sa.Column("callback", sa.JSON(), nullable=False))
+ batch_op.add_column(sa.Column("callback", sa.JSON(), nullable=True))
+
+ deadline_table = sa.table("deadline", sa.column("id"),
sa.column("callback", sa.JSON()))
+ for row_id, serialized in callback_data.items():
+ conn.execute(
+ sa.update(deadline_table).where(deadline_table.c.id ==
row_id).values(callback=serialized)
Review Comment:
One UPDATE round-trip per row. Could you batch with `executemany`? The
pattern at `0094_3_2_0_replace_deadline_inline_callback_with_fkey.py:232-237`
uses bulk update with bindparams:
```python
conn.execute(
deadline_table.update()
.where(deadline_table.c.id == sa.bindparam("deadline_id"))
.values(callback=sa.bindparam("callback")),
[{"deadline_id": rid, "callback": ser} for rid, ser in
callback_data.items()],
)
```
On 10k rows that's one statement instead of 10k.
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_modify_deadline_callback_schema.py:
##########
@@ -38,17 +41,113 @@
airflow_version = "3.1.0"
+_ASYNC_CALLBACK_CLASSNAME = "airflow.sdk.definitions.deadline.AsyncCallback"
+
+
def upgrade():
"""Replace deadline table's string callback and JSON callback_kwargs with
JSON callback."""
+ if context.is_offline_mode():
+ # In offline (--sql) mode there is no live connection so data cannot be
+ # migrated. Delete existing rows so the NOT NULL column can be added
+ # safely. This matches the offline behaviour of migration 0094.
+ print(
+ dedent("""
+ ------------
+ -- WARNING: Unable to migrate the data in the deadline table
+ -- while in offline mode! All rows in the deadline table will
+ -- be deleted.
+ ------------
+ """)
+ )
+ op.execute("DELETE FROM deadline")
+ with op.batch_alter_table("deadline", schema=None) as batch_op:
+ batch_op.drop_column("callback")
+ batch_op.drop_column("callback_kwargs")
+ batch_op.add_column(sa.Column("callback", sa.JSON(),
nullable=False))
+ return
+
+ conn = op.get_bind()
+
+ rows = conn.execute(sa.text("SELECT id, callback, callback_kwargs FROM
deadline")).fetchall()
Review Comment:
`fetchall()` here loads the entire `deadline` table into memory in one shot.
The neighbouring migration 0094 on the same table uses `conf.getint("database",
"migration_batch_size")` for exactly this case -- worth mirroring the
keyset-paginated pattern from
`0094_3_2_0_replace_deadline_inline_callback_with_fkey.py:_upgrade_mysql_sqlite`.
Normally `deadline` is small, but if the triggerer has been off long enough to
accumulate thousands of unprocessed rows this OOMs the alembic process
mid-upgrade.
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_modify_deadline_callback_schema.py:
##########
@@ -38,17 +41,113 @@
airflow_version = "3.1.0"
+_ASYNC_CALLBACK_CLASSNAME = "airflow.sdk.definitions.deadline.AsyncCallback"
+
+
def upgrade():
"""Replace deadline table's string callback and JSON callback_kwargs with
JSON callback."""
+ if context.is_offline_mode():
+ # In offline (--sql) mode there is no live connection so data cannot be
+ # migrated. Delete existing rows so the NOT NULL column can be added
+ # safely. This matches the offline behaviour of migration 0094.
+ print(
+ dedent("""
+ ------------
+ -- WARNING: Unable to migrate the data in the deadline table
+ -- while in offline mode! All rows in the deadline table will
+ -- be deleted.
+ ------------
+ """)
+ )
+ op.execute("DELETE FROM deadline")
+ with op.batch_alter_table("deadline", schema=None) as batch_op:
+ batch_op.drop_column("callback")
+ batch_op.drop_column("callback_kwargs")
+ batch_op.add_column(sa.Column("callback", sa.JSON(),
nullable=False))
+ return
+
+ conn = op.get_bind()
+
+ rows = conn.execute(sa.text("SELECT id, callback, callback_kwargs FROM
deadline")).fetchall()
+ callback_data: dict = {}
+ for row in rows:
+ path = row[1] or ""
+ kwargs = row[2]
+ if isinstance(kwargs, str):
+ kwargs = json.loads(kwargs) if kwargs else {}
+ elif kwargs is None:
+ kwargs = {}
+ callback_data[row[0]] = {
+ "__data__": {"path": path, "kwargs": kwargs},
+ "__classname__": _ASYNC_CALLBACK_CLASSNAME,
+ "__version__": 0,
+ }
+
with op.batch_alter_table("deadline", schema=None) as batch_op:
batch_op.drop_column("callback")
batch_op.drop_column("callback_kwargs")
- batch_op.add_column(sa.Column("callback", sa.JSON(), nullable=False))
+ batch_op.add_column(sa.Column("callback", sa.JSON(), nullable=True))
+
+ deadline_table = sa.table("deadline", sa.column("id"),
sa.column("callback", sa.JSON()))
+ for row_id, serialized in callback_data.items():
+ conn.execute(
+ sa.update(deadline_table).where(deadline_table.c.id ==
row_id).values(callback=serialized)
+ )
+
+ with op.batch_alter_table("deadline", schema=None) as batch_op:
+ batch_op.alter_column("callback", existing_type=sa.JSON(),
nullable=False)
def downgrade():
"""Replace deadline table's JSON callback with string callback and JSON
callback_kwargs."""
+ if context.is_offline_mode():
+ print(
+ dedent("""
+ ------------
+ -- WARNING: Unable to migrate the data in the deadline table
+ -- while in offline mode! All rows in the deadline table will
+ -- be deleted.
+ ------------
+ """)
+ )
+ op.execute("DELETE FROM deadline")
+ with op.batch_alter_table("deadline", schema=None) as batch_op:
+ batch_op.drop_column("callback")
+ batch_op.add_column(sa.Column("callback_kwargs", sa.JSON(),
nullable=True))
+ batch_op.add_column(sa.Column("callback", sa.String(length=500),
nullable=False))
+ return
+
+ conn = op.get_bind()
+
+ rows = conn.execute(sa.text("SELECT id, callback FROM
deadline")).fetchall()
+ callback_data: dict = {}
+ for row in rows:
+ cb = row[1]
+ if cb is None:
+ callback_data[row[0]] = ("", {})
+ continue
+ if isinstance(cb, str):
+ cb = json.loads(cb)
+ cb_inner = cb.get("__data__", cb)
+ callback_data[row[0]] = (cb_inner.get("path", ""),
cb_inner.get("kwargs", {}))
+
with op.batch_alter_table("deadline", schema=None) as batch_op:
batch_op.drop_column("callback")
batch_op.add_column(sa.Column("callback_kwargs", sa.JSON(),
nullable=True))
- batch_op.add_column(sa.Column("callback", sa.String(length=500),
nullable=False))
+ batch_op.add_column(sa.Column("callback", sa.String(length=500),
nullable=True))
Review Comment:
The downgrade restores `callback` to `String(length=500)` and back-fills
from the JSON envelope. If anyone's `__data__.path` is over 500 chars (legal in
the JSON column), the post-backfill UPDATE on PostgreSQL or MySQL strict-mode
raises `value too long for type character varying(500)` partway through. Edge
case but worth either truncating with a warning or documenting in the docstring
as a known limitation.
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_modify_deadline_callback_schema.py:
##########
@@ -38,17 +41,113 @@
airflow_version = "3.1.0"
+_ASYNC_CALLBACK_CLASSNAME = "airflow.sdk.definitions.deadline.AsyncCallback"
+
+
def upgrade():
Review Comment:
No regression test for this fix. Per the project's testing standards, every
bug fix needs a test. Could you add one that stamps `3bda03debd04`, inserts a
deadline row with old-shape `(callback VARCHAR, callback_kwargs JSON)`, runs
`alembic upgrade 808787349f22`, and verifies the row's `callback` is the
`{"__data__": {...}, "__classname__": ..., "__version__": 0}` envelope? Plus
the downgrade round-trip. Otherwise a future refactor can silently re-introduce
the `NotNullViolation`.
##########
airflow-core/src/airflow/migrations/versions/0080_3_1_0_modify_deadline_callback_schema.py:
##########
@@ -38,17 +41,113 @@
airflow_version = "3.1.0"
+_ASYNC_CALLBACK_CLASSNAME = "airflow.sdk.definitions.deadline.AsyncCallback"
+
+
def upgrade():
"""Replace deadline table's string callback and JSON callback_kwargs with
JSON callback."""
+ if context.is_offline_mode():
+ # In offline (--sql) mode there is no live connection so data cannot be
+ # migrated. Delete existing rows so the NOT NULL column can be added
+ # safely. This matches the offline behaviour of migration 0094.
+ print(
+ dedent("""
+ ------------
+ -- WARNING: Unable to migrate the data in the deadline table
+ -- while in offline mode! All rows in the deadline table will
+ -- be deleted.
+ ------------
+ """)
+ )
+ op.execute("DELETE FROM deadline")
+ with op.batch_alter_table("deadline", schema=None) as batch_op:
+ batch_op.drop_column("callback")
+ batch_op.drop_column("callback_kwargs")
+ batch_op.add_column(sa.Column("callback", sa.JSON(),
nullable=False))
+ return
+
+ conn = op.get_bind()
+
+ rows = conn.execute(sa.text("SELECT id, callback, callback_kwargs FROM
deadline")).fetchall()
+ callback_data: dict = {}
+ for row in rows:
+ path = row[1] or ""
+ kwargs = row[2]
Review Comment:
If `kwargs` comes back as a list/number/bool from a hand-edited DB it passes
through unchanged into `__data__.kwargs`, where downstream Pydantic expects a
dict. Probably fine in practice (app code only writes dicts), but `if not
isinstance(kwargs, dict): kwargs = {}` would defend against hand-edits cheaply.
--
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]
