From ad939298d3384e703eb4ac240087dc47f2813787 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Sun, 5 Jul 2020 00:20:34 +0100 Subject: [PATCH 1/2] Update migrations to ensure compatibility with Airflow 1.10.* closes https://github.com/apache/airflow/issues/9640 --- .../3c20cacc0044_add_dagrun_run_type.py | 4 +- .../versions/4ebbffe0a39a_merge_heads.py | 39 ------------------- ...add_rendered_task_instance_fields_table.py | 4 +- ...c8_task_reschedule_fk_on_cascade_delete.py | 4 +- ...dd_precision_to_execution_date_in_mysql.py | 4 +- .../versions/b0125267960b_merge_heads.py | 6 +-- ...5a55525161_increase_length_of_pool_name.py | 4 +- ...bbf4a7ad0465_remove_id_column_from_xcom.py | 4 +- .../cf5dc11e79ad_drop_user_and_chart.py | 4 +- 9 files changed, 17 insertions(+), 56 deletions(-) delete mode 100644 airflow/migrations/versions/4ebbffe0a39a_merge_heads.py diff --git a/airflow/migrations/versions/3c20cacc0044_add_dagrun_run_type.py b/airflow/migrations/versions/3c20cacc0044_add_dagrun_run_type.py index 9a4ea5a471460..e650e519835ff 100644 --- a/airflow/migrations/versions/3c20cacc0044_add_dagrun_run_type.py +++ b/airflow/migrations/versions/3c20cacc0044_add_dagrun_run_type.py @@ -20,7 +20,7 @@ Add DagRun run_type Revision ID: 3c20cacc0044 -Revises: 952da73b5eff +Revises: b25a55525161 Create Date: 2020-04-08 13:35:25.671327 """ @@ -38,7 +38,7 @@ # revision identifiers, used by Alembic. revision = "3c20cacc0044" -down_revision = "952da73b5eff" +down_revision = "b25a55525161" branch_labels = None depends_on = None diff --git a/airflow/migrations/versions/4ebbffe0a39a_merge_heads.py b/airflow/migrations/versions/4ebbffe0a39a_merge_heads.py deleted file mode 100644 index 28bc60716f507..0000000000000 --- a/airflow/migrations/versions/4ebbffe0a39a_merge_heads.py +++ /dev/null @@ -1,39 +0,0 @@ -# -# 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. - -"""Merge heads - -Revision ID: 4ebbffe0a39a -Revises: dd4ecb8fbee3, cf5dc11e79ad -Create Date: 2019-02-04 20:19:50.628137 - -""" - -# revision identifiers, used by Alembic. -revision = '4ebbffe0a39a' -down_revision = ('dd4ecb8fbee3', 'cf5dc11e79ad') -branch_labels = None -depends_on = None - - -def upgrade(): # noqa: D103 - pass - - -def downgrade(): # noqa: D103 - pass diff --git a/airflow/migrations/versions/852ae6c715af_add_rendered_task_instance_fields_table.py b/airflow/migrations/versions/852ae6c715af_add_rendered_task_instance_fields_table.py index a7b06fc469633..01357e7c6f706 100644 --- a/airflow/migrations/versions/852ae6c715af_add_rendered_task_instance_fields_table.py +++ b/airflow/migrations/versions/852ae6c715af_add_rendered_task_instance_fields_table.py @@ -19,7 +19,7 @@ """Add RenderedTaskInstanceFields table Revision ID: 852ae6c715af -Revises: b25a55525161 +Revises: a4c2fd67d16b Create Date: 2020-03-10 22:19:18.034961 """ @@ -29,7 +29,7 @@ # revision identifiers, used by Alembic. revision = '852ae6c715af' -down_revision = 'b25a55525161' +down_revision = 'a4c2fd67d16b' branch_labels = None depends_on = None diff --git a/airflow/migrations/versions/939bb1e647c8_task_reschedule_fk_on_cascade_delete.py b/airflow/migrations/versions/939bb1e647c8_task_reschedule_fk_on_cascade_delete.py index 797c1f3e66bba..4465ba1ae5b8c 100644 --- a/airflow/migrations/versions/939bb1e647c8_task_reschedule_fk_on_cascade_delete.py +++ b/airflow/migrations/versions/939bb1e647c8_task_reschedule_fk_on_cascade_delete.py @@ -18,7 +18,7 @@ """task reschedule fk on cascade delete Revision ID: 939bb1e647c8 -Revises: 4ebbffe0a39a +Revises: dd4ecb8fbee3 Create Date: 2019-02-04 20:21:50.669751 """ @@ -27,7 +27,7 @@ # revision identifiers, used by Alembic. revision = '939bb1e647c8' -down_revision = '4ebbffe0a39a' +down_revision = 'dd4ecb8fbee3' branch_labels = None depends_on = None diff --git a/airflow/migrations/versions/a66efa278eea_add_precision_to_execution_date_in_mysql.py b/airflow/migrations/versions/a66efa278eea_add_precision_to_execution_date_in_mysql.py index ecb589df4d4ab..7c122ef3df1c6 100644 --- a/airflow/migrations/versions/a66efa278eea_add_precision_to_execution_date_in_mysql.py +++ b/airflow/migrations/versions/a66efa278eea_add_precision_to_execution_date_in_mysql.py @@ -19,7 +19,7 @@ """Add Precision to execution_date in RenderedTaskInstanceFields table Revision ID: a66efa278eea -Revises: 8f966b9c467a +Revises: 952da73b5eff Create Date: 2020-06-16 21:44:02.883132 """ @@ -29,7 +29,7 @@ # revision identifiers, used by Alembic. revision = 'a66efa278eea' -down_revision = '8f966b9c467a' +down_revision = '952da73b5eff' branch_labels = None depends_on = None diff --git a/airflow/migrations/versions/b0125267960b_merge_heads.py b/airflow/migrations/versions/b0125267960b_merge_heads.py index d0f546e065b04..007b99a2a3921 100644 --- a/airflow/migrations/versions/b0125267960b_merge_heads.py +++ b/airflow/migrations/versions/b0125267960b_merge_heads.py @@ -16,17 +16,17 @@ # specific language governing permissions and limitations # under the License. -"""merge multiple heads +"""Merge the four heads back together Revision ID: 08364691d074 -Revises: a56c9515abdc, 74effc47d867, b3b105409875, bbf4a7ad0465 +Revises: a56c9515abdc, 004c1210f153, 74effc47d867, b3b105409875 Create Date: 2019-11-19 22:05:11.752222 """ # revision identifiers, used by Alembic. revision = '08364691d074' -down_revision = ('a56c9515abdc', '74effc47d867', 'b3b105409875', 'bbf4a7ad0465') +down_revision = ('a56c9515abdc', '004c1210f153', '74effc47d867', 'b3b105409875') branch_labels = None depends_on = None diff --git a/airflow/migrations/versions/b25a55525161_increase_length_of_pool_name.py b/airflow/migrations/versions/b25a55525161_increase_length_of_pool_name.py index 2e6b5d97b08cc..8d343bc90af68 100644 --- a/airflow/migrations/versions/b25a55525161_increase_length_of_pool_name.py +++ b/airflow/migrations/versions/b25a55525161_increase_length_of_pool_name.py @@ -19,7 +19,7 @@ """Increase length of pool name Revision ID: b25a55525161 -Revises: a4c2fd67d16b +Revises: bbf4a7ad0465 Create Date: 2020-03-09 08:48:14.534700 """ @@ -29,7 +29,7 @@ # revision identifiers, used by Alembic. revision = 'b25a55525161' -down_revision = 'a4c2fd67d16b' +down_revision = 'bbf4a7ad0465' branch_labels = None depends_on = None diff --git a/airflow/migrations/versions/bbf4a7ad0465_remove_id_column_from_xcom.py b/airflow/migrations/versions/bbf4a7ad0465_remove_id_column_from_xcom.py index aa6fbeedae545..f0098dc7d4ad5 100644 --- a/airflow/migrations/versions/bbf4a7ad0465_remove_id_column_from_xcom.py +++ b/airflow/migrations/versions/bbf4a7ad0465_remove_id_column_from_xcom.py @@ -19,7 +19,7 @@ """Remove id column from xcom Revision ID: bbf4a7ad0465 -Revises: 004c1210f153 +Revises: cf5dc11e79ad Create Date: 2019-10-29 13:53:09.445943 """ @@ -29,7 +29,7 @@ # revision identifiers, used by Alembic. revision = 'bbf4a7ad0465' -down_revision = '004c1210f153' +down_revision = 'cf5dc11e79ad' branch_labels = None depends_on = None diff --git a/airflow/migrations/versions/cf5dc11e79ad_drop_user_and_chart.py b/airflow/migrations/versions/cf5dc11e79ad_drop_user_and_chart.py index 1a48f3629142a..8f8db1884b9e5 100644 --- a/airflow/migrations/versions/cf5dc11e79ad_drop_user_and_chart.py +++ b/airflow/migrations/versions/cf5dc11e79ad_drop_user_and_chart.py @@ -18,7 +18,7 @@ """drop_user_and_chart Revision ID: cf5dc11e79ad -Revises: 41f5f12752f8 +Revises: a66efa278eea Create Date: 2019-01-24 15:30:35.834740 """ @@ -29,7 +29,7 @@ # revision identifiers, used by Alembic. revision = 'cf5dc11e79ad' -down_revision = '41f5f12752f8' +down_revision = 'a66efa278eea' branch_labels = None depends_on = None From 4aad12cf6d7ac425a50b8a5a1f541c02c8cec822 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Sun, 5 Jul 2020 16:48:04 +0100 Subject: [PATCH 2/2] Add safety checks for Master/main branch users --- .../3c20cacc0044_add_dagrun_run_type.py | 42 +++++++++++-------- ...bbf4a7ad0465_remove_id_column_from_xcom.py | 12 ++++-- .../cf5dc11e79ad_drop_user_and_chart.py | 10 +++-- 3 files changed, 40 insertions(+), 24 deletions(-) diff --git a/airflow/migrations/versions/3c20cacc0044_add_dagrun_run_type.py b/airflow/migrations/versions/3c20cacc0044_add_dagrun_run_type.py index e650e519835ff..cc3de10ac3227 100644 --- a/airflow/migrations/versions/3c20cacc0044_add_dagrun_run_type.py +++ b/airflow/migrations/versions/3c20cacc0044_add_dagrun_run_type.py @@ -28,6 +28,7 @@ import sqlalchemy as sa from alembic import op from sqlalchemy import Boolean, Column, Integer, PickleType, String +from sqlalchemy.engine.reflection import Inspector from sqlalchemy.ext.declarative import declarative_base from airflow.models.base import ID_LEN @@ -68,28 +69,33 @@ def upgrade(): """Apply Add DagRun run_type""" run_type_col_type = sa.String(length=50) - # Add nullable column - with op.batch_alter_table("dag_run") as batch_op: - batch_op.add_column(sa.Column("run_type", run_type_col_type, nullable=True)) + conn = op.get_bind() + inspector = Inspector.from_engine(conn) + dag_run_columns = [col.get('name') for col in inspector.get_columns("dag_run")] - # Generate run type for existing records - connection = op.get_bind() - sessionmaker = sa.orm.sessionmaker() - session = sessionmaker(bind=connection) + if "run_type" not in dag_run_columns: - for run_type in DagRunType: - session.query(DagRun).filter(DagRun.run_id.like(f"{run_type.value}__%")).update( - {DagRun.run_type: run_type.value}, synchronize_session=False - ) + # Add nullable column + with op.batch_alter_table("dag_run") as batch_op: + batch_op.add_column(sa.Column("run_type", run_type_col_type, nullable=True)) + + # Generate run type for existing records + sessionmaker = sa.orm.sessionmaker() + session = sessionmaker(bind=conn) - session.query(DagRun).filter(DagRun.run_type.is_(None)).update( - {DagRun.run_type: DagRunType.MANUAL.value}, synchronize_session=False - ) - session.commit() + for run_type in DagRunType: + session.query(DagRun).filter(DagRun.run_id.like(f"{run_type.value}__%")).update( + {DagRun.run_type: run_type.value}, synchronize_session=False + ) + + session.query(DagRun).filter(DagRun.run_type.is_(None)).update( + {DagRun.run_type: DagRunType.MANUAL.value}, synchronize_session=False + ) + session.commit() - # Make run_type not nullable - with op.batch_alter_table("dag_run") as batch_op: - batch_op.alter_column("run_type", type_=run_type_col_type, nullable=False) + # Make run_type not nullable + with op.batch_alter_table("dag_run") as batch_op: + batch_op.alter_column("run_type", type_=run_type_col_type, nullable=False) def downgrade(): diff --git a/airflow/migrations/versions/bbf4a7ad0465_remove_id_column_from_xcom.py b/airflow/migrations/versions/bbf4a7ad0465_remove_id_column_from_xcom.py index f0098dc7d4ad5..54c55e74022ee 100644 --- a/airflow/migrations/versions/bbf4a7ad0465_remove_id_column_from_xcom.py +++ b/airflow/migrations/versions/bbf4a7ad0465_remove_id_column_from_xcom.py @@ -26,6 +26,7 @@ from alembic import op from sqlalchemy import Column, Integer +from sqlalchemy.engine.reflection import Inspector # revision identifiers, used by Alembic. revision = 'bbf4a7ad0465' @@ -36,10 +37,15 @@ def upgrade(): """Apply Remove id column from xcom""" + conn = op.get_bind() + inspector = Inspector.from_engine(conn) + with op.batch_alter_table('xcom') as bop: - bop.drop_column('id') - bop.drop_index('idx_xcom_dag_task_date') - bop.create_primary_key('pk_xcom', ['dag_id', 'task_id', 'key', 'execution_date']) + xcom_columns = [col.get('name') for col in inspector.get_columns("xcom")] + if "id" in xcom_columns: + bop.drop_column('id') + bop.drop_index('idx_xcom_dag_task_date') + bop.create_primary_key('pk_xcom', ['dag_id', 'task_id', 'key', 'execution_date']) def downgrade(): diff --git a/airflow/migrations/versions/cf5dc11e79ad_drop_user_and_chart.py b/airflow/migrations/versions/cf5dc11e79ad_drop_user_and_chart.py index 8f8db1884b9e5..86e976983a154 100644 --- a/airflow/migrations/versions/cf5dc11e79ad_drop_user_and_chart.py +++ b/airflow/migrations/versions/cf5dc11e79ad_drop_user_and_chart.py @@ -43,12 +43,16 @@ def upgrade(): # noqa: D103 conn = op.get_bind() inspector = Inspector.from_engine(conn) + tables = inspector.get_table_names() - if 'known_event' in inspector.get_table_names() != 'sqlite': + if 'known_event' in tables: op.drop_constraint('known_event_user_id_fkey', 'known_event') - op.drop_table("chart") - op.drop_table("users") + if "chart" in tables: + op.drop_table("chart", ) + + if "users" in tables: + op.drop_table("users") def downgrade(): # noqa: D103