Update conf column in dag_run table type from bytea to JSON #44533

vatsrahul1001 · 2024-12-01T13:11:21Z

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

vatsrahul1001 · 2024-12-02T10:39:07Z

While Testing this I noticed an issue in the downgrade case. We are not removing the data from downgrade and just perform below conversion.

          ALTER TABLE dag_run
          ALTER COLUMN conf TYPE BYTEA
          USING CASE
              WHEN conf IS NOT NULL THEN CONVERT_TO(conf::TEXT, 'UTF8')
              ELSE NULL
          END

I am getting UnpicklingError error with this when I try to open DAG page.

Should we also move data to the archive table in case of a downgrade as well?
cc: @kaxil

kaxil · 2024-12-02T14:08:35Z

While Testing this I noticed an issue in the downgrade case. We are not removing the data from downgrade and just perform below conversion.
          ALTER TABLE dag_run
          ALTER COLUMN conf TYPE BYTEA
          USING CASE
              WHEN conf IS NOT NULL THEN CONVERT_TO(conf::TEXT, 'UTF8')
              ELSE NULL
          END
 
I am getting UnpicklingError error with this when I try to open DAG page.

Should we also move data to the archive table in case of a downgrade as well? cc: @kaxil

Yeah, it is different than XCom, because XCom code had handling of both JSON & pickle type -- for dagrun conf it would be different. For downgrade you might just want to insert all records from the archive table back here.

kaxil · 2024-12-02T14:09:17Z

airflow/models/dagrun.py

@@ -137,7 +138,7 @@ class DagRun(Base, LoggingMixin):
    triggered_by = Column(
        Enum(DagRunTriggeredByType, native_enum=False, length=50)
    )  # Airflow component that triggered the run.
-    conf = Column(PickleType)
+    conf = Column(JSON().with_variant(postgresql.JSONB, "postgresql"))


Do we not need any more handling than these?

Check if the DagRun edit view works in the FAB UI

I think you might need to change this line:

airflow/airflow/www/forms.py

Lines 131 to 132 in 3c1124e

if item.conf:

item.conf = json.loads(item.conf)

You will also need to add a newsfragment, something like:

https://github.com/astronomer/airflow/blob/feb39ed3788e22fd83a1cb294a824e60f4c13187/newsfragments/aip-72.significant.rst#L21-L37

fyi: the following is what I had to do for XComs:

PRs:

Remove XCom pickling

Migrate pickled data & change XCom value type to JSON

Handle XCom value errors in FAB's XCom list view

Use JSONB type for XCom.value column in PostgreSQL

…rate_conf_column_as_json

…onomer/airflow into dag_run_migrate_conf_column_as_json

airflow/migrations/versions/0050_3_0_0_remove_pickled_data_from_dagrun_table.py

In [PR](#44166) we added migration for removing pickled data from `xcom` table. During my testing I noticed with `SQLite` [insert](https://github.com/apache/airflow/blob/main/airflow/migrations/versions/0049_3_0_0_remove_pickled_data_from_xcom_table.py#L88) statement is not working in case of upgrade. Changing condition to `hex(substr(value, 1, 1)) = '80'` works. Tested [here](#44533 (comment)). related: #44166

jedcunningham · 2024-12-03T17:07:20Z

airflow/migrations/versions/0050_3_0_0_remove_pickled_data_from_dagrun_table.py

+    )
+
+    # Delete the pickled data from the dag_run table so that we can update the column type
+    conn.execute(text(f"DELETE FROM dag_run WHERE conf IS NOT NULL AND {condition}"))


We don't want to delete the rows right? That'll cascade down and delete other stuff too, like TIs.

Should we just set conf to null instead?

One difference with this situation vs with xcom is the xcom value is the "valuable" part of the row - this conf isn't though.

Yes, this is better.

Should we just set conf to null instead?

Yes, this is better.

Should we just set conf to null instead?

This is done

jedcunningham · 2024-12-03T17:10:57Z

airflow/migrations/versions/0050_3_0_0_remove_pickled_data_from_dagrun_table.py

