ephraimbuddy commented on code in PR #44790:
URL: https://github.com/apache/airflow/pull/44790#discussion_r1883514615
##########
airflow/migrations/versions/0027_2_10_3_fix_dag_schedule_dataset_alias_reference_naming.py:
##########
@@ -39,91 +41,130 @@
depends_on = None
airflow_version = "2.10.3"
-if TYPE_CHECKING:
- from alembic.operations.base import BatchOperations
- from sqlalchemy.sql.elements import conv
-
-
-def _rename_fk_constraint(
- *,
- batch_op: BatchOperations,
- original_name: str | conv,
- new_name: str | conv,
- referent_table: str,
- local_cols: list[str],
- remote_cols: list[str],
- ondelete: str,
-) -> None:
- batch_op.drop_constraint(original_name, type_="foreignkey")
- batch_op.create_foreign_key(
- constraint_name=new_name,
- referent_table=referent_table,
- local_cols=local_cols,
- remote_cols=remote_cols,
- ondelete=ondelete,
- )
-
def upgrade():
"""Rename dag_schedule_dataset_alias_reference constraint."""
- with op.batch_alter_table("dag_schedule_dataset_alias_reference",
schema=None) as batch_op:
- bind = op.get_context().bind
- insp = inspect(bind)
- fk_constraints = [fk["name"] for fk in
insp.get_foreign_keys("dag_schedule_dataset_alias_reference")]
-
- # "dsdar_dataset_alias_fkey" was the constraint name defined in the
model while "dsdar_dataset_fkey" is the one
- # defined in the previous migration.
- # Rename this constraint name if user is using the name
"dsdar_dataset_fkey".
- if "dsdar_dataset_fkey" in fk_constraints:
- _rename_fk_constraint(
- batch_op=batch_op,
- original_name="dsdar_dataset_fkey",
- new_name="dsdar_dataset_alias_fkey",
+ dialect = op.get_context().dialect.name
+ if dialect != "sqlite":
+ with op.batch_alter_table("dag_schedule_dataset_alias_reference",
schema=None) as batch_op:
+ if dialect == "postgresql":
+ op.execute(
+ "ALTER TABLE dag_schedule_dataset_alias_reference DROP
CONSTRAINT IF EXISTS dsdar_dataset_fkey"
+ )
+ op.execute(
+ "ALTER TABLE dag_schedule_dataset_alias_reference DROP
CONSTRAINT IF EXISTS dsdar_dag_fkey"
+ )
+ if dialect == "mysql":
+ mysql_drop_foreignkey_if_exists(
+ "dsdar_dataset_fkey",
"dag_schedule_dataset_alias_reference", op
+ )
+ mysql_drop_foreignkey_if_exists("dsdar_dag_fkey",
"dag_schedule_dataset_alias_reference", op)
+ batch_op.create_foreign_key(
+ constraint_name="dsdar_dataset_alias_fkey",
Review Comment:
There's a need to check if this fk name already exists before adding it.
Working on that now...
--
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]
