seanmuth commented on code in PR #69064:
URL: https://github.com/apache/airflow/pull/69064#discussion_r3492237653
##########
airflow-core/src/airflow/migrations/versions/0055_3_0_0_remove_pickled_data_from_dagrun_table.py:
##########
@@ -44,6 +45,28 @@
airflow_version = "3.0.0"
+def _json_safe(obj):
+ """Replace non-finite floats with their quoted string form before
json.dumps.
+
+ Pickled ``conf`` may contain ``float('nan')`` / ``inf`` / ``-inf``.
``json.dumps``
+ (with the default ``allow_nan=True``) emits the bare tokens ``NaN`` /
``Infinity`` /
+ ``-Infinity``, which PostgreSQL JSON/JSONB rejects. Quoting them to
strings mirrors
+ the SQL sanitization in migration 0049 (xcom) so the value is preserved
rather than
+ dropped by the per-row error handler below.
+ """
+ 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, dict):
+ return {k: _json_safe(v) for k, v in obj.items()}
+ if isinstance(obj, (list, tuple)):
+ return [_json_safe(v) for v in obj]
Review Comment:
Done — switched to `collections.abc.Mapping`/`Sequence` (excluding
`str`/`bytes` so strings aren't walked char-by-char), and added an explicit
`str` branch. `_json_safe` now also recurses into mapping keys. See the
follow-up commit.
*— posted via Claude Code (Claude Opus 4.8)*
##########
airflow-core/src/airflow/migrations/versions/0055_3_0_0_remove_pickled_data_from_dagrun_table.py:
##########
@@ -95,7 +118,12 @@ def upgrade():
try:
original_data = pickle.loads(pickle_data)
- json_data = json.dumps(original_data)
+ # Sanitize values legal in pickle but illegal in strict
JSON/JSONB so the
+ # row is preserved instead of dropped by the except below:
+ # * NaN / Infinity / -Infinity -> quoted strings (see
_json_safe).
+ # * the U+0000 (NUL) escape that json.dumps emits for
null bytes -> stripped
+ # (PostgreSQL JSON/JSONB cannot store it and it cannot
be quoted/kept).
+ json_data =
json.dumps(_json_safe(original_data)).replace('\\u0000', '')
Review Comment:
Good catch — confirmed: `json.dumps({"k": "foo\\u0000bar"})` then a blind
`.replace` of the NUL escape corrupts the value to `foo\bar`. Fixed by handling
NUL on the object **before** serialization: `_json_safe` now strips real null
bytes (`chr(0)`) from string values, and the post-dump `.replace` is gone. That
can't false-match a genuinely escaped sequence.
*— posted via Claude Code (Claude Opus 4.8)*
##########
airflow-core/src/airflow/migrations/versions/0055_3_0_0_remove_pickled_data_from_dagrun_table.py:
##########
@@ -95,7 +118,12 @@ def upgrade():
try:
original_data = pickle.loads(pickle_data)
- json_data = json.dumps(original_data)
+ # Sanitize values legal in pickle but illegal in strict
JSON/JSONB so the
+ # row is preserved instead of dropped by the except below:
+ # * NaN / Infinity / -Infinity -> quoted strings (see
_json_safe).
+ # * the U+0000 (NUL) escape that json.dumps emits for
null bytes -> stripped
+ # (PostgreSQL JSON/JSONB cannot store it and it cannot
be quoted/kept).
+ json_data =
json.dumps(_json_safe(original_data)).replace('\\u0000', '')
Review Comment:
Right — the SQL `replace(…, U+0000, '')` had the identical false-match.
Fixed in all three dialects with a protect-strip-restore sentinel swap: escaped
backslashes (`\\`) → `chr(1)`, then strip the now-unambiguous active NUL
escape, then restore. `chr(1)` is safe as a sentinel because `json.dumps`
escapes control bytes, so a raw `chr(1)` never appears in the stored JSON text.
Added a regression test asserting a literal `\u0000` in the data survives (xcom
all dialects + dag_run.conf).
*— posted via Claude Code (Claude Opus 4.8)*
##########
airflow-core/src/airflow/migrations/versions/0049_3_0_0_remove_pickled_data_from_xcom_table.py:
##########
@@ -41,13 +41,93 @@
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" so the production
calls below stay
+# byte-for-byte identical to the original inline SQL. 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, quoted) and the U+0000 (NUL) escape (stripped).
+_XCOM_PG_SANITIZE_SQL = r"""
+ UPDATE xcom
+ SET value = convert_to(
+ regexp_replace(
+ -- Strip the U+0000 (NUL) escape; illegal in
JSON/JSONB and not quotable.
+ replace(convert_from(value, 'UTF8'), '\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.
+ '([:,\[]\s*|^)(NaN|-?Infinity)(?=\s*[,}\]]|$)',
+ '\1"\2"',
+ 'g'
+ ),
+ 'UTF8'
+ )
+ WHERE value IS NOT NULL AND get_byte(value, 0) != 128
+ """
+_XCOM_MYSQL_SANITIZE_SQL = """
+ UPDATE xcom
+ SET value = CONVERT(
+ REGEXP_REPLACE(
+ -- Strip the U+0000 (NUL) escape before the cast
(illegal JSON; see PostgreSQL branch).
+ REPLACE(CONVERT(value USING utf8mb4), '\\\\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"',
+ 1,
+ 0,
+ 'c'
+ ) USING BINARY
+ )
+ WHERE value IS NOT NULL AND HEX(SUBSTRING(value, 1, 1)) != '80'
+ """
+_XCOM_SQLITE_SANITIZE_SQL = """
+ UPDATE xcom
+ SET value = CAST(
+ REPLACE(
+ REPLACE(
+ -- Step 1: replace NaN first so it doesn't
interfere with Infinity.
+ REPLACE(REPLACE(CAST(value AS TEXT), '\\u0000',
''), '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("UPDATE xcom", f"UPDATE {table}")
Review Comment:
`.format()` / `string.Template` don't fit here: the SQL bodies contain
literal `{`/`}` (regex like `[,}\]]`, `[ ]*`) that `.format()` would try to
parse, and `$` collides both with the regex `$` anchors and MySQL's `$1"$2"`
replacement. I switched the substitution to an explicit `__TABLE__` token
instead of matching `"UPDATE xcom"`, which reads more clearly.
*— posted via Claude Code (Claude Opus 4.8)*
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
