Copilot commented on code in PR #64838:
URL: https://github.com/apache/airflow/pull/64838#discussion_r3173682258
##########
airflow-core/src/airflow/migrations/utils.py:
##########
@@ -61,3 +69,192 @@ def ignore_sqlite_value_error():
if op.get_bind().dialect.name == "sqlite":
return contextlib.suppress(ValueError)
return contextlib.nullcontext()
+
+
+def create_index_if_not_exists(op, index_name, table_name, columns,
unique=False) -> None:
+ """
+ Create an index if it does not already exist.
+
+ MySQL does not support CREATE INDEX IF NOT EXISTS, so a stored procedure
is used.
+ PostgreSQL and SQLite support it natively.
+ """
+ dialect_name = get_dialect_name(op)
+
+ if dialect_name == "mysql":
+ unique_kw = "UNIQUE " if unique else ""
+ col_list = ", ".join(f"`{c}`" for c in columns)
+ op.execute(
+ text(f"""
+ DROP PROCEDURE IF EXISTS CreateIndexIfNotExists;
+ CREATE PROCEDURE CreateIndexIfNotExists()
+ BEGIN
+ IF NOT EXISTS (
+ SELECT 1
+ FROM information_schema.STATISTICS
+ WHERE
+ TABLE_SCHEMA = DATABASE() AND
+ TABLE_NAME = '{table_name}' AND
+ INDEX_NAME = '{index_name}'
+ ) THEN
+ CREATE {unique_kw}INDEX `{index_name}` ON `{table_name}`
({col_list});
+ END IF;
+ END;
+ CALL CreateIndexIfNotExists();
+ DROP PROCEDURE IF EXISTS CreateIndexIfNotExists;
+ """)
+ )
+ else:
+ op.create_index(index_name, table_name, columns, unique=unique,
if_not_exists=True)
+
+
+def drop_index_if_exists(op, index_name, table_name) -> None:
+ """
+ Drop an index if it exists.
+
+ Works in both online and offline mode by using raw SQL for PostgreSQL and
MySQL.
+ SQLite and PostgreSQL support DROP INDEX IF EXISTS natively.
+ MySQL requires a stored procedure since it does not support IF EXISTS for
DROP INDEX.
+ """
+ dialect_name = get_dialect_name(op)
+
+ if dialect_name == "mysql":
+ op.execute(
+ text(f"""
+ CREATE PROCEDURE DropIndexIfExists()
+ BEGIN
+ IF EXISTS (
+ SELECT 1
+ FROM information_schema.STATISTICS
+ WHERE
+ TABLE_SCHEMA = DATABASE() AND
+ TABLE_NAME = '{table_name}' AND
+ INDEX_NAME = '{index_name}'
+ ) THEN
+ DROP INDEX `{index_name}` ON `{table_name}`;
+ END IF;
+ END;
+ CALL DropIndexIfExists();
+ DROP PROCEDURE DropIndexIfExists;
+ """)
Review Comment:
`drop_index_if_exists()` creates the MySQL stored procedure
`DropIndexIfExists()` without first doing `DROP PROCEDURE IF EXISTS ...`. If
the migration run is interrupted after `CREATE PROCEDURE` but before the final
drop, rerunning can fail because the procedure already exists. Add a `DROP
PROCEDURE IF EXISTS DropIndexIfExists;` prelude (similar to
`create_index_if_not_exists()`).
##########
airflow-core/src/airflow/migrations/utils.py:
##########
@@ -61,3 +69,192 @@ def ignore_sqlite_value_error():
if op.get_bind().dialect.name == "sqlite":
return contextlib.suppress(ValueError)
return contextlib.nullcontext()
+
+
+def create_index_if_not_exists(op, index_name, table_name, columns,
unique=False) -> None:
+ """
+ Create an index if it does not already exist.
+
+ MySQL does not support CREATE INDEX IF NOT EXISTS, so a stored procedure
is used.
+ PostgreSQL and SQLite support it natively.
+ """
+ dialect_name = get_dialect_name(op)
+
+ if dialect_name == "mysql":
+ unique_kw = "UNIQUE " if unique else ""
+ col_list = ", ".join(f"`{c}`" for c in columns)
+ op.execute(
+ text(f"""
+ DROP PROCEDURE IF EXISTS CreateIndexIfNotExists;
+ CREATE PROCEDURE CreateIndexIfNotExists()
+ BEGIN
+ IF NOT EXISTS (
+ SELECT 1
+ FROM information_schema.STATISTICS
+ WHERE
+ TABLE_SCHEMA = DATABASE() AND
+ TABLE_NAME = '{table_name}' AND
+ INDEX_NAME = '{index_name}'
+ ) THEN
+ CREATE {unique_kw}INDEX `{index_name}` ON `{table_name}`
({col_list});
+ END IF;
+ END;
+ CALL CreateIndexIfNotExists();
+ DROP PROCEDURE IF EXISTS CreateIndexIfNotExists;
+ """)
+ )
+ else:
+ op.create_index(index_name, table_name, columns, unique=unique,
if_not_exists=True)
+
+
+def drop_index_if_exists(op, index_name, table_name) -> None:
+ """
+ Drop an index if it exists.
+
+ Works in both online and offline mode by using raw SQL for PostgreSQL and
MySQL.
+ SQLite and PostgreSQL support DROP INDEX IF EXISTS natively.
+ MySQL requires a stored procedure since it does not support IF EXISTS for
DROP INDEX.
+ """
+ dialect_name = get_dialect_name(op)
+
+ if dialect_name == "mysql":
+ op.execute(
+ text(f"""
+ CREATE PROCEDURE DropIndexIfExists()
+ BEGIN
+ IF EXISTS (
+ SELECT 1
+ FROM information_schema.STATISTICS
+ WHERE
+ TABLE_SCHEMA = DATABASE() AND
+ TABLE_NAME = '{table_name}' AND
+ INDEX_NAME = '{index_name}'
+ ) THEN
+ DROP INDEX `{index_name}` ON `{table_name}`;
+ END IF;
+ END;
+ CALL DropIndexIfExists();
+ DROP PROCEDURE DropIndexIfExists;
+ """)
+ )
+ else:
+ # PostgreSQL and SQLite both support DROP INDEX IF EXISTS
+ op.drop_index(index_name, table_name=table_name, if_exists=True)
+
+
+def drop_unique_constraints_on_columns(op, table_name, columns) -> None:
+ """
+ Drop all unique constraints covering any of the given columns, regardless
of constraint name.
+
+ Works in both online and offline mode by using raw SQL for PostgreSQL and
MySQL.
+ SQLite falls back to batch mode and requires a live connection.
+ """
+ import sqlalchemy as sa
+
+ dialect_name = get_dialect_name(op)
+
+ if dialect_name == "postgresql":
+ cols_array = ", ".join(f"'{c}'" for c in columns)
+ op.execute(
+ text(f"""
+ DO $$
+ DECLARE r record;
+ BEGIN
+ FOR r IN
+ SELECT DISTINCT tc.constraint_name
+ FROM information_schema.table_constraints tc
+ JOIN information_schema.key_column_usage kcu
+ ON tc.constraint_name = kcu.constraint_name
+ AND tc.table_schema = kcu.table_schema
+ WHERE tc.table_name = '{table_name}'
+ AND tc.constraint_type = 'UNIQUE'
+ AND kcu.column_name = ANY(ARRAY[{cols_array}]::text[])
+ LOOP
+ EXECUTE 'ALTER TABLE ' || quote_ident('{table_name}') || '
DROP CONSTRAINT IF EXISTS '
+ || quote_ident(r.constraint_name);
+ END LOOP;
+ END $$
+ """)
+ )
+ elif dialect_name == "mysql":
+ cols_in = ", ".join(f"'{c}'" for c in columns)
+ op.execute(
+ text(f"""
+ CREATE PROCEDURE DropUniqueOnColumns()
+ BEGIN
+ DECLARE done INT DEFAULT FALSE;
+ DECLARE v_name VARCHAR(255);
+ DECLARE cur CURSOR FOR
+ SELECT DISTINCT kcu.CONSTRAINT_NAME
Review Comment:
In the MySQL branch of `drop_unique_constraints_on_columns()`, the helper
creates a stored procedure `DropUniqueOnColumns()` but doesn't `DROP PROCEDURE
IF EXISTS` before creating it. If a previous run failed mid-way, reruns can
error on `CREATE PROCEDURE` due to the leftover procedure. Consider adding a
defensive drop prelude.
##########
airflow-core/src/airflow/migrations/utils.py:
##########
@@ -61,3 +69,192 @@ def ignore_sqlite_value_error():
if op.get_bind().dialect.name == "sqlite":
Review Comment:
`ignore_sqlite_value_error()` calls `op.get_bind().dialect.name`, which will
raise in Alembic offline mode (where `get_bind()` is `None`). Since this module
now has `get_dialect_name(op)` for offline-safe dialect detection, use that
here as well so offline `--show-sql-only` migrations don't crash on non-SQLite
dialects.
##########
airflow-core/src/airflow/migrations/utils.py:
##########
@@ -61,3 +69,192 @@ def ignore_sqlite_value_error():
if op.get_bind().dialect.name == "sqlite":
return contextlib.suppress(ValueError)
return contextlib.nullcontext()
+
+
+def create_index_if_not_exists(op, index_name, table_name, columns,
unique=False) -> None:
+ """
+ Create an index if it does not already exist.
+
+ MySQL does not support CREATE INDEX IF NOT EXISTS, so a stored procedure
is used.
+ PostgreSQL and SQLite support it natively.
+ """
+ dialect_name = get_dialect_name(op)
+
+ if dialect_name == "mysql":
+ unique_kw = "UNIQUE " if unique else ""
+ col_list = ", ".join(f"`{c}`" for c in columns)
+ op.execute(
+ text(f"""
+ DROP PROCEDURE IF EXISTS CreateIndexIfNotExists;
+ CREATE PROCEDURE CreateIndexIfNotExists()
+ BEGIN
+ IF NOT EXISTS (
+ SELECT 1
+ FROM information_schema.STATISTICS
+ WHERE
+ TABLE_SCHEMA = DATABASE() AND
+ TABLE_NAME = '{table_name}' AND
+ INDEX_NAME = '{index_name}'
+ ) THEN
+ CREATE {unique_kw}INDEX `{index_name}` ON `{table_name}`
({col_list});
+ END IF;
+ END;
+ CALL CreateIndexIfNotExists();
+ DROP PROCEDURE IF EXISTS CreateIndexIfNotExists;
+ """)
+ )
+ else:
+ op.create_index(index_name, table_name, columns, unique=unique,
if_not_exists=True)
+
+
+def drop_index_if_exists(op, index_name, table_name) -> None:
+ """
+ Drop an index if it exists.
+
+ Works in both online and offline mode by using raw SQL for PostgreSQL and
MySQL.
+ SQLite and PostgreSQL support DROP INDEX IF EXISTS natively.
+ MySQL requires a stored procedure since it does not support IF EXISTS for
DROP INDEX.
+ """
+ dialect_name = get_dialect_name(op)
+
+ if dialect_name == "mysql":
+ op.execute(
+ text(f"""
+ CREATE PROCEDURE DropIndexIfExists()
+ BEGIN
+ IF EXISTS (
+ SELECT 1
+ FROM information_schema.STATISTICS
+ WHERE
+ TABLE_SCHEMA = DATABASE() AND
+ TABLE_NAME = '{table_name}' AND
+ INDEX_NAME = '{index_name}'
+ ) THEN
+ DROP INDEX `{index_name}` ON `{table_name}`;
+ END IF;
+ END;
+ CALL DropIndexIfExists();
+ DROP PROCEDURE DropIndexIfExists;
+ """)
+ )
+ else:
+ # PostgreSQL and SQLite both support DROP INDEX IF EXISTS
+ op.drop_index(index_name, table_name=table_name, if_exists=True)
+
+
+def drop_unique_constraints_on_columns(op, table_name, columns) -> None:
+ """
+ Drop all unique constraints covering any of the given columns, regardless
of constraint name.
+
+ Works in both online and offline mode by using raw SQL for PostgreSQL and
MySQL.
+ SQLite falls back to batch mode and requires a live connection.
+ """
+ import sqlalchemy as sa
+
+ dialect_name = get_dialect_name(op)
+
+ if dialect_name == "postgresql":
+ cols_array = ", ".join(f"'{c}'" for c in columns)
+ op.execute(
+ text(f"""
+ DO $$
+ DECLARE r record;
+ BEGIN
+ FOR r IN
+ SELECT DISTINCT tc.constraint_name
+ FROM information_schema.table_constraints tc
+ JOIN information_schema.key_column_usage kcu
+ ON tc.constraint_name = kcu.constraint_name
+ AND tc.table_schema = kcu.table_schema
+ WHERE tc.table_name = '{table_name}'
+ AND tc.constraint_type = 'UNIQUE'
+ AND kcu.column_name = ANY(ARRAY[{cols_array}]::text[])
+ LOOP
+ EXECUTE 'ALTER TABLE ' || quote_ident('{table_name}') || '
DROP CONSTRAINT IF EXISTS '
+ || quote_ident(r.constraint_name);
+ END LOOP;
+ END $$
+ """)
+ )
+ elif dialect_name == "mysql":
+ cols_in = ", ".join(f"'{c}'" for c in columns)
+ op.execute(
+ text(f"""
+ CREATE PROCEDURE DropUniqueOnColumns()
+ BEGIN
+ DECLARE done INT DEFAULT FALSE;
+ DECLARE v_name VARCHAR(255);
+ DECLARE cur CURSOR FOR
+ SELECT DISTINCT kcu.CONSTRAINT_NAME
+ FROM information_schema.KEY_COLUMN_USAGE kcu
+ JOIN information_schema.TABLE_CONSTRAINTS tc
+ ON kcu.CONSTRAINT_NAME = tc.CONSTRAINT_NAME
+ AND kcu.TABLE_SCHEMA = tc.TABLE_SCHEMA
+ AND kcu.TABLE_NAME = tc.TABLE_NAME
+ WHERE kcu.TABLE_NAME = '{table_name}'
+ AND kcu.TABLE_SCHEMA = DATABASE()
+ AND tc.CONSTRAINT_TYPE = 'UNIQUE'
+ AND kcu.COLUMN_NAME IN ({cols_in});
+ DECLARE CONTINUE HANDLER FOR NOT FOUND SET done = TRUE;
+ OPEN cur;
+ drop_loop: LOOP
+ FETCH cur INTO v_name;
+ IF done THEN LEAVE drop_loop; END IF;
+ SET @stmt = CONCAT('ALTER TABLE `{table_name}` DROP INDEX
`', v_name, '`');
+ PREPARE s FROM @stmt;
+ EXECUTE s;
+ DEALLOCATE PREPARE s;
+ END LOOP;
+ CLOSE cur;
+ END;
+ CALL DropUniqueOnColumns();
+ DROP PROCEDURE DropUniqueOnColumns;
+ """)
+ )
+ else:
+ # SQLite — batch mode rewrites the table; requires a live connection
+ with op.batch_alter_table(table_name, schema=None) as batch_op:
+ for uq in
sa.inspect(op.get_bind()).get_unique_constraints(table_name):
+ if any(col in uq["column_names"] for col in columns):
+ batch_op.drop_constraint(uq["name"], type_="unique")
+
+
+def drop_unique_constraint_if_exists(op, table_name, constraint_name) -> None:
+ """
+ Drop a unique constraint by name if it exists.
+
+ Works in both online and offline mode by using raw SQL for PostgreSQL and
MySQL.
+ SQLite falls back to batch mode and requires a live connection.
+ """
+ dialect_name = get_dialect_name(op)
+
+ if dialect_name == "postgresql":
+ op.execute(text(f'ALTER TABLE "{table_name}" DROP CONSTRAINT IF EXISTS
"{constraint_name}"'))
+ elif dialect_name == "mysql":
+ op.execute(
+ text(f"""
Review Comment:
In the MySQL branch of `drop_unique_constraint_if_exists()`, the helper
creates a stored procedure `DropUniqueIfExists()` but doesn't `DROP PROCEDURE
IF EXISTS` before creating it. A failed migration run could leave the procedure
behind and make reruns fail on `CREATE PROCEDURE`. Add a defensive `DROP
PROCEDURE IF EXISTS DropUniqueIfExists;` prelude (consistent with
`create_index_if_not_exists()`).
##########
airflow-core/src/airflow/migrations/utils.py:
##########
@@ -61,3 +69,192 @@ def ignore_sqlite_value_error():
if op.get_bind().dialect.name == "sqlite":
return contextlib.suppress(ValueError)
Review Comment:
There is a `from alembic import op` import inside
`ignore_sqlite_value_error()`. Imports inside functions are discouraged in this
codebase (and make this helper harder to reason about for offline/online
behavior). Move the import to module scope (or accept `op` as a parameter) and
reuse `get_dialect_name()` for the dialect check.
##########
airflow-core/src/airflow/migrations/utils.py:
##########
@@ -61,3 +69,192 @@ def ignore_sqlite_value_error():
if op.get_bind().dialect.name == "sqlite":
return contextlib.suppress(ValueError)
return contextlib.nullcontext()
+
+
+def create_index_if_not_exists(op, index_name, table_name, columns,
unique=False) -> None:
+ """
+ Create an index if it does not already exist.
+
+ MySQL does not support CREATE INDEX IF NOT EXISTS, so a stored procedure
is used.
Review Comment:
This PR's description focuses on SQLite FK PRAGMA teardown, but this file
also adds new cross-dialect index/constraint helpers that don't appear to be
used by the migrations changed here. Consider updating the PR description to
cover these additions (and adding tests), or splitting them into a separate PR
to keep scope/review risk lower.
##########
airflow-core/src/airflow/migrations/utils.py:
##########
@@ -61,3 +69,192 @@ def ignore_sqlite_value_error():
if op.get_bind().dialect.name == "sqlite":
return contextlib.suppress(ValueError)
return contextlib.nullcontext()
+
+
+def create_index_if_not_exists(op, index_name, table_name, columns,
unique=False) -> None:
+ """
+ Create an index if it does not already exist.
+
+ MySQL does not support CREATE INDEX IF NOT EXISTS, so a stored procedure
is used.
+ PostgreSQL and SQLite support it natively.
+ """
+ dialect_name = get_dialect_name(op)
+
+ if dialect_name == "mysql":
+ unique_kw = "UNIQUE " if unique else ""
+ col_list = ", ".join(f"`{c}`" for c in columns)
+ op.execute(
+ text(f"""
+ DROP PROCEDURE IF EXISTS CreateIndexIfNotExists;
+ CREATE PROCEDURE CreateIndexIfNotExists()
+ BEGIN
+ IF NOT EXISTS (
+ SELECT 1
+ FROM information_schema.STATISTICS
+ WHERE
+ TABLE_SCHEMA = DATABASE() AND
+ TABLE_NAME = '{table_name}' AND
+ INDEX_NAME = '{index_name}'
+ ) THEN
+ CREATE {unique_kw}INDEX `{index_name}` ON `{table_name}`
({col_list});
+ END IF;
+ END;
+ CALL CreateIndexIfNotExists();
+ DROP PROCEDURE IF EXISTS CreateIndexIfNotExists;
+ """)
+ )
+ else:
+ op.create_index(index_name, table_name, columns, unique=unique,
if_not_exists=True)
+
+
+def drop_index_if_exists(op, index_name, table_name) -> None:
+ """
+ Drop an index if it exists.
+
+ Works in both online and offline mode by using raw SQL for PostgreSQL and
MySQL.
+ SQLite and PostgreSQL support DROP INDEX IF EXISTS natively.
+ MySQL requires a stored procedure since it does not support IF EXISTS for
DROP INDEX.
+ """
+ dialect_name = get_dialect_name(op)
+
+ if dialect_name == "mysql":
+ op.execute(
+ text(f"""
+ CREATE PROCEDURE DropIndexIfExists()
+ BEGIN
+ IF EXISTS (
+ SELECT 1
+ FROM information_schema.STATISTICS
+ WHERE
+ TABLE_SCHEMA = DATABASE() AND
+ TABLE_NAME = '{table_name}' AND
+ INDEX_NAME = '{index_name}'
+ ) THEN
+ DROP INDEX `{index_name}` ON `{table_name}`;
+ END IF;
+ END;
+ CALL DropIndexIfExists();
+ DROP PROCEDURE DropIndexIfExists;
+ """)
+ )
+ else:
+ # PostgreSQL and SQLite both support DROP INDEX IF EXISTS
+ op.drop_index(index_name, table_name=table_name, if_exists=True)
+
+
+def drop_unique_constraints_on_columns(op, table_name, columns) -> None:
+ """
+ Drop all unique constraints covering any of the given columns, regardless
of constraint name.
+
+ Works in both online and offline mode by using raw SQL for PostgreSQL and
MySQL.
+ SQLite falls back to batch mode and requires a live connection.
+ """
+ import sqlalchemy as sa
+
+ dialect_name = get_dialect_name(op)
Review Comment:
`drop_unique_constraints_on_columns()` introduces a function-local `import
sqlalchemy as sa`. Per project guidelines, prefer module-level imports unless
there's a concrete circular-import/lazy-load need; moving this import to the
top of the module will avoid repeated imports and makes dependencies clearer.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
