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 8ef3202e64f0279754f9b7c8b02199f3ca78c521 Author: Hemkumar Chheda <[email protected]> AuthorDate: Thu Sep 10 14:05:11 2026 +0530 Do not corrupt XCom values that already parse as JSON in the bytea to JSONB migration (#71926) * Substitute null for non-finite floats in the XCom bytea to JSONB migration The sanitizer in 0049 wrapped NaN/Infinity/-Infinity in unescaped double quotes, which is only valid at the top level of a JSON document. When the XCom value is a JSON string wrapping another JSON document every interior quote is backslash-escaped, so the injected quote closes the wrapping string and the bytea -> JSONB cast aborts. null needs no quoting, so it is correct at any nesting depth. Also moves the run of spaces after the delimiter inside group 1 in the MySQL pattern, matching PostgreSQL, so the replacement stops dropping it. closes: #71921 * Skip the non-finite float rewrite for XCom values that already parse as JSON If a value parses as JSON, any NaN or Infinity in it can only be inside a string, so the rewrite has nothing to fix and the value is left byte for byte alone. That is the shape reported in #71921 and it means the migration no longer edits the contents of a JSON string at all. The sanitizer is now two statements per dialect. The U+0000 strip runs first and unconditionally, only on rows that contain the escape; the non-finite rewrite runs second, only on rows that do not parse. The order matters: a value invalid only because of the escape would otherwise look like it needs the rewrite and be corrupted by it. Rows that do need the rewrite still get bare null, since a value can be invalid at the top level and also wrap an escaped document, and only null is valid at every nesting depth. PostgreSQL has no non-throwing JSON validator before 16 and 14 is still supported, so the check goes through a session-local pg_temp function with a cheap regex pre-test. MySQL and SQLite use the native JSON_VALID, with a JSON1 probe and an unguarded fallback on SQLite as in migration 0117. The guard also shrinks the documented SQLite false positive, since a value holding "NaN detected" now parses and is skipped. closes: #71921 * Narrow the SQLite JSON1 probe to OperationalError in migration 0049 A missing json_valid() raises OperationalError. Catching every exception also swallowed connection failures and silently sanitized without the guard, rewriting rows that already parse. (cherry picked from commit 6b60399e52b1b571c2a0b2cf1c0b7d4581edcc6f) --- ...49_3_0_0_remove_pickled_data_from_xcom_table.py | 227 +++++++++++++-------- ..._3_0_0_remove_pickled_data_from_dagrun_table.py | 5 +- ...est_0049_remove_pickled_data_from_xcom_table.py | 169 +++++++++++---- 3 files changed, 277 insertions(+), 124 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 367a68178bc..fd9d77c4878 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 @@ -42,112 +42,180 @@ 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""" +# against an isolated table; ``table`` defaults to "xcom" for the production calls. +# +# Two things round-trip through pickle but are illegal in strict JSON/JSONB: +# +# 1. The U+0000 (NUL) escape. JSONB cannot represent it and it cannot be quoted, so it is +# stripped. Escaped backslashes are protected first so a literal ``\u0000`` in the data survives. +# 2. Non-finite floats (NaN / Infinity / -Infinity), rewritten to bare ``null``. +# +# Step 2 is skipped for values that already parse as JSON, where the token can only be inside a +# string: an XCom holding json.dumps'd JSON has its interior quotes escaped, so rewriting it +# would change data for no reason. Rows that do need it get ``null`` and not a quoted ``"NaN"``, +# because a raw quote is only valid at the top level. On a value that is invalid at the top level +# and also wraps an escaped document, the quote closes that string early and leaves the token +# bare, which aborts the cast. +# +# Step 1 runs first, or a value invalid only because of the escape looks like it needs step 2. +_XCOM_PG_STRIP_NUL_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( - replace(convert_from(value, 'UTF8'), '\\', chr(1)), - '\u0000', '' - ), - chr(1), '\\' + replace(convert_from(value, 'UTF8'), '\\', chr(1)), + '\u0000', '' ), - -- 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. + chr(1), '\\' + ), + 'UTF8' + ) + WHERE value IS NOT NULL AND get_byte(value, 0) != 128 + -- chr(1) never appears in the JSON text, json.dumps escapes control bytes. + -- The chain is an identity without the escape, so those rows are skipped + -- rather than rewritten. + AND position(convert_to('\u0000', 'UTF8') in value) > 0 + """ +# PostgreSQL has no non-throwing JSON validator before 16 (``pg_input_is_valid``) and 14 is still +# supported, so the guard goes through a session-local function. The regex test comes first so +# the cast is only attempted for rows that have a token at all. +_XCOM_PG_NEEDS_NAN_FIX_SQL = r""" + CREATE OR REPLACE FUNCTION pg_temp._airflow_xcom_needs_nan_fix(txt text) + RETURNS boolean AS $$ + BEGIN + IF txt !~ '(NaN|Infinity)' THEN + RETURN false; + END IF; + PERFORM txt::jsonb; + RETURN false; + EXCEPTION WHEN others THEN + RETURN true; + END; + $$ LANGUAGE plpgsql + """ +_XCOM_PG_SANITIZE_SQL = r""" + UPDATE __TABLE__ + SET value = convert_to( + regexp_replace( + convert_from(value, 'UTF8'), + -- Group 1 is the preceding delimiter, or ^ for a bare scalar value. The + -- closing delimiter is a lookahead rather than a consuming group so that + -- consecutive tokens in an array ([NaN, Infinity]) each match. NaN and + -- Infinity share one pass to avoid a second table scan. '([:,\[]\s*|^)(NaN|-?Infinity)(?=\s*[,}\]]|$)', - '\1"\2"', + '\1null', 'g' ), 'UTF8' ) WHERE value IS NOT NULL AND get_byte(value, 0) != 128 + AND pg_temp._airflow_xcom_needs_nan_fix(convert_from(value, 'UTF8')) """ -_XCOM_MYSQL_SANITIZE_SQL = """ +_XCOM_MYSQL_STRIP_NUL_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( - REPLACE(CONVERT(value USING utf8mb4), '\\\\\\\\', CHAR(1)), - '\\\\u0000', '' - ), - CHAR(1), '\\\\\\\\' + REPLACE(CONVERT(value USING utf8mb4), '\\\\\\\\', CHAR(1)), + '\\\\u0000', '' ), - -- 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"', + CHAR(1), '\\\\\\\\' + ) USING BINARY + ) + WHERE value IS NOT NULL AND HEX(SUBSTRING(value, 1, 1)) != '80' + AND LOCATE('\\\\u0000', value) > 0 + """ +_XCOM_MYSQL_SANITIZE_SQL = """ + UPDATE __TABLE__ + SET value = CONVERT( + REGEXP_REPLACE( + CONVERT(value USING utf8mb4), + -- Same grouping and lookahead as PostgreSQL. The run of spaces after the + -- delimiter is inside group 1 so the replacement puts it back; outside the + -- group it was dropped, which is invisible in a top-level document but not + -- when the document is itself a JSON string. 'c' forces case-sensitive + -- matching (NaN != nan). + -- Python escaping: \\\\[ -> SQL \\[ -> regex \\[ -> literal [ + '([:,\\\\[][ ]*|^)(NaN|-?Infinity)(?=[ ]*[,}\\\\]]|$)', + '$1null', 1, 0, 'c' ) USING BINARY ) WHERE value IS NOT NULL AND HEX(SUBSTRING(value, 1, 1)) != '80' + AND (LOCATE('NaN', value) > 0 OR LOCATE('Infinity', value) > 0) + AND NOT JSON_VALID(CONVERT(value USING utf8mb4)) + """ +_XCOM_SQLITE_STRIP_NUL_SQL = """ + UPDATE __TABLE__ + SET value = CAST( + REPLACE( + REPLACE( + REPLACE(CAST(value AS TEXT), '\\\\', char(1)), + '\\u0000', '' + ), + char(1), '\\\\' + ) AS BLOB) + WHERE value IS NOT NULL AND hex(substr(value, 1, 1)) != '80' + AND instr(CAST(value AS TEXT), '\\u0000') > 0 """ +# SQLite has no REGEXP_REPLACE, so this is a plain substring replace that cannot tell a token in +# a JSON syntax position from the same text inside a string. The json_valid() guard removes most +# of that risk: a value holding "NaN detected" still parses and is left alone. Only a value that +# fails to parse for some other reason can still be altered, and the result is valid JSON either +# way, so the migration completes. json_valid() needs the JSON1 extension, which SQLite builds +# before 3.38 may not have, so the guard is substituted in rather than inlined. _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"' + -- -Infinity first, or the bare Infinity step leaves '-null' behind. + REPLACE(CAST(value AS TEXT), '-Infinity', 'null'), + 'Infinity', 'null' ), - -- Step 3: fix the -"Infinity" artifact left by step 2. - '-"Infinity"', '"-Infinity"' + 'NaN', 'null' ) 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' + __GUARD__ """ +_SQLITE_JSON_VALID_GUARD = "AND NOT json_valid(CAST(value AS TEXT))" + + +def _xcom_pg_sanitize_statements(table: str = "xcom") -> list[str]: + return [ + _XCOM_PG_STRIP_NUL_SQL.replace("__TABLE__", table), + _XCOM_PG_NEEDS_NAN_FIX_SQL, + _XCOM_PG_SANITIZE_SQL.replace("__TABLE__", table), + ] -def _xcom_pg_sanitize_sql(table: str = "xcom") -> str: - return _XCOM_PG_SANITIZE_SQL.replace("__TABLE__", table) +def _xcom_mysql_sanitize_statements(table: str = "xcom") -> list[str]: + return [ + _XCOM_MYSQL_STRIP_NUL_SQL.replace("__TABLE__", table), + _XCOM_MYSQL_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_statements(table: str = "xcom", json1: bool = True) -> list[str]: + return [ + _XCOM_SQLITE_STRIP_NUL_SQL.replace("__TABLE__", table), + _XCOM_SQLITE_SANITIZE_SQL.replace("__TABLE__", table).replace( + "__GUARD__", _SQLITE_JSON_VALID_GUARD if json1 else "" + ), + ] -def _xcom_sqlite_sanitize_sql(table: str = "xcom") -> str: - return _XCOM_SQLITE_SANITIZE_SQL.replace("__TABLE__", table) +def _sqlite_has_json1(conn) -> bool: + """Whether this SQLite build provides json_valid() (the JSON1 extension).""" + try: + conn.execute(text("SELECT json_valid('{}')")).fetchone() + except sa.exc.OperationalError: + # Only a build without JSON1 lands here; a broader catch would hide a real connection error. + print("SQLite JSON functions unavailable; sanitizing without the json_valid() guard.") + return False + return True def upgrade(): @@ -156,7 +224,7 @@ def upgrade(): # 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 values illegal in strict JSON/JSONB (quote NaN/Infinity, strip the U+0000 NUL escape) + # 4. Sanitize values illegal in strict JSON/JSONB (strip the U+0000 NUL escape, null out NaN/Infinity) # 5. Update the XCom.value column type to JSON from LargeBinary/LongBlob conn = op.get_bind() @@ -221,16 +289,11 @@ 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 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. + # Sanitize values that are legal in the pickled blob but illegal in strict JSON/JSONB + # before changing the column type. See the statements at the top of this module. if dialect == "postgresql": - conn.execute(text(_xcom_pg_sanitize_sql())) + for stmt in _xcom_pg_sanitize_statements(): + conn.execute(text(stmt)) op.execute( """ @@ -243,7 +306,8 @@ def upgrade(): """ ) elif dialect == "mysql": - conn.execute(text(_xcom_mysql_sanitize_sql())) + for stmt in _xcom_mysql_sanitize_statements(): + conn.execute(text(stmt)) 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)") @@ -251,7 +315,8 @@ def upgrade(): op.alter_column("xcom", "value_json", existing_type=sa.JSON(), new_column_name="value") elif dialect == "sqlite": - conn.execute(text(_xcom_sqlite_sanitize_sql())) + for stmt in _xcom_sqlite_sanitize_statements(json1=_sqlite_has_json1(conn)): + conn.execute(text(stmt)) # 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 b662290ef50..c9004a0a1c0 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 @@ -53,8 +53,9 @@ def _json_safe(obj): 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); + * non-finite floats (NaN / inf / -inf) -> quoted strings. This runs on the deserialized + object, so json.dumps handles the escaping at any depth. Migration 0049 (xcom) rewrites + serialized text with a regex and cannot, so it substitutes ``null`` there; * embedded U+0000 (NUL) characters in strings -> stripped, since PostgreSQL JSON/JSONB cannot store them. 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 index 5d18c7f9a28..66fbd25e31a 100644 --- 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 @@ -19,11 +19,12 @@ """ 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. +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 leave a value that already parses +as JSON untouched, including one that wraps another JSON document with its interior quotes +escaped, and it must not corrupt a literal backslash-u-0000 in the data. These tests run the +migration's own per-dialect SQL against an isolated table. """ from __future__ import annotations @@ -47,14 +48,36 @@ _BS = chr(92) _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} +_EXPECTED = {"d": "Foo", "a": None, "b": None, "c": None, "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. +# Row 2: a string that literally contains backslash-u-0000 (no null byte). It serializes to an +# escaped backslash sequence and must survive unchanged. _LITERAL_VALUE = "x" + _BS + "u0000y" _LITERAL_RAW = json.dumps({"k": _LITERAL_VALUE}) _LITERAL_EXPECTED = {"k": _LITERAL_VALUE} +# Rows 3 and 4: a task pushed already-serialized JSON, so the value is a JSON string wrapping +# another document with its interior quotes escaped. Both already parse, so the non-finite +# rewrite must skip them: quoting the token would close the wrapping string early and abort the +# cast, and nulling it would rewrite data the migration has no reason to touch. +_INNER = json.dumps({"amount": 604441.0, "commission": float("nan"), "rate": float("-inf")}) +_ESCAPED_RAW = json.dumps(_INNER) +_NESTED_RAW = json.dumps({"report": _INNER}) + +# Row 5: invalid at the top level and wrapping an escaped document. The rewrite has to run, so +# the inner tokens go too; ``null`` keeps the result parseable where a quote would not. +_MIXED_RAW = json.dumps({"top": float("nan"), "report": _INNER}) +_MIXED_EXPECTED = { + "top": None, + "report": json.dumps({"amount": 604441.0, "commission": None, "rate": None}), +} + +# Row 6: invalid only because of the NUL escape, and wrapping an escaped document. Proves the +# strip runs before the rewrite: stripping first makes the value parse, so the inner NaN is +# preserved. Guarding before stripping would send this row through the rewrite and corrupt it. +_ORDER_RAW = json.dumps({"n": "x" + chr(0) + "y", "report": _INNER}) +_ORDER_EXPECTED = {"n": "xy", "report": _INNER} + # 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 = ( @@ -67,56 +90,120 @@ _spec.loader.exec_module(_migration) # type: ignore[union-attr] _TABLE = "_test_xcom_sanitize" +# id -> serialized value, inserted into the isolated table by every dialect test. +_ROWS = { + 1: _RAW, + 2: _LITERAL_RAW, + 3: _ESCAPED_RAW, + 4: _NESTED_RAW, + 5: _MIXED_RAW, + 6: _ORDER_RAW, +} + + +def _assert_sanitized(rows: dict[int, str]) -> None: + """Check the sanitized text of every row. Rows 3, 4 and 6 compare the inner document as an + exact string, so dropped whitespace inside it fails here rather than passing a loads() check. + """ + assert json.loads(rows[1]) == _EXPECTED + assert json.loads(rows[2]) == _LITERAL_EXPECTED + assert json.loads(rows[3]) == _INNER + assert json.loads(rows[4]) == {"report": _INNER} + assert json.loads(rows[5]) == _MIXED_EXPECTED + assert json.loads(rows[6]) == _ORDER_EXPECTED + # Rows 3 and 4 already parsed, so they must be byte-identical to what was stored. + assert rows[3] == _ESCAPED_RAW + assert rows[4] == _NESTED_RAW -def test_sqlite_sanitize_quotes_nonfinite_strips_nul_and_keeps_literal(): - """SQLite branch: real sanitize SQL on an in-memory db. Backend-independent.""" + +def _sqlite_sanitized(json1: bool = True) -> dict[int, str]: 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))) + for row_id, value in _ROWS.items(): + conn.execute( + sa.text(f"INSERT INTO {_TABLE} (id, value) VALUES (:i, :v)"), + {"i": row_id, "v": value.encode("utf-8")}, + ) + for stmt in _migration._xcom_sqlite_sanitize_statements(_TABLE, json1=json1): + conn.execute(sa.text(stmt)) # 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()) + # It also re-serializes, so the assertions read the stored text instead. + rows = conn.execute( + sa.text(f"SELECT id, CAST(value AS TEXT), json(CAST(value AS TEXT)) FROM {_TABLE}") + ).all() + return {row[0]: row[1] for row in rows} + + +def test_sqlite_sanitize(): + """SQLite branch: real sanitize SQL on an in-memory db. Backend-independent.""" + _assert_sanitized(_sqlite_sanitized()) + + +def test_sqlite_sanitize_without_json1(): + """Without JSON1 there is no guard, so already-valid values are rewritten too. The result + still has to be valid JSON, which is what keeps the migration completing on old builds. + """ + rows = _sqlite_sanitized(json1=False) assert json.loads(rows[1]) == _EXPECTED - assert json.loads(rows[2]) == _LITERAL_EXPECTED + assert json.loads(rows[3]) == _MIXED_EXPECTED["report"] + + +def test_sqlite_has_json1_probe(): + """The probe reports True on a build with JSON1 and swallows only the missing-function error.""" + engine = sa.create_engine("sqlite://") + with engine.connect() as conn: + assert _migration._sqlite_has_json1(conn) is True + + class _NoJson1: + def execute(self, *args, **kwargs): + raise sa.exc.OperationalError( + "SELECT json_valid('{}')", {}, Exception("no such function: json_valid") + ) + + assert _migration._sqlite_has_json1(_NoJson1()) is False + + +def test_sqlite_has_json1_probe_propagates_other_errors(): + """A failure that is not a missing function must surface instead of downgrading the sanitize.""" + + class _Broken: + def execute(self, *args, **kwargs): + raise sa.exc.InterfaceError("SELECT json_valid('{}')", {}, Exception("connection gone")) + + with pytest.raises(sa.exc.InterfaceError): + _migration._sqlite_has_json1(_Broken()) @pytest.mark.db_test class TestPostgresSanitize: @pytest.mark.backend("postgres") - def test_nul_blocks_jsonb_cast_until_sanitized_and_literal_survives(self): + def test_nul_and_nan_block_jsonb_cast_until_sanitized(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}, - ) + for row_id, value in _ROWS.items(): + conn.execute( + sa.text(f"INSERT INTO {_TABLE} VALUES (:i, convert_to(:v, 'UTF8'))"), + {"i": row_id, "v": value}, + ) try: - # Before sanitizing, the JSONB cast fails on the NUL escape (the reported bug). + # Before sanitizing, the JSONB cast fails (the reported upgrade failure). 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. + # pg_temp is per-session, so the helper and the UPDATE share one connection. with settings.engine.begin() as conn: - conn.execute(sa.text(_migration._xcom_pg_sanitize_sql(_TABLE))) + for stmt in _migration._xcom_pg_sanitize_statements(_TABLE): + conn.execute(sa.text(stmt)) 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 + _assert_sanitized(rows) finally: with settings.engine.begin() as conn: conn.execute(sa.text(drop)) @@ -125,26 +212,26 @@ class TestPostgresSanitize: @pytest.mark.db_test class TestMysqlSanitize: @pytest.mark.backend("mysql") - def test_sanitize_allows_json_cast_and_literal_survives(self): + def test_sanitize_allows_json_cast(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}, - ) + for row_id, value in _ROWS.items(): + conn.execute( + sa.text(f"INSERT INTO {_TABLE} VALUES (:i, CONVERT(:v USING utf8mb4))"), + {"i": row_id, "v": value}, + ) try: with settings.engine.begin() as conn: - conn.execute(sa.text(_migration._xcom_mysql_sanitize_sql(_TABLE))) + for stmt in _migration._xcom_mysql_sanitize_statements(_TABLE): + conn.execute(sa.text(stmt)) 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 + _assert_sanitized(rows) finally: with settings.engine.begin() as conn: conn.execute(sa.text(drop))