+
+    insert_query = text(
+        """
+                       INSERT INTO dag_run (dag_id, run_id, queued_at, logical_date, start_date, end_date, state, creating_job_id, external_trigger, run_type, triggered_by, conf, data_interval_start, data_interval_end, last_scheduling_decision, log_template_id, updated_at, clear_number, backfill_id, dag_version_id)


This doesn't feel right - what if I upgraded 2 weeks ago and have since done a db clean?!

This will only run if _dag_run_archive exits.

Removed this

jedcunningham · 2024-12-03T17:18:33Z

airflow/migrations/versions/0050_3_0_0_remove_pickled_data_from_dagrun_table.py

+        sa.Column("logical_date", sa.TIMESTAMP(timezone=True), nullable=False),
+        sa.Column("start_date", sa.TIMESTAMP(timezone=True), nullable=True),
+        sa.Column("end_date", sa.TIMESTAMP(timezone=True), nullable=True),
+        sa.Column("state", sa.String(length=50), nullable=True),


Not sure we really need a full record for the archive. I also kinda question if we really need one to keep the old pickled confs even.

We can simplify and nuke

nuking will simplify a lot.

We are now not archiving the record.

jedcunningham · 2024-12-05T22:32:32Z

airflow/migrations/versions/0050_3_0_0_remove_pickled_data_from_dagrun_table.py

+    dialect = conn.dialect.name
+
+    # Update the dag_run.conf column value to NULL
+    conn.execute(text("UPDATE dag_run set conf=null WHERE conf IS NOT NULL"))


This just gets rid of all of the values? They may not all be pickled, right?

I assumed that if we upgrade to this new version, all values in the configuration would be pickled. Should we add a condition to check for pickled data?

Since DagRun.conf uses PickleType, all values should be pickled

airflow/airflow/models/dagrun.py

Line 140 in b69441d

conf = Column(PickleType)

jedcunningham · 2024-12-05T22:34:36Z

airflow/migrations/versions/0050_3_0_0_remove_pickled_data_from_dagrun_table.py

+            END
+            """
+        )
+    elif dialect == "mysql":


And for mysql/sqlite, even if we had data, we don't move any of it from the old to new column.

update command would be executed for mysql/sqlite as well. Also I was getting this error earlier when we were moving data after conversion.

@jedcunningham as discussed I have update mysql and sqlite commands for downgrade

…rate_conf_column_as_json

kaxil · 2024-12-12T11:16:55Z

airflow/migrations/versions/0051_3_0_0_remove_pickled_data_from_dagrun_table.py

