You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by hu...@apache.org on 2024/04/01 13:36:37 UTC
(airflow) branch main updated: Add a migration script for encrypted trigger kwargs (#38358)
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 <hu...@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 |
+---------------------------------+-------------------+-------------------+--------------------------------------------------------------+