0101 Migration upgrade fails on MySQL with "Out of sort memory"

[vatsrahul1001]

Apache Airflow version

main (development)

If "Other Airflow 3 version" selected, which one?

No response

What happened?

Migration 0101_3_2_0_ui_improvements_for_deadlines fails on MySQL

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/python/bin/airflow", line 10, in <module>
    sys.exit(main())
  File "/opt/airflow/airflow-core/src/airflow/__main__.py", line 55, in main
    args.func(args)
  File "/opt/airflow/airflow-core/src/airflow/cli/cli_config.py", line 49, in command
    return func(*args, **kwargs)
  File "/opt/airflow/airflow-core/src/airflow/utils/cli.py", line 113, in wrapper
    return f(*args, **kwargs)
  File "/opt/airflow/airflow-core/src/airflow/utils/providers_configuration_loader.py", line 54, in wrapped_function
    return func(*args, **kwargs)
  File "/opt/airflow/airflow-core/src/airflow/cli/commands/db_command.py", line 207, in migratedb
    run_db_migrate_command(args, db.upgradedb, _REVISION_HEADS_MAP)
  File "/opt/airflow/airflow-core/src/airflow/cli/commands/db_command.py", line 134, in run_db_migrate_command
    command(
  File "/opt/airflow/airflow-core/src/airflow/utils/session.py", line 100, in wrapper
    return func(*args, session=session, **kwargs)  # type: ignore[arg-type]
  File "/opt/airflow/airflow-core/src/airflow/utils/db.py", line 1248, in upgradedb
    _run_upgradedb(config, to_revision, session)
  File "/opt/airflow/airflow-core/src/airflow/utils/db.py", line 1171, in _run_upgradedb
    command.upgrade(config, revision=to_revision or "heads")
  File "/usr/python/lib/python3.10/site-packages/alembic/command.py", line 483, in upgrade
    script.run_env()
  File "/usr/python/lib/python3.10/site-packages/alembic/script/base.py", line 545, in run_env
    util.load_python_file(self.dir, "env.py")
  File "/usr/python/lib/python3.10/site-packages/alembic/util/pyfiles.py", line 116, in load_python_file
    module = load_module_py(module_id, path)
  File "/usr/python/lib/python3.10/site-packages/alembic/util/pyfiles.py", line 136, in load_module_py
    spec.loader.exec_module(module)  # type: ignore
  File "<frozen importlib._bootstrap_external>", line 883, in exec_module
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "/opt/airflow/airflow-core/src/airflow/migrations/env.py", line 138, in <module>
    run_migrations_online()
  File "/opt/airflow/airflow-core/src/airflow/migrations/env.py", line 132, in run_migrations_online
    context.run_migrations()
  File "<string>", line 8, in run_migrations
  File "/usr/python/lib/python3.10/site-packages/alembic/runtime/environment.py", line 969, in run_migrations
    self.get_context().run_migrations(**kw)
  File "/usr/python/lib/python3.10/site-packages/alembic/runtime/migration.py", line 626, in run_migrations
    step.migration_fn(**kw)
  File "/opt/airflow/airflow-core/src/airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py", line 145, in upgrade
    migrate_existing_deadline_alert_data_from_serialized_dag()
  File "/opt/airflow/airflow-core/src/airflow/migrations/versions/0101_3_2_0_ui_improvements_for_deadlines.py", line 391, in migrate_existing_deadline_alert_data_from_serialized_dag
    result = conn.execute(
  File "/usr/python/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1419, in execute
    return meth(
  File "/usr/python/lib/python3.10/site-packages/sqlalchemy/sql/elements.py", line 527, in _execute_on_connection
    return connection._execute_clauseelement(
  File "/usr/python/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1641, in _execute_clauseelement
    ret = self._execute_context(
  File "/usr/python/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1846, in _execute_context
    return self._exec_single_context(
  File "/usr/python/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1986, in _exec_single_context
    self._handle_dbapi_exception(
  File "/usr/python/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 2363, in _handle_dbapi_exception
    raise sqlalchemy_exception.with_traceback(exc_info[2]) from e
  File "/usr/python/lib/python3.10/site-packages/sqlalchemy/engine/base.py", line 1967, in _exec_single_context
    self.dialect.do_execute(
  File "/usr/python/lib/python3.10/site-packages/sqlalchemy/engine/default.py", line 952, in do_execute
    cursor.execute(statement, parameters)
  File "/usr/python/lib/python3.10/site-packages/MySQLdb/cursors.py", line 179, in execute
    res = self._query(mogrified_query)
  File "/usr/python/lib/python3.10/site-packages/MySQLdb/cursors.py", line 331, in _query
    self._do_get_result(db)
  File "/usr/python/lib/python3.10/site-packages/MySQLdb/cursors.py", line 136, in _do_get_result
    self._result = result = self._get_result()
  File "/usr/python/lib/python3.10/site-packages/MySQLdb/cursors.py", line 363, in _get_result
    return self._get_db().store_result()
sqlalchemy.exc.OperationalError: (MySQLdb.OperationalError) (1038, 'Out of sort memory, consider increasing server sort buffer size')
[SQL:
                    SELECT sd.id, sd.dag_id, sd.data, sd.data_compressed, sd.created_at
                    FROM serialized_dag sd
                    INNER JOIN (
                        SELECT id, dag_id
                        FROM serialized_dag
                        WHERE (data IS NOT NULL OR data_compressed IS NOT NULL)
                          AND dag_id > %s
                        ORDER BY dag_id
                        LIMIT %s
                    ) AS subq ON sd.id = subq.id
                    ORDER BY sd.dag_id
                ]
[parameters: ('', 10000)]
(Background on this error at: https://sqlalche.me/e/20/e3q8)

I had to increase sort_buffer_size default size from default 256KB to 256MB as workaround

What you think should happen instead?

No response

How to reproduce

1. Start Airflow on MySQL at version 3.1.8
2. Run airflow db migrate to upgrade to 3.2.0
3. Migration 0101 fails with sort buffer error

Operating System

OS

Versions of Apache Airflow Providers

No response

Deployment

Official Apache Airflow Helm Chart

Deployment details

No response

Anything else?

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[wjddn279]

@vatsrahul1001

I have resolved the same issue before. #55589

The root cause is that when large-sized data is involved in the query, the sort_buffer becomes full. We confirmed that the "Out of sort memory" error occurs not only when selecting the data, but also simply by including it in the WHERE clause.

I changed the query so that the large data column is not included in the query, and the final sorting is performed in Python. Confirmed that the error no longer occurs.

I tested this and confirmed that this PR, which adds an index, still produces the error.

[vatsrahul1001]

Thanks @wjddn279 PR looks good to me