uranusjr commented on code in PR #69064:
URL: https://github.com/apache/airflow/pull/69064#discussion_r3484457711
##########
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:
Why not just make the SQL a template string and use `.format()` instead?
--
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]
