diff --git a/airflow-core/src/airflow/migrations/versions/0017_2_9_2_fix_inconsistency_between_ORM_and_migration_files.py b/airflow-core/src/airflow/migrations/versions/0017_2_9_2_fix_inconsistency_between_ORM_and_migration_files.py index 64531a957b302..26d2e86176cf1 100644 --- a/airflow-core/src/airflow/migrations/versions/0017_2_9_2_fix_inconsistency_between_ORM_and_migration_files.py +++ b/airflow-core/src/airflow/migrations/versions/0017_2_9_2_fix_inconsistency_between_ORM_and_migration_files.py @@ -28,7 +28,7 @@ from __future__ import annotations import sqlalchemy as sa -from alembic import op +from alembic import context, op from sqlalchemy import literal # revision identifiers, used by Alembic. @@ -39,40 +39,59 @@ airflow_version = "2.9.2" -def upgrade(): - """Apply Update missing constraints.""" - conn = op.get_bind() - if conn.dialect.name == "mysql": - # TODO: Rewrite these queries to use alembic when lowest MYSQL version supports IF EXISTS - conn.execute( - sa.text(""" - set @var=if((SELECT true FROM information_schema.TABLE_CONSTRAINTS WHERE - CONSTRAINT_SCHEMA = DATABASE() AND - TABLE_NAME = 'connection' AND - CONSTRAINT_NAME = 'unique_conn_id' AND - CONSTRAINT_TYPE = 'UNIQUE') = true,'ALTER TABLE connection - DROP INDEX unique_conn_id','select 1'); +def _mysql_drop_unique_constraint_if_exists(conn, table: str, index_name: str) -> None: + """ + Drop a MySQL unique constraint only if it is actually present. - prepare stmt from @var; - execute stmt; - deallocate prepare stmt; - """) - ) - # Dropping the below and recreating cause there's no IF NOT EXISTS in mysql + MySQL has no ``DROP INDEX IF EXISTS``, and PyMySQL does not enable + ``CLIENT.MULTI_STATEMENTS``, so it cannot run the + ``prepare``/``execute``/``deallocate prepare`` script in a single + ``cursor.execute()`` call, so the existence check and the drop are issued as two + separate single statements. In offline (``--sql``) mode there is no live connection + to query information_schema against, so the guarded dynamic SQL is emitted as literal + script text instead, to be run later through a real SQL client that supports + multi-statement scripts. + """ + if context.is_offline_mode(): conn.execute( - sa.text(""" + sa.text(f""" set @var=if((SELECT true FROM information_schema.TABLE_CONSTRAINTS WHERE CONSTRAINT_SCHEMA = DATABASE() AND - TABLE_NAME = 'connection' AND - CONSTRAINT_NAME = 'connection_conn_id_uq' AND - CONSTRAINT_TYPE = 'UNIQUE') = true,'ALTER TABLE connection - DROP INDEX connection_conn_id_uq','select 1'); + TABLE_NAME = '{table}' AND + CONSTRAINT_NAME = '{index_name}' AND + CONSTRAINT_TYPE = 'UNIQUE') = true,'ALTER TABLE {table} + DROP INDEX {index_name}','select 1'); prepare stmt from @var; execute stmt; deallocate prepare stmt; """) ) + return + existing_indexes = { + row[0] + for row in conn.execute( + sa.text(""" + SELECT CONSTRAINT_NAME FROM information_schema.TABLE_CONSTRAINTS + WHERE CONSTRAINT_SCHEMA = DATABASE() + AND TABLE_NAME = :table + AND CONSTRAINT_TYPE = 'UNIQUE' + """), + {"table": table}, + ) + } + if index_name in existing_indexes: + conn.execute(sa.text(f"ALTER TABLE {table} DROP INDEX {index_name}")) + + +def upgrade(): + """Apply Update missing constraints.""" + conn = op.get_bind() + if conn.dialect.name == "mysql": + # TODO: Rewrite these queries to use alembic when lowest MYSQL version supports IF EXISTS + # Dropping the below and recreating cause there's no IF NOT EXISTS in mysql + for index_name in ("unique_conn_id", "connection_conn_id_uq"): + _mysql_drop_unique_constraint_if_exists(conn, "connection", index_name) elif conn.dialect.name == "sqlite": # SQLite does not support DROP CONSTRAINT # We have to recreate the table without the constraint @@ -121,63 +140,14 @@ def upgrade(): batch_op.drop_constraint("task_reschedule_dr_fkey", type_="foreignkey") if conn.dialect.name == "mysql": - conn.execute( - sa.text(""" - set @var=if((SELECT true FROM information_schema.TABLE_CONSTRAINTS WHERE - CONSTRAINT_SCHEMA = DATABASE() AND - TABLE_NAME = 'dag_run' AND - CONSTRAINT_NAME = 'dag_run_dag_id_execution_date_uq' AND - CONSTRAINT_TYPE = 'UNIQUE') = true,'ALTER TABLE dag_run - DROP INDEX dag_run_dag_id_execution_date_uq','select 1'); - - prepare stmt from @var; - execute stmt; - deallocate prepare stmt; - """) - ) - conn.execute( - sa.text(""" - set @var=if((SELECT true FROM information_schema.TABLE_CONSTRAINTS WHERE - CONSTRAINT_SCHEMA = DATABASE() AND - TABLE_NAME = 'dag_run' AND - CONSTRAINT_NAME = 'dag_run_dag_id_run_id_uq' AND - CONSTRAINT_TYPE = 'UNIQUE') = true,'ALTER TABLE dag_run - DROP INDEX dag_run_dag_id_run_id_uq','select 1'); - - prepare stmt from @var; - execute stmt; - deallocate prepare stmt; - """) - ) # below we drop and recreate the constraints because there's no IF NOT EXISTS - conn.execute( - sa.text(""" - set @var=if((SELECT true FROM information_schema.TABLE_CONSTRAINTS WHERE - CONSTRAINT_SCHEMA = DATABASE() AND - TABLE_NAME = 'dag_run' AND - CONSTRAINT_NAME = 'dag_run_dag_id_execution_date_key' AND - CONSTRAINT_TYPE = 'UNIQUE') = true,'ALTER TABLE dag_run - DROP INDEX dag_run_dag_id_execution_date_key','select 1'); - - prepare stmt from @var; - execute stmt; - deallocate prepare stmt; - """) - ) - conn.execute( - sa.text(""" - set @var=if((SELECT true FROM information_schema.TABLE_CONSTRAINTS WHERE - CONSTRAINT_SCHEMA = DATABASE() AND - TABLE_NAME = 'dag_run' AND - CONSTRAINT_NAME = 'dag_run_dag_id_run_id_key' AND - CONSTRAINT_TYPE = 'UNIQUE') = true,'ALTER TABLE dag_run - DROP INDEX dag_run_dag_id_run_id_key','select 1'); - - prepare stmt from @var; - execute stmt; - deallocate prepare stmt; - """) - ) + for index_name in ( + "dag_run_dag_id_execution_date_uq", + "dag_run_dag_id_run_id_uq", + "dag_run_dag_id_execution_date_key", + "dag_run_dag_id_run_id_key", + ): + _mysql_drop_unique_constraint_if_exists(conn, "dag_run", index_name) with op.batch_alter_table("callback_request", schema=None) as batch_op: batch_op.alter_column( "processor_subdir",