Fix migration when xcom has NaN values

airflow-core/docs/img/airflow_erd.sha256

@@ -1 +1 @@
-88337ee3cd515161de57099c2582d4b26d22666a30387fa32aec89e86d7beeb3
+db83973c6b3584fbb7e06a204d3d33021f6efbfa7b00b9092fbe4af0a3a42533

airflow-core/src/airflow/migrations/versions/0049_3_0_0_remove_pickled_data_from_xcom_table.py

@@ -47,7 +47,8 @@ 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. Update the XCom.value column type to JSON from LargeBinary/LongBlob
+    # 4. Sanitize NaN values in JSON (convert to string)
+    # 5. Update the XCom.value column type to JSON from LargeBinary/LongBlob
 
     conn = op.get_bind()
     dialect = conn.dialect.name
@@ -111,8 +112,16 @@ 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}"))
 
-    # Update the value column type to JSON
+    # Update the values from nan to nan string
     if dialect == "postgresql":
+        conn.execute(
+            text("""
+                UPDATE xcom
+                SET value = convert_to(replace(convert_from(value, 'UTF8'), 'NaN', '"nan"'), 'UTF8')
+                WHERE value IS NOT NULL AND get_byte(value, 0) != 128
+            """)
+        )
+
         op.execute(
             """
             ALTER TABLE xcom
@@ -124,11 +133,27 @@ def upgrade():
             """
         )
     elif dialect == "mysql":
+        conn.execute(
+            text("""
+                UPDATE xcom
+                SET value = CONVERT(REPLACE(CONVERT(value USING utf8mb4), 'NaN', '"nan"') USING BINARY)
+                WHERE value IS NOT NULL AND HEX(SUBSTRING(value, 1, 1)) != '80'
+            """)
+        )
+
         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)")
         op.drop_column("xcom", "value")
         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(CAST(value AS TEXT), 'NaN', '"nan"') AS BLOB)
+                WHERE value IS NOT NULL AND hex(substr(value, 1, 1)) != '80'
+            """)
+        )
         # 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")