+            """
+            ALTER TABLE dag_run
+            ALTER COLUMN conf TYPE JSONB
+            USING CAST(CONVERT_FROM(conf, 'UTF8') AS JSONB);


Since everything is already null, do we need this line?

Assuming we don't, we can have same code for all 3 DB types

apart from JSONB (Postgres) vs JSON (MySQL & SQLite) where you can do something like:

JSON().with_variant(postgresql.JSONB, "postgresql")

Implemented same logic for all 3 DB with condition JSON().with_variant(postgresql.JSONB, "postgresql")

kaxil · 2024-12-12T11:21:34Z

airflow/migrations/versions/0051_3_0_0_remove_pickled_data_from_dagrun_table.py

+        op.add_column("dag_run", sa.Column("conf_json", sa.JSON(), nullable=True))
+        op.drop_column("dag_run", "conf")
+        op.alter_column("dag_run", "conf_json", existing_type=sa.JSON(), new_column_name="conf")


Do we need to drop the column? Since all values are NULL, can’t we just change the type?

For SQLite ALTER column type is not supported.

…onomer/airflow into dag_run_migrate_conf_column_as_json

kaxil · 2024-12-13T18:05:16Z

airflow/migrations/versions/0052_3_0_0_remove_pickled_data_from_dagrun_table.py

+    conn.execute(text("UPDATE dag_run set conf=null WHERE conf IS NOT NULL"))
+
+    # Update the dag_run.conf column value to NULL to avoid issues during the type change
+    conn.execute(text("UPDATE dag_run set conf=null WHERE conf IS NOT NULL"))


kaxil · 2024-12-13T18:06:30Z

airflow/migrations/versions/0052_3_0_0_remove_pickled_data_from_dagrun_table.py

+    conn.execute(text("UPDATE dag_run set conf=null WHERE conf IS NOT NULL"))
+
+    conf_type = sa.JSON().with_variant(postgresql.JSONB, "postgresql")
+
+    with op.batch_alter_table("dag_run", schema=None) as batch_op:
+        # Drop the existing column for SQLite (due to its limitations)
+        batch_op.drop_column("conf")


Do we need to set conf=null, if we are going to drop that column anyway?

Worth for you to take a look too @jedcunningham @ephraimbuddy

…rate_conf_column_as_json

remove pickled data from dag run table

b91446b

vatsrahul1001 requested review from potiuk, XD-DENG and ashb as code owners December 1, 2024 13:11

boring-cyborg bot added area:db-migrations PRs with DB migration kind:documentation labels Dec 1, 2024

vatsrahul1001 marked this pull request as draft December 1, 2024 13:11

vatsrahul1001 marked this pull request as ready for review December 2, 2024 10:32

vatsrahul1001 marked this pull request as draft December 2, 2024 10:33

vatsrahul1001 marked this pull request as ready for review December 2, 2024 13:23

kaxil reviewed Dec 2, 2024

View reviewed changes

Merge branch 'main' of github.com:astronomer/airflow into dag_run_mig…

cbde2ed

…rate_conf_column_as_json

vatsrahul1001 added the airflow3.0:breaking Candidates for Airflow 3.0 that contain breaking changes label Dec 3, 2024

fix downgrade + add news fragement

e698769

vatsrahul1001 requested review from ryanahamilton, bbovenzi, pierrejeambrun and jscheffl as code owners December 3, 2024 10:28

vatsrahul1001 and others added 3 commits December 3, 2024 17:03

Merge branch 'main' into dag_run_migrate_conf_column_as_json

f7f9155

remove archive table if exits after downgrade

18c03fb

Merge branch 'dag_run_migrate_conf_column_as_json' of github.com:astr…

2b6abf7

…onomer/airflow into dag_run_migrate_conf_column_as_json

kaxil reviewed Dec 3, 2024

View reviewed changes

airflow/migrations/versions/0050_3_0_0_remove_pickled_data_from_dagrun_table.py Outdated Show resolved Hide resolved

vatsrahul1001 mentioned this pull request Dec 3, 2024

Fix SQLite condition for detecting pickled data #44611

Merged

jedcunningham reviewed Dec 3, 2024

View reviewed changes

removing archiving data

a54c435

vatsrahul1001 requested a review from jedcunningham December 5, 2024 15:12

vatsrahul1001 requested a review from kaxil December 5, 2024 15:12

jedcunningham reviewed Dec 5, 2024

View reviewed changes

vatsrahul1001 requested a review from jedcunningham December 6, 2024 05:08

vatsrahul1001 and others added 5 commits December 10, 2024 15:30

implementing review comments

6b3789d

Merge branch 'main' of github.com:astronomer/airflow into dag_run_mig…

07258f5

…rate_conf_column_as_json

fixing static check

8b88d63

Merge branch 'main' into dag_run_migrate_conf_column_as_json

730470b

fixing static checks

c2e0d09

vatsrahul1001 added the full tests needed We need to run full set of tests for this PR to merge label Dec 12, 2024

Merge branch 'main' into dag_run_migrate_conf_column_as_json

c2b8dbd

vatsrahul1001 requested a review from ephraimbuddy as a code owner December 12, 2024 09:38

kaxil reviewed Dec 12, 2024

View reviewed changes

vatsrahul1001 added 3 commits December 13, 2024 15:57

simplying upgrade and downgrade as per review

c8fa27d

simplying upgrade and downgrade as per review

d27b7ed

Merge branch 'dag_run_migrate_conf_column_as_json' of github.com:astr…

603ca17

…onomer/airflow into dag_run_migrate_conf_column_as_json

vatsrahul1001 requested a review from kaxil December 13, 2024 15:33

fixing failures

9406702

kaxil reviewed Dec 13, 2024

View reviewed changes

vatsrahul1001 added 4 commits December 14, 2024 08:18

Merge branch 'main' of github.com:astronomer/airflow into dag_run_mig…

01c453c

…rate_conf_column_as_json

Merge branch 'main' of github.com:astronomer/airflow into dag_run_mig…

219b694

…rate_conf_column_as_json

removing setting conf to null

22a5dba

refactor approach to migrate values in conf

70ce5df

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Update conf column in dag_run table type from bytea to JSON #44533

Update conf column in dag_run table type from bytea to JSON #44533

vatsrahul1001 commented Dec 1, 2024

vatsrahul1001 commented Dec 2, 2024

kaxil commented Dec 2, 2024

kaxil Dec 2, 2024

kaxil Dec 2, 2024

kaxil Dec 2, 2024

kaxil Dec 2, 2024 •

edited

Loading

kaxil Dec 2, 2024 •

edited

Loading

jedcunningham Dec 3, 2024

jedcunningham Dec 3, 2024

vatsrahul1001 Dec 3, 2024

vatsrahul1001 Dec 4, 2024

jedcunningham Dec 3, 2024

vatsrahul1001 Dec 3, 2024

vatsrahul1001 Dec 4, 2024

jedcunningham Dec 3, 2024

kaxil Dec 3, 2024

vatsrahul1001 Dec 3, 2024

vatsrahul1001 Dec 4, 2024

jedcunningham Dec 5, 2024

vatsrahul1001 Dec 6, 2024

vatsrahul1001 Dec 10, 2024

kaxil Dec 13, 2024

jedcunningham Dec 5, 2024

vatsrahul1001 Dec 6, 2024

vatsrahul1001 Dec 10, 2024

kaxil Dec 12, 2024

kaxil Dec 12, 2024

kaxil Dec 12, 2024

vatsrahul1001 Dec 13, 2024

kaxil Dec 12, 2024

vatsrahul1001 Dec 13, 2024

kaxil Dec 13, 2024

kaxil Dec 13, 2024

kaxil Dec 13, 2024

Update conf column in dag_run table type from bytea to JSON #44533

Are you sure you want to change the base?

Update conf column in dag_run table type from bytea to JSON #44533

Conversation

vatsrahul1001 commented Dec 1, 2024

vatsrahul1001 commented Dec 2, 2024

kaxil commented Dec 2, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kaxil Dec 2, 2024 • edited Loading

Choose a reason for hiding this comment

kaxil Dec 2, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kaxil Dec 2, 2024 •

edited

Loading

kaxil Dec 2, 2024 •

edited

Loading