aaron-y-chen commented on code in PR #69985:
URL: https://github.com/apache/airflow/pull/69985#discussion_r3654044544
##########
airflow-core/src/airflow/migrations/versions/0094_3_2_0_replace_deadline_inline_callback_with_fkey.py:
##########
@@ -54,12 +54,102 @@
_ASYNC_CALLBACK_CLASSNAME = "airflow.sdk.definitions.deadline.AsyncCallback"
+def _serialize_extended(value):
+ """
+ Encode ``value`` in Airflow's extended-JSON format, matching
``BaseSerialization.serialize``.
+
+ ``callback.data`` is an ``ExtendedJSON`` column, so on read the runtime
runs
+ ``BaseSerialization.deserialize``, which requires every nested dict to be
wrapped as
+ ``{"__type": "dict", "__var": {...}}``. We must write that same wrapping
here (recursively)
+ rather than embedding raw nested dicts -- otherwise deserialize raises
``KeyError('__var')``
+ and crashes the scheduler. The dict/list/primitive subset below is all
that callback data
+ contains; the logic is inlined so the migration stays independent of
runtime serialization code.
+ """
+ if isinstance(value, dict):
+ return {"__type": "dict", "__var": {str(k): _serialize_extended(v) for
k, v in value.items()}}
+ if isinstance(value, list):
+ return [_serialize_extended(v) for v in value]
+ return value
+
+
+# Recursive extended-JSON encoder as a session-local (auto-dropped) SQL
function, so the
+# Postgres CTE path can wrap nested callback kwargs the same way
``_serialize_extended`` does.
+_PG_ENCODE_EXTENDED_DDL = dedent("""
+ CREATE OR REPLACE FUNCTION pg_temp.encode_extended(node jsonb) RETURNS
jsonb AS $$
+ DECLARE k text; v jsonb; out jsonb := '{}'::jsonb;
+ BEGIN
+ IF jsonb_typeof(node) = 'object' THEN
+ FOR k, v IN SELECT * FROM jsonb_each(node) LOOP
+ out := out || jsonb_build_object(k, pg_temp.encode_extended(v));
+ END LOOP;
+ RETURN jsonb_build_object('__type', 'dict', '__var', out);
+ ELSIF jsonb_typeof(node) = 'array' THEN
+ RETURN (SELECT jsonb_agg(pg_temp.encode_extended(e)) FROM
jsonb_array_elements(node) e);
+ END IF;
+ RETURN node;
+ END;
+ $$ LANGUAGE plpgsql;
+""")
+
+
+def _deserialize_extended(value):
+ """
+ Inverse of :func:`_serialize_extended`: unwrap extended-JSON back to plain
values.
+
+ Used by ``downgrade`` to rebuild the old inline callback ``kwargs`` (which
were stored
+ raw). Lenient: already-raw dicts (e.g. produced by the pre-fix version of
this migration)
+ pass through unchanged, so downgrade is correct regardless of which
version upgraded.
+ """
+ if isinstance(value, dict):
+ if "__type" in value and "__var" in value:
+ if value["__type"] == "dict" and isinstance(value["__var"], dict):
+ return {k: _deserialize_extended(v) for k, v in
value["__var"].items()}
+ return value
+ return {k: _deserialize_extended(v) for k, v in value.items()}
+ if isinstance(value, list):
+ return [_deserialize_extended(v) for v in value]
+ return value
+
+
+# SQL inverse of ``pg_temp.encode_extended`` for the Postgres downgrade path.
Lenient in the
+# same way as ``_deserialize_extended`` so it handles data written by either
version of upgrade.
+_PG_DECODE_EXTENDED_DDL = dedent("""
+ CREATE OR REPLACE FUNCTION pg_temp.decode_extended(node jsonb) RETURNS
jsonb AS $$
+ DECLARE k text; v jsonb; out jsonb := '{}'::jsonb;
+ BEGIN
+ IF jsonb_typeof(node) = 'object' THEN
+ IF node ? '__type' AND node ? '__var' THEN
+ IF node->>'__type' = 'dict' AND jsonb_typeof(node->'__var') =
'object' THEN
+ FOR k, v IN SELECT * FROM jsonb_each(node->'__var') LOOP
+ out := out || jsonb_build_object(k, pg_temp.decode_extended(v));
+ END LOOP;
+ RETURN out;
+ END IF;
+ RETURN node;
+ ELSE
+ FOR k, v IN SELECT * FROM jsonb_each(node) LOOP
+ out := out || jsonb_build_object(k, pg_temp.decode_extended(v));
+ END LOOP;
+ RETURN out;
+ END IF;
+ ELSIF jsonb_typeof(node) = 'array' THEN
+ RETURN (SELECT jsonb_agg(pg_temp.decode_extended(e)) FROM
jsonb_array_elements(node) e);
Review Comment:
Same as above.
##########
airflow-core/src/airflow/migrations/versions/0094_3_2_0_replace_deadline_inline_callback_with_fkey.py:
##########
@@ -54,12 +54,102 @@
_ASYNC_CALLBACK_CLASSNAME = "airflow.sdk.definitions.deadline.AsyncCallback"
+def _serialize_extended(value):
+ """
+ Encode ``value`` in Airflow's extended-JSON format, matching
``BaseSerialization.serialize``.
+
+ ``callback.data`` is an ``ExtendedJSON`` column, so on read the runtime
runs
+ ``BaseSerialization.deserialize``, which requires every nested dict to be
wrapped as
+ ``{"__type": "dict", "__var": {...}}``. We must write that same wrapping
here (recursively)
+ rather than embedding raw nested dicts -- otherwise deserialize raises
``KeyError('__var')``
+ and crashes the scheduler. The dict/list/primitive subset below is all
that callback data
+ contains; the logic is inlined so the migration stays independent of
runtime serialization code.
+ """
+ if isinstance(value, dict):
+ return {"__type": "dict", "__var": {str(k): _serialize_extended(v) for
k, v in value.items()}}
+ if isinstance(value, list):
+ return [_serialize_extended(v) for v in value]
+ return value
+
+
+# Recursive extended-JSON encoder as a session-local (auto-dropped) SQL
function, so the
+# Postgres CTE path can wrap nested callback kwargs the same way
``_serialize_extended`` does.
+_PG_ENCODE_EXTENDED_DDL = dedent("""
+ CREATE OR REPLACE FUNCTION pg_temp.encode_extended(node jsonb) RETURNS
jsonb AS $$
+ DECLARE k text; v jsonb; out jsonb := '{}'::jsonb;
+ BEGIN
+ IF jsonb_typeof(node) = 'object' THEN
+ FOR k, v IN SELECT * FROM jsonb_each(node) LOOP
+ out := out || jsonb_build_object(k, pg_temp.encode_extended(v));
+ END LOOP;
+ RETURN jsonb_build_object('__type', 'dict', '__var', out);
+ ELSIF jsonb_typeof(node) = 'array' THEN
+ RETURN (SELECT jsonb_agg(pg_temp.encode_extended(e)) FROM
jsonb_array_elements(node) e);
Review Comment:
I wonder if `node = []`, whether the entire return value would be `None`. I
guess that's not what you want, right? Should we handle this case here?
--
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]
