This is an automated email from the ASF dual-hosted git repository.
uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 54f4300bcd0 Guard JSONB migration against invalid bytes conversion
(#69064)
54f4300bcd0 is described below
commit 54f4300bcd03e2ab500c197993092119b6e8ddf0
Author: Sean Muth <[email protected]>
AuthorDate: Wed Jul 1 19:40:12 2026 -0500
Guard JSONB migration against invalid bytes conversion (#69064)
Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>
---
...49_3_0_0_remove_pickled_data_from_xcom_table.py | 190 +++++++++++++--------
..._3_0_0_remove_pickled_data_from_dagrun_table.py | 37 +++-
...est_0049_remove_pickled_data_from_xcom_table.py | 150 ++++++++++++++++
...t_0055_remove_pickled_data_from_dagrun_table.py | 143 ++++++++++++++++
4 files changed, 450 insertions(+), 70 deletions(-)
diff --git
a/airflow-core/src/airflow/migrations/versions/0049_3_0_0_remove_pickled_data_from_xcom_table.py
b/airflow-core/src/airflow/migrations/versions/0049_3_0_0_remove_pickled_data_from_xcom_table.py
index 202a7aadbd7..367a68178bc 100644
---
a/airflow-core/src/airflow/migrations/versions/0049_3_0_0_remove_pickled_data_from_xcom_table.py
+++
b/airflow-core/src/airflow/migrations/versions/0049_3_0_0_remove_pickled_data_from_xcom_table.py
@@ -41,13 +41,122 @@ depends_on = None
airflow_version = "3.0.0"
+# --- Value-sanitization SQL, factored out so migration tests can run the real
statements
+# against an isolated table via the helpers below; ``table`` defaults to
"xcom" for the
+# production calls. Both classes of value that are legal in the pickled blob
but illegal in
+# strict JSON/JSONB are handled: non-finite floats (NaN/Infinity/-Infinity)
are quoted, and
+# the active U+0000 (NUL) escape is stripped (escaped backslashes are
protected first so a
+# literal U+0000 escape embedded in the data survives).
+_XCOM_PG_SANITIZE_SQL = r"""
+ UPDATE __TABLE__
+ SET value = convert_to(
+ regexp_replace(
+ -- Strip the active U+0000 (NUL) escape (illegal in
JSON/JSONB, not quotable).
+ -- Protect escaped backslashes with chr(1) first so an
embedded literal
+ -- backslash + u0000 in the data is preserved; chr(1)
is safe because
+ -- json.dumps escapes control bytes, so it never
appears in the JSON text.
+ replace(
+ replace(
+ replace(convert_from(value, 'UTF8'), '\\',
chr(1)),
+ '\u0000', ''
+ ),
+ chr(1), '\\'
+ ),
+ -- Group 1 captures the preceding delimiter (:, comma,
or [)
+ -- or ^ for a bare scalar value (the entire XCom value
is just NaN).
+ -- A lookahead is used for the closing delimiter
instead of a
+ -- consuming group so that consecutive tokens in an
array
+ -- (e.g. [NaN, Infinity]) are each matched
independently.
+ -- NaN and Infinity are done in the same query to
avoid another table scan.
+ '([:,\[]\s*|^)(NaN|-?Infinity)(?=\s*[,}\]]|$)',
+ '\1"\2"',
+ 'g'
+ ),
+ 'UTF8'
+ )
+ WHERE value IS NOT NULL AND get_byte(value, 0) != 128
+ """
+_XCOM_MYSQL_SANITIZE_SQL = """
+ UPDATE __TABLE__
+ SET value = CONVERT(
+ REGEXP_REPLACE(
+ -- Strip the active U+0000 (NUL) escape (illegal JSON;
see PostgreSQL branch).
+ -- Protect escaped backslashes with CHAR(1) first so
an embedded literal
+ -- backslash + u0000 in the data is preserved.
+ REPLACE(
+ REPLACE(
+ REPLACE(CONVERT(value USING utf8mb4),
'\\\\\\\\', CHAR(1)),
+ '\\\\u0000', ''
+ ),
+ CHAR(1), '\\\\\\\\'
+ ),
+ -- Same lookahead strategy as PostgreSQL (see above).
+ -- Python string escaping: \\\\[ → SQL \\[ → regex \\[
→ literal [
+ -- and \\\\] inside the character class → SQL \\] →
regex \\] → literal ]
+ -- The 'c' flag enforces case-sensitive matching (NaN
≠ nan).
+ -- NaN and Infinity are done in the same query to
avoid another table scan.
+ '(:|,|\\\\[|^)[ ]*(NaN|-?Infinity)(?=[ ]*[,}\\\\]]|$)',
+ '$1"$2"',
+ 1,
+ 0,
+ 'c'
+ ) USING BINARY
+ )
+ WHERE value IS NOT NULL AND HEX(SUBSTRING(value, 1, 1)) != '80'
+ """
+_XCOM_SQLITE_SANITIZE_SQL = """
+ UPDATE __TABLE__
+ SET value = CAST(
+ REPLACE(
+ REPLACE(
+ -- Step 1: replace NaN first so it doesn't
interfere with Infinity.
+ REPLACE(
+ -- Protect escaped backslashes with char(1),
strip the active
+ -- U+0000 (NUL) escape, then restore (see
PostgreSQL branch).
+ REPLACE(
+ REPLACE(
+ REPLACE(CAST(value AS TEXT), '\\\\',
char(1)),
+ '\\u0000', ''
+ ),
+ char(1), '\\\\'
+ ),
+ 'NaN', '"NaN"'
+ ),
+ -- Step 2: replace Infinity (also matches the
Infinity in -Infinity,
+ -- turning -Infinity into -"Infinity").
+ 'Infinity', '"Infinity"'
+ ),
+ -- Step 3: fix the -"Infinity" artifact left by step 2.
+ '-"Infinity"', '"-Infinity"'
+ ) AS BLOB)
+ -- NOTE: SQLite lacks REGEXP_REPLACE, so plain REPLACE is used.
+ -- This is a substring operation and will incorrectly alter
XCom values
+ -- that contain the literal text 'NaN' or 'Infinity' inside a
JSON string
+ -- (e.g. {"msg": "NaN detected"}). In practice such values
are rare and
+ -- SQLite is not recommended for production deployments.
+ WHERE value IS NOT NULL AND hex(substr(value, 1, 1)) != '80'
+ """
+
+
+def _xcom_pg_sanitize_sql(table: str = "xcom") -> str:
+ return _XCOM_PG_SANITIZE_SQL.replace("__TABLE__", table)
+
+
+def _xcom_mysql_sanitize_sql(table: str = "xcom") -> str:
+ return _XCOM_MYSQL_SANITIZE_SQL.replace("__TABLE__", table)
+
+
+def _xcom_sqlite_sanitize_sql(table: str = "xcom") -> str:
+ return _XCOM_SQLITE_SANITIZE_SQL.replace("__TABLE__", table)
+
+
def upgrade():
"""Apply Remove pickled data from xcom table."""
# Summary of the change:
# 1. Create an archived table (`_xcom_archive`) to store the current
"pickled" data in the xcom table
# 2. Extract and archive the pickled data using the condition
# 3. Delete the pickled data from the xcom table so that we can update the
column type
- # 4. Sanitize non-standard JSON tokens (NaN, Infinity, -Infinity) to
quoted strings
+ # 4. Sanitize values illegal in strict JSON/JSONB (quote NaN/Infinity,
strip the U+0000 NUL escape)
# 5. Update the XCom.value column type to JSON from LargeBinary/LongBlob
conn = op.get_bind()
@@ -112,31 +221,16 @@ def upgrade():
# Delete the pickled data from the xcom table so that we can update the
column type
conn.execute(text(f"DELETE FROM xcom WHERE value IS NOT NULL AND
{condition}"))
- # Sanitize non-standard JSON tokens (NaN, Infinity, -Infinity) to quoted
strings.
- # These are valid Python float representations but illegal in strict JSON;
they must
- # be quoted before the column type is changed to JSON/JSONB.
+ # Sanitize values that round-trip through pickle but are illegal in strict
JSON/JSONB
+ # before changing the column type:
+ # * NaN / Infinity / -Infinity -> quoted strings (valid Python floats,
illegal JSON).
+ # * the U+0000 (NUL) escape -> stripped. PostgreSQL JSON/JSONB cannot
represent it
+ # ("unsupported Unicode escape sequence ... cannot be converted to
text") and,
+ # unlike the non-finite floats, it cannot be quoted/kept, so it is
removed.
+ # json.dumps() emits a literal NUL as the 6-char escape, never a raw
0x00 byte, so
+ # stripping the escape covers values produced by normal XCom
serialization.
if dialect == "postgresql":
- conn.execute(
- text(r"""
- UPDATE xcom
- SET value = convert_to(
- regexp_replace(
- convert_from(value, 'UTF8'),
- -- Group 1 captures the preceding delimiter (:, comma,
or [)
- -- or ^ for a bare scalar value (the entire XCom value
is just NaN).
- -- A lookahead is used for the closing delimiter
instead of a
- -- consuming group so that consecutive tokens in an
array
- -- (e.g. [NaN, Infinity]) are each matched
independently.
- -- NaN and Infinity are done in the same query to
avoid another table scan.
- '([:,\[]\s*|^)(NaN|-?Infinity)(?=\s*[,}\]]|$)',
- '\1"\2"',
- 'g'
- ),
- 'UTF8'
- )
- WHERE value IS NOT NULL AND get_byte(value, 0) != 128
- """)
- )
+ conn.execute(text(_xcom_pg_sanitize_sql()))
op.execute(
"""
@@ -149,27 +243,7 @@ def upgrade():
"""
)
elif dialect == "mysql":
- conn.execute(
- text("""
- UPDATE xcom
- SET value = CONVERT(
- REGEXP_REPLACE(
- CONVERT(value USING utf8mb4),
- -- Same lookahead strategy as PostgreSQL (see above).
- -- Python string escaping: \\\\[ → SQL \\[ → regex \\[
→ literal [
- -- and \\\\] inside the character class → SQL \\] →
regex \\] → literal ]
- -- The 'c' flag enforces case-sensitive matching (NaN
≠ nan).
- -- NaN and Infinity are done in the same query to
avoid another table scan.
- '(:|,|\\\\[|^)[ ]*(NaN|-?Infinity)(?=[ ]*[,}\\\\]]|$)',
- '$1"$2"',
- 1,
- 0,
- 'c'
- ) USING BINARY
- )
- WHERE value IS NOT NULL AND HEX(SUBSTRING(value, 1, 1)) != '80'
- """)
- )
+ conn.execute(text(_xcom_mysql_sanitize_sql()))
op.add_column("xcom", sa.Column("value_json", sa.JSON(),
nullable=True))
op.execute("UPDATE xcom SET value_json = CAST(value AS CHAR CHARACTER
SET utf8mb4)")
@@ -177,29 +251,7 @@ def upgrade():
op.alter_column("xcom", "value_json", existing_type=sa.JSON(),
new_column_name="value")
elif dialect == "sqlite":
- conn.execute(
- text("""
- UPDATE xcom
- SET value = CAST(
- REPLACE(
- REPLACE(
- -- Step 1: replace NaN first so it doesn't
interfere with Infinity.
- REPLACE(CAST(value AS TEXT), 'NaN', '"NaN"'),
- -- Step 2: replace Infinity (also matches the
Infinity in -Infinity,
- -- turning -Infinity into -"Infinity").
- 'Infinity', '"Infinity"'
- ),
- -- Step 3: fix the -"Infinity" artifact left by step 2.
- '-"Infinity"', '"-Infinity"'
- ) AS BLOB)
- -- NOTE: SQLite lacks REGEXP_REPLACE, so plain REPLACE is used.
- -- This is a substring operation and will incorrectly alter
XCom values
- -- that contain the literal text 'NaN' or 'Infinity' inside a
JSON string
- -- (e.g. {"msg": "NaN detected"}). In practice such values
are rare and
- -- SQLite is not recommended for production deployments.
- WHERE value IS NOT NULL AND hex(substr(value, 1, 1)) != '80'
- """)
- )
+ conn.execute(text(_xcom_sqlite_sanitize_sql()))
# Rename the existing `value` column to `value_old`
with op.batch_alter_table("xcom", schema=None) as batch_op:
batch_op.alter_column("value", new_column_name="value_old")
diff --git
a/airflow-core/src/airflow/migrations/versions/0055_3_0_0_remove_pickled_data_from_dagrun_table.py
b/airflow-core/src/airflow/migrations/versions/0055_3_0_0_remove_pickled_data_from_dagrun_table.py
index 31b71e2bedc..f527dab936e 100644
---
a/airflow-core/src/airflow/migrations/versions/0055_3_0_0_remove_pickled_data_from_dagrun_table.py
+++
b/airflow-core/src/airflow/migrations/versions/0055_3_0_0_remove_pickled_data_from_dagrun_table.py
@@ -28,7 +28,9 @@ Create Date: 2024-12-01 08:33:15.425141
from __future__ import annotations
import json
+import math
import pickle
+from collections.abc import Mapping, Sequence
from textwrap import dedent
import sqlalchemy as sa
@@ -44,6 +46,37 @@ depends_on = None
airflow_version = "3.0.0"
+def _json_safe(obj):
+ """
+ Make a pickled conf value safe for strict JSON/JSONB before json.dumps.
+
+ Pickled ``conf`` can hold values that round-trip through pickle but are
illegal in
+ strict JSON/JSONB:
+
+ * non-finite floats (NaN / inf / -inf) -> quoted strings, mirroring the SQL
+ sanitization in migration 0049 (xcom);
+ * embedded U+0000 (NUL) characters in strings -> stripped, since PostgreSQL
+ JSON/JSONB cannot store them.
+
+ NUL is handled here, on the object before serialization, rather than on
the dumped
+ text: a blind string replace on the JSON output would also corrupt a
genuinely
+ escaped backslash sequence (an embedded literal backslash followed by
``u0000``).
+ """
+ if isinstance(obj, float):
+ if math.isnan(obj):
+ return "NaN"
+ if math.isinf(obj):
+ return "Infinity" if obj > 0 else "-Infinity"
+ return obj
+ if isinstance(obj, str):
+ return obj.replace(chr(0), "")
+ if isinstance(obj, Mapping):
+ return {_json_safe(k): _json_safe(v) for k, v in obj.items()}
+ if isinstance(obj, Sequence) and not isinstance(obj, (bytes, bytearray)):
+ return [_json_safe(v) for v in obj]
+ return obj
+
+
def upgrade():
"""Apply remove pickled data from dagrun table."""
conn = op.get_bind()
@@ -95,7 +128,9 @@ def upgrade():
try:
original_data = pickle.loads(pickle_data)
- json_data = json.dumps(original_data)
+ # _json_safe quotes non-finite floats and strips embedded
NUL chars so the
+ # row is preserved instead of dropped by the except below.
+ json_data = json.dumps(_json_safe(original_data))
conn.execute(
text("""
UPDATE dag_run
diff --git
a/airflow-core/tests/unit/migrations/test_0049_remove_pickled_data_from_xcom_table.py
b/airflow-core/tests/unit/migrations/test_0049_remove_pickled_data_from_xcom_table.py
new file mode 100644
index 00000000000..5d18c7f9a28
--- /dev/null
+++
b/airflow-core/tests/unit/migrations/test_0049_remove_pickled_data_from_xcom_table.py
@@ -0,0 +1,150 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Regression tests for migration 0049 (eed27faa34e3) value sanitization.
+
+The 2.x -> 3.x conversion of ``xcom.value`` from pickled bytea to JSON/JSONB
must not choke
+on values that are legal in the pickled blob but illegal in strict JSON/JSONB:
non-finite
+floats (NaN/Infinity/-Infinity) and the U+0000 (NUL) escape. It must also NOT
corrupt a
+genuinely escaped backslash sequence (a literal backslash-u-0000 in the data).
These tests
+run the migration's own per-dialect sanitization SQL against an isolated table.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+from pathlib import Path
+
+import pytest
+import sqlalchemy as sa
+
+from airflow import settings
+
+from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
+
+# A single backslash, built via chr() so no literal escape appears in the
source.
+_BS = chr(92)
+
+# Row 1: every value class the sanitizer must clean. chr(0) is a real embedded
null byte;
+# json.dumps serializes it to the 6-char NUL escape, which is what the
migration must strip.
+_RAW = json.dumps(
+ {"d": "F" + chr(0) + "oo", "a": float("nan"), "b": float("inf"), "c":
float("-inf"), "ok": 1.5}
+)
+_EXPECTED = {"d": "Foo", "a": "NaN", "b": "Infinity", "c": "-Infinity", "ok":
1.5}
+
+# Row 2: a string that literally contains backslash-u-0000 (no null byte). It
serializes to
+# an escaped backslash sequence (\\u0000) and MUST survive sanitization
unchanged.
+_LITERAL_VALUE = "x" + _BS + "u0000y"
+_LITERAL_RAW = json.dumps({"k": _LITERAL_VALUE})
+_LITERAL_EXPECTED = {"k": _LITERAL_VALUE}
+
+# Migration filenames start with a digit so they cannot be imported via the
normal import
+# system; load the module by file path instead.
+_MIGRATION_PATH = (
+ Path(AIRFLOW_CORE_SOURCES_PATH)
+ /
"airflow/migrations/versions/0049_3_0_0_remove_pickled_data_from_xcom_table.py"
+)
+_spec = importlib.util.spec_from_file_location("migration_0049",
_MIGRATION_PATH)
+_migration = importlib.util.module_from_spec(_spec) # type: ignore[arg-type]
+_spec.loader.exec_module(_migration) # type: ignore[union-attr]
+
+_TABLE = "_test_xcom_sanitize"
+
+
+def test_sqlite_sanitize_quotes_nonfinite_strips_nul_and_keeps_literal():
+ """SQLite branch: real sanitize SQL on an in-memory db.
Backend-independent."""
+ engine = sa.create_engine("sqlite://")
+ with engine.begin() as conn:
+ conn.execute(sa.text(f"CREATE TABLE {_TABLE} (id INTEGER PRIMARY KEY,
value BLOB)"))
+ conn.execute(
+ sa.text(f"INSERT INTO {_TABLE} (id, value) VALUES (1, :v)"),
+ {"v": _RAW.encode("utf-8")},
+ )
+ conn.execute(
+ sa.text(f"INSERT INTO {_TABLE} (id, value) VALUES (2, :v)"),
+ {"v": _LITERAL_RAW.encode("utf-8")},
+ )
+ conn.execute(sa.text(_migration._xcom_sqlite_sanitize_sql(_TABLE)))
+ # json(...) mirrors the migration's own conversion and raises if still
invalid JSON.
+ rows = dict(conn.execute(sa.text(f"SELECT id, json(CAST(value AS
TEXT)) FROM {_TABLE}")).all())
+ assert json.loads(rows[1]) == _EXPECTED
+ assert json.loads(rows[2]) == _LITERAL_EXPECTED
+
+
[email protected]_test
+class TestPostgresSanitize:
+ @pytest.mark.backend("postgres")
+ def test_nul_blocks_jsonb_cast_until_sanitized_and_literal_survives(self):
+ drop = f"DROP TABLE IF EXISTS {_TABLE}"
+ cast = f"SELECT CAST(CONVERT_FROM(value, 'UTF8') AS JSONB) FROM
{_TABLE}"
+ with settings.engine.begin() as conn:
+ conn.execute(sa.text(drop))
+ conn.execute(sa.text(f"CREATE TABLE {_TABLE} (id int PRIMARY KEY,
value bytea)"))
+ conn.execute(sa.text(f"INSERT INTO {_TABLE} VALUES (1,
convert_to(:v, 'UTF8'))"), {"v": _RAW})
+ conn.execute(
+ sa.text(f"INSERT INTO {_TABLE} VALUES (2, convert_to(:v,
'UTF8'))"),
+ {"v": _LITERAL_RAW},
+ )
+ try:
+ # Before sanitizing, the JSONB cast fails on the NUL escape (the
reported bug).
+ with settings.engine.connect() as conn:
+ with pytest.raises(sa.exc.DataError):
+ conn.execute(sa.text(cast)).all()
+ conn.rollback()
+ # After the migration's sanitize SQL, the cast succeeds and values
are correct.
+ with settings.engine.begin() as conn:
+ conn.execute(sa.text(_migration._xcom_pg_sanitize_sql(_TABLE)))
+ conn.execute(sa.text(cast)).all()
+ rows = dict(
+ conn.execute(sa.text(f"SELECT id, CONVERT_FROM(value,
'UTF8') FROM {_TABLE}")).all()
+ )
+ assert json.loads(rows[1]) == _EXPECTED
+ assert json.loads(rows[2]) == _LITERAL_EXPECTED
+ finally:
+ with settings.engine.begin() as conn:
+ conn.execute(sa.text(drop))
+
+
[email protected]_test
+class TestMysqlSanitize:
+ @pytest.mark.backend("mysql")
+ def test_sanitize_allows_json_cast_and_literal_survives(self):
+ drop = f"DROP TABLE IF EXISTS {_TABLE}"
+ cast = f"SELECT CAST(CONVERT(value USING utf8mb4) AS JSON) FROM
{_TABLE}"
+ with settings.engine.begin() as conn:
+ conn.execute(sa.text(drop))
+ conn.execute(sa.text(f"CREATE TABLE {_TABLE} (id int PRIMARY KEY,
value LONGBLOB)"))
+ conn.execute(sa.text(f"INSERT INTO {_TABLE} VALUES (1, CONVERT(:v
USING utf8mb4))"), {"v": _RAW})
+ conn.execute(
+ sa.text(f"INSERT INTO {_TABLE} VALUES (2, CONVERT(:v USING
utf8mb4))"),
+ {"v": _LITERAL_RAW},
+ )
+ try:
+ with settings.engine.begin() as conn:
+
conn.execute(sa.text(_migration._xcom_mysql_sanitize_sql(_TABLE)))
+ conn.execute(sa.text(cast)).all() # must not raise (bare NaN
would be rejected)
+ rows = dict(
+ conn.execute(sa.text(f"SELECT id, CONVERT(value USING
utf8mb4) FROM {_TABLE}")).all()
+ )
+ assert json.loads(rows[1]) == _EXPECTED
+ assert json.loads(rows[2]) == _LITERAL_EXPECTED
+ finally:
+ with settings.engine.begin() as conn:
+ conn.execute(sa.text(drop))
diff --git
a/airflow-core/tests/unit/migrations/test_0055_remove_pickled_data_from_dagrun_table.py
b/airflow-core/tests/unit/migrations/test_0055_remove_pickled_data_from_dagrun_table.py
new file mode 100644
index 00000000000..aeca8563a01
--- /dev/null
+++
b/airflow-core/tests/unit/migrations/test_0055_remove_pickled_data_from_dagrun_table.py
@@ -0,0 +1,143 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Unit tests for migration 0055 (e39a26ac59f6) conf sanitization.
+
+The 2.x -> 3.x conversion of ``dag_run.conf`` from pickled bytea to JSON/JSONB
happens
+Python-side (``json.dumps`` + a per-row insert). ``_json_safe`` quotes
non-finite floats
+and strips embedded NUL characters so confs carrying those values are
preserved instead of
+being dropped by the migration's per-row error handler. NUL is handled on the
object (not
+on the dumped text) so a genuinely escaped backslash sequence is not
corrupted. These are
+pure-Python tests; no database is required.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+from collections import OrderedDict
+from pathlib import Path
+
+import pytest
+
+from tests_common.test_utils.paths import AIRFLOW_CORE_SOURCES_PATH
+
+# A single backslash, built via chr() so no literal escape appears in the
source.
+_BS = chr(92)
+# The 6-char escape json.dumps emits for an embedded null byte.
+_NUL_ESCAPE = _BS + "u0000"
+
+_MIGRATION_PATH = (
+ Path(AIRFLOW_CORE_SOURCES_PATH)
+ /
"airflow/migrations/versions/0055_3_0_0_remove_pickled_data_from_dagrun_table.py"
+)
+_spec = importlib.util.spec_from_file_location("migration_0055",
_MIGRATION_PATH)
+_migration = importlib.util.module_from_spec(_spec) # type: ignore[arg-type]
+_spec.loader.exec_module(_migration) # type: ignore[union-attr]
+
+_json_safe = _migration._json_safe
+
+
[email protected](
+ ("value", "expected"),
+ [
+ (float("nan"), "NaN"),
+ (float("inf"), "Infinity"),
+ (float("-inf"), "-Infinity"),
+ (1.5, 1.5),
+ (0.0, 0.0),
+ (-2.0, -2.0),
+ ("plain", "plain"),
+ (42, 42),
+ (None, None),
+ (True, True),
+ ],
+)
+def test_json_safe_scalars(value, expected):
+ assert _json_safe(value) == expected
+
+
+def test_json_safe_strips_null_bytes_in_strings():
+ assert _json_safe("foo" + chr(0) + "bar") == "foobar"
+ assert _json_safe(chr(0)) == ""
+
+
+def test_json_safe_preserves_literal_backslash_u_text():
+ """A string literally containing backslash-u-0000 (no null byte) must
survive intact."""
+ literal = "foo" + _NUL_ESCAPE + "bar"
+ assert _json_safe(literal) == literal
+ # and round-trips through json without corruption
+ assert json.loads(json.dumps(_json_safe({"k": literal}))) == {"k": literal}
+
+
+def test_json_safe_recurses_into_mappings_and_sequences():
+ data = OrderedDict(
+ [
+ ("f", float("nan")),
+ ("lst", [float("inf"), 1, {"deep": float("-inf")}]),
+ ("tpl", (float("nan"), 2)),
+ ("nul" + chr(0), "v" + chr(0)),
+ ("keep", 3.14),
+ ]
+ )
+ assert _json_safe(data) == {
+ "f": "NaN",
+ "lst": ["Infinity", 1, {"deep": "-Infinity"}],
+ "tpl": ["NaN", 2], # tuples normalize to lists, like json.dumps would
+ "nul": "v", # NUL stripped from both key and value
+ "keep": 3.14,
+ }
+
+
+def test_json_safe_does_not_explode_strings_into_chars():
+ assert _json_safe("hello") == "hello"
+
+
+def _reject_constant(token):
+ raise AssertionError(f"non-finite token survived sanitization: {token!r}")
+
+
+def test_full_pipeline_yields_strict_valid_json():
+ """Mirror the migration's exact serialization:
json.dumps(_json_safe(...))."""
+ original = {
+ "d": "F" + chr(0) + "oo", # real embedded null byte
+ "lit": "x" + _NUL_ESCAPE + "y", # literal backslash-u-0000 text, must
survive
+ "a": float("nan"),
+ "b": float("inf"),
+ "c": float("-inf"),
+ "ok": 1.5,
+ }
+ json_data = json.dumps(_json_safe(original))
+
+ # parse_constant fires on any surviving bare NaN/Infinity/-Infinity token.
+ parsed = json.loads(json_data, parse_constant=_reject_constant)
+ assert parsed == {
+ "d": "Foo",
+ "lit": "x" + _NUL_ESCAPE + "y",
+ "a": "NaN",
+ "b": "Infinity",
+ "c": "-Infinity",
+ "ok": 1.5,
+ }
+
+
+def test_finite_floats_are_untouched():
+ original = {"x": 1.25, "y": [0.0, -3.5], "z": 1000000.0}
+ json_data = json.dumps(_json_safe(original))
+ assert json.loads(json_data, parse_constant=_reject_constant) == original
