This is an automated email from the ASF dual-hosted git repository. husseinawala 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 6f0e5bce59 Add a migration script for encrypted trigger kwargs (#38358) 6f0e5bce59 is described below commit 6f0e5bce590de567d5010220d2ca392b22630a7b Author: Hussein Awala <huss...@awala.fr> AuthorDate: Mon Apr 1 15:36:29 2024 +0200 Add a migration script for encrypted trigger kwargs (#38358) --- .../0140_2_9_0_update_trigger_kwargs_type.py | 50 ++++++++++++++++++++++ airflow/utils/db.py | 41 +++++++++++++++++- docs/apache-airflow/img/airflow_erd.sha256 | 2 +- docs/apache-airflow/img/airflow_erd.svg | 8 ++-- docs/apache-airflow/migrations-ref.rst | 4 +- 5 files changed, 98 insertions(+), 7 deletions(-) diff --git a/airflow/migrations/versions/0140_2_9_0_update_trigger_kwargs_type.py b/airflow/migrations/versions/0140_2_9_0_update_trigger_kwargs_type.py new file mode 100644 index 0000000000..dbde1201e4 --- /dev/null +++ b/airflow/migrations/versions/0140_2_9_0_update_trigger_kwargs_type.py @@ -0,0 +1,50 @@ +# +# 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. + +"""update trigger kwargs type + +Revision ID: 1949afb29106 +Revises: ee1467d4aa35 +Create Date: 2024-03-17 22:09:09.406395 + +""" +import sqlalchemy as sa + +from airflow.models.trigger import Trigger +from alembic import op + +from airflow.utils.sqlalchemy import ExtendedJSON + +# revision identifiers, used by Alembic. +revision = "1949afb29106" +down_revision = "ee1467d4aa35" +branch_labels = None +depends_on = None +airflow_version = "2.9.0" + + +def upgrade(): + """Update trigger kwargs type to string""" + with op.batch_alter_table("trigger") as batch_op: + batch_op.alter_column("kwargs", type_=sa.Text(), ) + + +def downgrade(): + """Unapply update trigger kwargs type to string""" + with op.batch_alter_table("trigger") as batch_op: + batch_op.alter_column("kwargs", type_=ExtendedJSON(), postgresql_using="kwargs::json") diff --git a/airflow/utils/db.py b/airflow/utils/db.py index 6cc0d7159b..26697bb118 100644 --- a/airflow/utils/db.py +++ b/airflow/utils/db.py @@ -90,7 +90,7 @@ _REVISION_HEADS_MAP = { "2.7.0": "405de8318b3a", "2.8.0": "10b52ebd31f7", "2.8.1": "88344c1d9134", - "2.9.0": "1fd565369930", + "2.9.0": "1949afb29106", } @@ -972,6 +972,33 @@ def synchronize_log_template(*, session: Session = NEW_SESSION) -> None: session.add(LogTemplate(filename=filename, elasticsearch_id=elasticsearch_id)) +def encrypt_trigger_kwargs(*, session: Session) -> None: + """Encrypt trigger kwargs.""" + from airflow.models.trigger import Trigger + from airflow.serialization.serialized_objects import BaseSerialization + + for trigger in session.query(Trigger): + # convert serialized dict to string and encrypt it + trigger.kwargs = BaseSerialization.deserialize(json.loads(trigger.encrypted_kwargs)) + session.commit() + + +def decrypt_trigger_kwargs(*, session: Session) -> None: + """Decrypt trigger kwargs.""" + from airflow.models.trigger import Trigger + from airflow.serialization.serialized_objects import BaseSerialization + + if not inspect(session.bind).has_table(Trigger.__tablename__): + # table does not exist, nothing to do + # this can happen when we downgrade to an old version before the Trigger table was added + return + + for trigger in session.query(Trigger): + # decrypt the string and convert it to serialized dict + trigger.encrypted_kwargs = json.dumps(BaseSerialization.serialize(trigger.kwargs)) + session.commit() + + def check_conn_id_duplicates(session: Session) -> Iterable[str]: """ Check unique conn_id in connection table. @@ -1639,6 +1666,12 @@ def upgradedb( _reserialize_dags(session=session) add_default_pool_if_not_exists(session=session) synchronize_log_template(session=session) + if _revision_greater( + config, + _REVISION_HEADS_MAP["2.9.0"], + _get_current_revision(session=session), + ): + encrypt_trigger_kwargs(session=session) @provide_session @@ -1711,6 +1744,12 @@ def downgrade(*, to_revision, from_revision=None, show_sql_only=False, session: else: log.info("Applying downgrade migrations.") command.downgrade(config, revision=to_revision, sql=show_sql_only) + if _revision_greater( + config, + _REVISION_HEADS_MAP["2.9.0"], + to_revision, + ): + decrypt_trigger_kwargs(session=session) def drop_airflow_models(connection): diff --git a/docs/apache-airflow/img/airflow_erd.sha256 b/docs/apache-airflow/img/airflow_erd.sha256 index 7a1c0a9fba..09f84daea2 100644 --- a/docs/apache-airflow/img/airflow_erd.sha256 +++ b/docs/apache-airflow/img/airflow_erd.sha256 @@ -1 +1 @@ -3ec33f4a14388277f9aba431c06c3bfa9d044ab2eae466aa394aa9618d2f3eb5 \ No newline at end of file +2a24225537326f38be5df14e0b7a8dca867122093e0fa932f1a11ac12d1fb11c \ No newline at end of file diff --git a/docs/apache-airflow/img/airflow_erd.svg b/docs/apache-airflow/img/airflow_erd.svg index 144668b900..dc32fe0566 100644 --- a/docs/apache-airflow/img/airflow_erd.svg +++ b/docs/apache-airflow/img/airflow_erd.svg @@ -1361,28 +1361,28 @@ <g id="edge41" class="edge"> <title>task_instance--xcom</title> <path fill="none" stroke="#7f7f7f" stroke-dasharray="5,2" d="M1198.1,-805.22C1228.72,-801.45 1260.55,-798.25 1290.36,-795.96"/> -<text text-anchor="start" x="1259.36" y="-784.76" font-family="Times,serif" font-size="14.00">0..N</text> +<text text-anchor="start" x="1280.36" y="-784.76" font-family="Times,serif" font-size="14.00">1</text> <text text-anchor="start" x="1198.1" y="-794.02" font-family="Times,serif" font-size="14.00">1</text> </g> <!-- task_instance--xcom --> <g id="edge42" class="edge"> <title>task_instance--xcom</title> <path fill="none" stroke="#7f7f7f" stroke-dasharray="5,2" d="M1198.1,-818.52C1228.72,-815.1 1260.55,-811.78 1290.36,-808.91"/> -<text text-anchor="start" x="1280.36" y="-812.71" font-family="Times,serif" font-size="14.00">1</text> +<text text-anchor="start" x="1259.36" y="-812.71" font-family="Times,serif" font-size="14.00">0..N</text> <text text-anchor="start" x="1198.1" y="-822.32" font-family="Times,serif" font-size="14.00">1</text> </g> <!-- task_instance--xcom --> <g id="edge43" class="edge"> <title>task_instance--xcom</title> <path fill="none" stroke="#7f7f7f" stroke-dasharray="5,2" d="M1198.1,-831.83C1228.72,-828.75 1260.55,-825.32 1290.36,-821.86"/> -<text text-anchor="start" x="1259.36" y="-825.66" font-family="Times,serif" font-size="14.00">0..N</text> +<text text-anchor="start" x="1280.36" y="-825.66" font-family="Times,serif" font-size="14.00">1</text> <text text-anchor="start" x="1198.1" y="-835.63" font-family="Times,serif" font-size="14.00">1</text> </g> <!-- task_instance--xcom --> <g id="edge44" class="edge"> <title>task_instance--xcom</title> <path fill="none" stroke="#7f7f7f" stroke-dasharray="5,2" d="M1198.1,-845.13C1228.72,-842.41 1260.55,-838.85 1290.36,-834.82"/> -<text text-anchor="start" x="1280.36" y="-838.62" font-family="Times,serif" font-size="14.00">1</text> +<text text-anchor="start" x="1259.36" y="-838.62" font-family="Times,serif" font-size="14.00">0..N</text> <text text-anchor="start" x="1198.1" y="-848.93" font-family="Times,serif" font-size="14.00">1</text> </g> <!-- rendered_task_instance_fields --> diff --git a/docs/apache-airflow/migrations-ref.rst b/docs/apache-airflow/migrations-ref.rst index a28f2ee766..13c70abe9d 100644 --- a/docs/apache-airflow/migrations-ref.rst +++ b/docs/apache-airflow/migrations-ref.rst @@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru +---------------------------------+-------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=================================+===================+===================+==============================================================+ -| ``ee1467d4aa35`` (head) | ``b4078ac230a1`` | ``2.9.0`` | add display name for dag and task instance | +| ``1949afb29106`` (head) | ``ee1467d4aa35`` | ``2.9.0`` | update trigger kwargs type | ++---------------------------------+-------------------+-------------------+--------------------------------------------------------------+ +| ``ee1467d4aa35`` | ``b4078ac230a1`` | ``2.9.0`` | add display name for dag and task instance | +---------------------------------+-------------------+-------------------+--------------------------------------------------------------+ | ``b4078ac230a1`` | ``8e1c784a4fc7`` | ``2.9.0`` | Change value column type to longblob in xcom table for mysql | +---------------------------------+-------------------+-------------------+--------------------------------------------------------------+