hkc-8010 commented on code in PR #71926:
URL: https://github.com/apache/airflow/pull/71926#discussion_r3976277889
##########
airflow-core/src/airflow/migrations/versions/0049_3_0_0_remove_pickled_data_from_xcom_table.py:
##########
@@ -42,112 +42,179 @@
# --- 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 Exception:
Review Comment:
Narrowed it. A build without JSON1 raises `sqlalchemy.exc.OperationalError`
(`no such function: json_valid`), so `Exception` was wider than the case it was
there for, and it also caught a connection failure and quietly sanitized
without the guard, which rewrites rows that already parse.
The probe had no coverage, so I added two tests: one asserting True on a
real connection and False on an OperationalError, and one asserting anything
else propagates. The second one fails on the old `except Exception`.
--
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]
