diff --git a/airflow-core/newsfragments/71249.improvement.rst b/airflow-core/newsfragments/71249.improvement.rst new file mode 100644 index 0000000000000..43085779f4f35 --- /dev/null +++ b/airflow-core/newsfragments/71249.improvement.rst @@ -0,0 +1 @@ +The collation used for the asset ``name``, ``uri`` and ``group`` columns on MySQL is now configurable through ``[database] sql_engine_collation_for_asset_names``. It still defaults to ``latin1_general_cs``, so existing databases are unaffected. Set it if your MySQL-compatible engine does not provide that collation. diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index c065b277716c4..10022889dd1fc 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -638,6 +638,17 @@ database: type: string example: ~ default: ~ + sql_engine_collation_for_asset_names: + description: | + Collation for the ``name``, ``uri`` and ``group`` columns of the asset tables on + ``mysql`` and ``mariadb``. These columns hold ASCII values and are indexed at 1500 + characters, so a single-byte charset is used to stay within the maximum index size. + Override this if your database engine does not provide ``latin1_general_cs`` -- + for example TiDB, which supports ``latin1_bin`` instead. + version_added: 3.4.0 + type: string + example: "latin1_bin" + default: "latin1_general_cs" sql_alchemy_pool_enabled: description: | If SQLAlchemy should pool database connections. diff --git a/airflow-core/src/airflow/migrations/utils.py b/airflow-core/src/airflow/migrations/utils.py index 4eeaf373c6a87..226928ca3b843 100644 --- a/airflow-core/src/airflow/migrations/utils.py +++ b/airflow-core/src/airflow/migrations/utils.py @@ -61,3 +61,16 @@ def ignore_sqlite_value_error(): if op.get_bind().dialect.name == "sqlite": return contextlib.suppress(ValueError) return contextlib.nullcontext() + + +def asset_name_collation() -> str: + """ + Return the MySQL collation for asset name/uri/group columns. + + Mirrors ``airflow.models.base.get_asset_str_field``. Migrations resolve it at + run time rather than hard-coding it, because MySQL-compatible engines do not + all ship ``latin1_general_cs`` and a fresh install replays these migrations. + """ + from airflow.configuration import conf + + return conf.get("database", "sql_engine_collation_for_asset_names", fallback="latin1_general_cs") diff --git a/airflow-core/src/airflow/migrations/versions/0000_2_6_2_squashed_migrations.py b/airflow-core/src/airflow/migrations/versions/0000_2_6_2_squashed_migrations.py index 739e4e7f88768..49f30872754d0 100644 --- a/airflow-core/src/airflow/migrations/versions/0000_2_6_2_squashed_migrations.py +++ b/airflow-core/src/airflow/migrations/versions/0000_2_6_2_squashed_migrations.py @@ -34,6 +34,7 @@ from sqlalchemy.dialects.mysql import MEDIUMTEXT from airflow.migrations.db_types import StringID +from airflow.migrations.utils import asset_name_collation from airflow.utils.sqlalchemy import ExtendedJSON, UtcDateTime # revision identifiers, used by Alembic. @@ -208,7 +209,7 @@ def upgrade() -> None: length=3000, # latin1 allows for more indexed length in mysql # and this field should only be ascii chars - collation="latin1_general_cs", + collation=asset_name_collation(), ), "mysql", ), diff --git a/airflow-core/src/airflow/migrations/versions/0022_2_10_0_add_dataset_alias.py b/airflow-core/src/airflow/migrations/versions/0022_2_10_0_add_dataset_alias.py index 0d4a9efe0ebac..0cbe6841a4f62 100644 --- a/airflow-core/src/airflow/migrations/versions/0022_2_10_0_add_dataset_alias.py +++ b/airflow-core/src/airflow/migrations/versions/0022_2_10_0_add_dataset_alias.py @@ -30,6 +30,8 @@ import sqlalchemy as sa from alembic import op +from airflow.migrations.utils import asset_name_collation + # revision identifiers, used by Alembic. revision = "05e19f3176be" down_revision = "d482b7261ff9" @@ -46,7 +48,7 @@ def upgrade(): sa.Column( "name", sa.String(length=3000).with_variant( - sa.String(length=3000, collation="latin1_general_cs"), "mysql" + sa.String(length=3000, collation=asset_name_collation()), "mysql" ), nullable=False, ), diff --git a/airflow-core/src/airflow/migrations/versions/0036_3_0_0_add_name_field_to_dataset_model.py b/airflow-core/src/airflow/migrations/versions/0036_3_0_0_add_name_field_to_dataset_model.py index b1f925dbffca2..a5a5cb5512b95 100644 --- a/airflow-core/src/airflow/migrations/versions/0036_3_0_0_add_name_field_to_dataset_model.py +++ b/airflow-core/src/airflow/migrations/versions/0036_3_0_0_add_name_field_to_dataset_model.py @@ -39,6 +39,8 @@ import sqlalchemy as sa from alembic import op +from airflow.migrations.utils import asset_name_collation + # revision identifiers, used by Alembic. revision = "0d9e73a75ee4" down_revision = "44eabb1904b4" @@ -47,7 +49,7 @@ airflow_version = "3.0.0" _STRING_COLUMN_TYPE = sa.String(length=1500).with_variant( - sa.String(length=1500, collation="latin1_general_cs"), + sa.String(length=1500, collation=asset_name_collation()), "mysql", ) @@ -127,7 +129,7 @@ def downgrade(): batch_op.alter_column( "uri", type_=sa.String(length=3000).with_variant( - sa.String(length=3000, collation="latin1_general_cs"), + sa.String(length=3000, collation=asset_name_collation()), "mysql", ), nullable=False, diff --git a/airflow-core/src/airflow/migrations/versions/0038_3_0_0_add_asset_active.py b/airflow-core/src/airflow/migrations/versions/0038_3_0_0_add_asset_active.py index 2a992cab4126e..39b8d8151692b 100644 --- a/airflow-core/src/airflow/migrations/versions/0038_3_0_0_add_asset_active.py +++ b/airflow-core/src/airflow/migrations/versions/0038_3_0_0_add_asset_active.py @@ -30,6 +30,8 @@ import sqlalchemy as sa from alembic import op +from airflow.migrations.utils import asset_name_collation + # revision identifiers, used by Alembic. revision = "5a5d66100783" down_revision = "c3389cd7793f" @@ -38,7 +40,7 @@ airflow_version = "3.0.0" _STRING_COLUMN_TYPE = sa.String(length=1500).with_variant( - sa.String(length=1500, collation="latin1_general_cs"), + sa.String(length=1500, collation=asset_name_collation()), "mysql", ) diff --git a/airflow-core/src/airflow/migrations/versions/0039_3_0_0_tweak_assetaliasmodel_to_match_asset.py b/airflow-core/src/airflow/migrations/versions/0039_3_0_0_tweak_assetaliasmodel_to_match_asset.py index d0067f1288255..79ded55d6d83d 100644 --- a/airflow-core/src/airflow/migrations/versions/0039_3_0_0_tweak_assetaliasmodel_to_match_asset.py +++ b/airflow-core/src/airflow/migrations/versions/0039_3_0_0_tweak_assetaliasmodel_to_match_asset.py @@ -42,6 +42,8 @@ import sqlalchemy as sa from alembic import op +from airflow.migrations.utils import asset_name_collation + # Revision identifiers, used by Alembic. revision = "fb2d4922cd79" down_revision = "5a5d66100783" @@ -50,7 +52,7 @@ airflow_version = "3.0.0" _STRING_COLUMN_TYPE = sa.String(length=1500).with_variant( - sa.String(length=1500, collation="latin1_general_cs"), + sa.String(length=1500, collation=asset_name_collation()), "mysql", ) @@ -76,7 +78,7 @@ def downgrade(): batch_op.alter_column( "name", type_=sa.String(length=3000).with_variant( - sa.String(length=3000, collation="latin1_general_cs"), + sa.String(length=3000, collation=asset_name_collation()), "mysql", ), nullable=False, diff --git a/airflow-core/src/airflow/migrations/versions/0054_3_0_0_add_asset_reference_models.py b/airflow-core/src/airflow/migrations/versions/0054_3_0_0_add_asset_reference_models.py index 4a34da03ecc2a..0eb50af20adb7 100644 --- a/airflow-core/src/airflow/migrations/versions/0054_3_0_0_add_asset_reference_models.py +++ b/airflow-core/src/airflow/migrations/versions/0054_3_0_0_add_asset_reference_models.py @@ -30,6 +30,7 @@ from alembic import op from airflow.migrations.db_types import StringID +from airflow.migrations.utils import asset_name_collation from airflow.utils.sqlalchemy import UtcDateTime # revision identifiers, used by Alembic. @@ -40,7 +41,7 @@ airflow_version = "3.0.0" ASSET_STR_FIELD = sa.String(length=1500).with_variant( - sa.String(length=1500, collation="latin1_general_cs"), "mysql" + sa.String(length=1500, collation=asset_name_collation()), "mysql" ) diff --git a/airflow-core/src/airflow/migrations/versions/0088_3_2_0_replace_asset_trigger_table_with_asset.py b/airflow-core/src/airflow/migrations/versions/0088_3_2_0_replace_asset_trigger_table_with_asset.py index 81c134741f4d2..17ab84969d16a 100644 --- a/airflow-core/src/airflow/migrations/versions/0088_3_2_0_replace_asset_trigger_table_with_asset.py +++ b/airflow-core/src/airflow/migrations/versions/0088_3_2_0_replace_asset_trigger_table_with_asset.py @@ -30,6 +30,8 @@ import sqlalchemy as sa from alembic import op +from airflow.migrations.utils import asset_name_collation + # revision identifiers, used by Alembic. revision = "15d84ca19038" down_revision = "509b94a1042d" @@ -38,7 +40,7 @@ airflow_version = "3.2.0" _STRING_COLUMN_TYPE = sa.String(length=1500).with_variant( - sa.String(length=1500, collation="latin1_general_cs"), + sa.String(length=1500, collation=asset_name_collation()), "mysql", ) diff --git a/airflow-core/src/airflow/models/asset.py b/airflow-core/src/airflow/models/asset.py index 750aac8d2b73b..ccfbc72c268c4 100644 --- a/airflow-core/src/airflow/models/asset.py +++ b/airflow-core/src/airflow/models/asset.py @@ -30,7 +30,6 @@ Index, Integer, PrimaryKeyConstraint, - String, Table, delete, select, @@ -39,7 +38,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from airflow._shared.timezones import timezone -from airflow.models.base import Base, StringID +from airflow.models.base import ASSET_STR_FIELD, Base, StringID from airflow.utils.sqlalchemy import UtcDateTime if TYPE_CHECKING: @@ -147,15 +146,7 @@ class AssetWatcherModel(Base): """A table to store asset watchers.""" name: Mapped[str] = mapped_column( - String(length=1500).with_variant( - String( - length=1500, - # latin1 allows for more indexed length in mysql - # and this field should only be ascii chars - collation="latin1_general_cs", - ), - "mysql", - ), + ASSET_STR_FIELD, nullable=False, ) asset_id: Mapped[int] = mapped_column(Integer, primary_key=True, nullable=False) @@ -198,27 +189,11 @@ class AssetAliasModel(Base): id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) name: Mapped[str] = mapped_column( - String(length=1500).with_variant( - String( - length=1500, - # latin1 allows for more indexed length in mysql - # and this field should only be ascii chars - collation="latin1_general_cs", - ), - "mysql", - ), + ASSET_STR_FIELD, nullable=False, ) group: Mapped[str] = mapped_column( - String(length=1500).with_variant( - String( - length=1500, - # latin1 allows for more indexed length in mysql - # and this field should only be ascii chars - collation="latin1_general_cs", - ), - "mysql", - ), + ASSET_STR_FIELD, default="", nullable=False, ) @@ -279,39 +254,15 @@ class AssetModel(Base): id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) name: Mapped[str] = mapped_column( - String(length=1500).with_variant( - String( - length=1500, - # latin1 allows for more indexed length in mysql - # and this field should only be ascii chars - collation="latin1_general_cs", - ), - "mysql", - ), + ASSET_STR_FIELD, nullable=False, ) uri: Mapped[str] = mapped_column( - String(length=1500).with_variant( - String( - length=1500, - # latin1 allows for more indexed length in mysql - # and this field should only be ascii chars - collation="latin1_general_cs", - ), - "mysql", - ), + ASSET_STR_FIELD, nullable=False, ) group: Mapped[str] = mapped_column( - String(length=1500).with_variant( - String( - length=1500, - # latin1 allows for more indexed length in mysql - # and this field should only be ascii chars - collation="latin1_general_cs", - ), - "mysql", - ), + ASSET_STR_FIELD, default=str, nullable=False, ) @@ -408,27 +359,11 @@ class AssetActive(Base): """ name: Mapped[str] = mapped_column( - String(length=1500).with_variant( - String( - length=1500, - # latin1 allows for more indexed length in mysql - # and this field should only be ascii chars - collation="latin1_general_cs", - ), - "mysql", - ), + ASSET_STR_FIELD, nullable=False, ) uri: Mapped[str] = mapped_column( - String(length=1500).with_variant( - String( - length=1500, - # latin1 allows for more indexed length in mysql - # and this field should only be ascii chars - collation="latin1_general_cs", - ), - "mysql", - ), + ASSET_STR_FIELD, nullable=False, ) @@ -456,15 +391,7 @@ class DagScheduleAssetNameReference(Base): """Reference from a DAG to an asset name reference of which it is a consumer.""" name: Mapped[str] = mapped_column( - String(length=1500).with_variant( - String( - length=1500, - # latin1 allows for more indexed length in mysql - # and this field should only be ascii chars - collation="latin1_general_cs", - ), - "mysql", - ), + ASSET_STR_FIELD, primary_key=True, nullable=False, ) @@ -502,15 +429,7 @@ class DagScheduleAssetUriReference(Base): """Reference from a DAG to an asset URI reference of which it is a consumer.""" uri: Mapped[str] = mapped_column( - String(length=1500).with_variant( - String( - length=1500, - # latin1 allows for more indexed length in mysql - # and this field should only be ascii chars - collation="latin1_general_cs", - ), - "mysql", - ), + ASSET_STR_FIELD, primary_key=True, nullable=False, ) diff --git a/airflow-core/src/airflow/models/base.py b/airflow-core/src/airflow/models/base.py index 1c7af7b275ab3..2840bf7e8f8a8 100644 --- a/airflow-core/src/airflow/models/base.py +++ b/airflow-core/src/airflow/models/base.py @@ -83,6 +83,23 @@ def get_id_collation_args(): COLLATION_ARGS: dict[str, Any] = get_id_collation_args() +def get_asset_str_field(length: int = 1500) -> String: + """ + Build the string type used for asset name/uri/group columns. + + On MySQL these carry an explicit latin1 collation: the values are ASCII, and + a 1-byte-per-character charset keeps the 1500-char unique indexes inside the + 3072-byte index limit that utf8mb4 would blow past. The collation is + overridable because MySQL-compatible engines do not all ship + ``latin1_general_cs`` (TiDB, for one, accepts only ``latin1_bin``). + """ + collation = conf.get("database", "sql_engine_collation_for_asset_names", fallback="latin1_general_cs") + return String(length=length).with_variant(String(length=length, collation=collation), "mysql") + + +ASSET_STR_FIELD: String = get_asset_str_field() + + def StringID(*, length=ID_LEN, **kwargs) -> String: return String(length=length, **kwargs, **COLLATION_ARGS)