This is an automated email from the ASF dual-hosted git repository.
vatsrahul1001 pushed a commit to branch v3-3-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-3-test by this push:
new 1ab3d9fe9f8 Do not corrupt XCom values that already parse as JSON in
the bytea to JSONB migration (#71926) (#72886)
1ab3d9fe9f8 is described below
commit 1ab3d9fe9f8d7851e79b7e19acc9dd6837ac405a
Author: Rahul Vats <[email protected]>
AuthorDate: Thu Sep 10 21:00:14 2026 +0530
Do not corrupt XCom values that already parse as JSON in the bytea to JSONB
migration (#71926) (#72886)
* Guard JSONB migration against invalid bytes conversion (#69064)
Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>
(cherry picked from commit 54f4300bcd03e2ab500c197993092119b6e8ddf0)
* 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)
---------
Co-authored-by: Sean Muth <[email protected]>
Co-authored-by: Claude Opus 4.8 (1M context) <[email protected]>
Co-authored-by: Hemkumar Chheda <[email protected]>
---
...49_3_0_0_remove_pickled_data_from_xcom_table.py | 255 +++++++++++++++------
..._3_0_0_remove_pickled_data_from_dagrun_table.py | 38 ++-
...est_0049_remove_pickled_data_from_xcom_table.py | 237 +++++++++++++++++++
...t_0055_remove_pickled_data_from_dagrun_table.py | 143 ++++++++++++
4 files changed, 603 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..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
@@ -41,13 +41,190 @@ 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; ``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(
+ replace(
+ replace(
+ replace(convert_from(value, 'UTF8'), '\\', chr(1)),
+ '\u0000', ''
+ ),
+ 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*[,}\]]|$)',
+ '\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_STRIP_NUL_SQL = """
+ UPDATE __TABLE__
+ SET value = CONVERT(
+ REPLACE(
+ REPLACE(
+ REPLACE(CONVERT(value USING utf8mb4), '\\\\\\\\',
CHAR(1)),
+ '\\\\u0000', ''
+ ),
+ 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(
+ -- -Infinity first, or the bare Infinity step
leaves '-null' behind.
+ REPLACE(CAST(value AS TEXT), '-Infinity', 'null'),
+ 'Infinity', 'null'
+ ),
+ 'NaN', 'null'
+ ) AS BLOB)
+ 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_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_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 _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():
"""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 (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()
@@ -112,31 +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 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 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(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
- """)
- )
+ for stmt in _xcom_pg_sanitize_statements():
+ conn.execute(text(stmt))
op.execute(
"""
@@ -149,27 +306,8 @@ 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'
- """)
- )
+ 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)")
@@ -177,29 +315,8 @@ 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'
- """)
- )
+ 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 95638963c05..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
@@ -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,38 @@ 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. 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.
+
+ 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 +129,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..66fbd25e31a
--- /dev/null
+++
b/airflow-core/tests/unit/migrations/test_0049_remove_pickled_data_from_xcom_table.py
@@ -0,0 +1,237 @@
+#
+# 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 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
+
+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": 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 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 = (
+ 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"
+
+# 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 _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)"))
+ 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.
+ # 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[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())
+
+
[email protected]_test
+class TestPostgresSanitize:
+ @pytest.mark.backend("postgres")
+ 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)"))
+ 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 (the reported upgrade
failure).
+ with settings.engine.connect() as conn:
+ with pytest.raises(sa.exc.DataError):
+ conn.execute(sa.text(cast)).all()
+ conn.rollback()
+ # pg_temp is per-session, so the helper and the UPDATE share one
connection.
+ with settings.engine.begin() as conn:
+ 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_sanitized(rows)
+ 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(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)"))
+ 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:
+ 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_sanitized(rows)
+ 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
