Migration fails on XCom table due to invalid JSONB conversion with NaN values

[IrakliJani]

Apache Airflow version

3.0.1

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

No response

What happened?

During a database migration (from 2.10.2 to 3.0.1) the migration fails when attempting to cast the value column in the xcom table to JSONB. The root cause is the presence of NaN values in pickled data, which is not valid JSON and cannot be cast to JSONB.

What you think should happen instead?

Instead of failing the migration, the process should gracefully handle non-migratable pickled values like NaN. These could be replaced with a placeholder (e.g., null or a specific string like "UNSUPPORTED_VALUE(NaN)") to indicate non-migratable content. This would allow the migration to proceed without data loss, while still flagging problematic entries for further inspection if needed.

How to reproduce

Have a XCOM value that has NaN in it in a dict.

Operating System

docker / slim

Versions of Apache Airflow Providers

No response

Deployment

Official Apache Airflow Helm Chart

Deployment details

docker-compose, nothing special, the same

Anything else?

stack trace:

Traceback (most recent call last):
  File "/home/airflow/.local/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1910, in _execute_context
    self.dialect.do_execute(
  File "/home/airflow/.local/lib/python3.12/site-packages/sqlalchemy/engine/default.py", line 736, in do_execute
    cursor.execute(statement, parameters)
psycopg2.errors.InvalidTextRepresentation: invalid input syntax for type json
DETAIL:  Token "NaN" is invalid.
CONTEXT:  JSON data, line 1: [{"day": "2024-06-07", "ArticleCountMetric": NaN...

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

Traceback (most recent call last):
  File "/home/airflow/.local/bin/airflow", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/home/airflow/.local/lib/python3.12/site-packages/airflow/__main__.py", line 55, in main
    args.func(args)
  File "/home/airflow/.local/lib/python3.12/site-packages/airflow/cli/cli_config.py", line 48, in command
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/airflow/.local/lib/python3.12/site-packages/airflow/utils/cli.py", line 111, in wrapper
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/home/airflow/.local/lib/python3.12/site-packages/airflow/utils/providers_configuration_loader.py", line 55, in wrapped_function
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/airflow/.local/lib/python3.12/site-packages/airflow/cli/commands/db_command.py", line 197, in migratedb
    run_db_migrate_command(args, db.upgradedb, _REVISION_HEADS_MAP)
  File "/home/airflow/.local/lib/python3.12/site-packages/airflow/cli/commands/db_command.py", line 125, in run_db_migrate_command
    command(
  File "/home/airflow/.local/lib/python3.12/site-packages/airflow/utils/session.py", line 101, in wrapper
    return func(*args, session=session, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/airflow/.local/lib/python3.12/site-packages/airflow/utils/db.py", line 1142, in upgradedb
    command.upgrade(config, revision=to_revision or "heads")
  File "/home/airflow/.local/lib/python3.12/site-packages/alembic/command.py", line 408, in upgrade
    script.run_env()
  File "/home/airflow/.local/lib/python3.12/site-packages/alembic/script/base.py", line 586, in run_env
    util.load_python_file(self.dir, "env.py")
  File "/home/airflow/.local/lib/python3.12/site-packages/alembic/util/pyfiles.py", line 95, in load_python_file
    module = load_module_py(module_id, path)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/airflow/.local/lib/python3.12/site-packages/alembic/util/pyfiles.py", line 113, in load_module_py
    spec.loader.exec_module(module)  # type: ignore
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<frozen importlib._bootstrap_external>", line 999, in exec_module
  File "<frozen importlib._bootstrap>", line 488, in _call_with_frames_removed
  File "/home/airflow/.local/lib/python3.12/site-packages/airflow/migrations/env.py", line 138, in <module>
    run_migrations_online()
  File "/home/airflow/.local/lib/python3.12/site-packages/airflow/migrations/env.py", line 132, in run_migrations_online
    context.run_migrations()
  File "<string>", line 8, in run_migrations
  File "/home/airflow/.local/lib/python3.12/site-packages/alembic/runtime/environment.py", line 946, in run_migrations
    self.get_context().run_migrations(**kw)
  File "/home/airflow/.local/lib/python3.12/site-packages/alembic/runtime/migration.py", line 623, in run_migrations
    step.migration_fn(**kw)
  File "/home/airflow/.local/lib/python3.12/site-packages/airflow/migrations/versions/0049_3_0_0_remove_pickled_data_from_xcom_table.py", line 101, in upgrade
    op.execute(
  File "<string>", line 8, in execute
  File "<string>", line 3, in execute
  File "/home/airflow/.local/lib/python3.12/site-packages/alembic/operations/ops.py", line 2552, in execute
    return operations.invoke(op)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/airflow/.local/lib/python3.12/site-packages/alembic/operations/base.py", line 441, in invoke
    return fn(self, operation)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/airflow/.local/lib/python3.12/site-packages/alembic/operations/toimpl.py", line 223, in execute_sql
    operations.migration_context.impl.execute(
  File "/home/airflow/.local/lib/python3.12/site-packages/alembic/ddl/impl.py", line 253, in execute
    self._exec(sql, execution_options)
  File "/home/airflow/.local/lib/python3.12/site-packages/alembic/ddl/impl.py", line 246, in _exec
    return conn.execute(construct, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/airflow/.local/lib/python3.12/site-packages/sqlalchemy/future/engine.py", line 286, in execute
    return self._execute_20(
           ^^^^^^^^^^^^^^^^^
  File "/home/airflow/.local/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1710, in _execute_20
    return meth(self, args_10style, kwargs_10style, execution_options)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/airflow/.local/lib/python3.12/site-packages/sqlalchemy/sql/elements.py", line 334, in _execute_on_connection
    return connection._execute_clauseelement(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/airflow/.local/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1577, in _execute_clauseelement
    ret = self._execute_context(
          ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/airflow/.local/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1953, in _execute_context
    self._handle_dbapi_exception(
  File "/home/airflow/.local/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 2134, in _handle_dbapi_exception
    util.raise_(
  File "/home/airflow/.local/lib/python3.12/site-packages/sqlalchemy/util/compat.py", line 211, in raise_
    raise exception
  File "/home/airflow/.local/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1910, in _execute_context
    self.dialect.do_execute(
  File "/home/airflow/.local/lib/python3.12/site-packages/sqlalchemy/engine/default.py", line 736, in do_execute
    cursor.execute(statement, parameters)
sqlalchemy.exc.DataError: (psycopg2.errors.InvalidTextRepresentation) invalid input syntax for type json
DETAIL:  Token "NaN" is invalid.
CONTEXT:  JSON data, line 1: [{"day": "2024-06-07", "ArticleCountMetric": NaN...

[SQL:
            ALTER TABLE xcom
            ALTER COLUMN value TYPE JSONB
            USING CASE
                WHEN value IS NOT NULL THEN CAST(CONVERT_FROM(value, 'UTF8') AS JSONB)
                ELSE NULL
            END
            ]
(Background on this error at: https://sqlalche.me/e/14/9h9h)
Sentry is attempting to send 2 pending events
Waiting up to 2 seconds
Press Ctrl-C to quit

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

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.