This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch backport-xcom-migration-chain-v3-3-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit a4f25dc454ebc9c4b3e6bbc69ba4e80e9c9e9a20 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]> (cherry picked from commit 54f4300bcd03e2ab500c197993092119b6e8ddf0) --- ...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 95638963c05..b662290ef50 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
