Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use Serialization to support more objects in XCom #16786

Closed
wants to merge 1 commit into from
Closed

Use Serialization to support more objects in XCom #16786

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions airflow/models/xcom.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,9 @@ def serialize_value(value: Any):
if conf.getboolean('core', 'enable_xcom_pickling'):
return pickle.dumps(value)
try:
return json.dumps(value).encode('UTF-8')
from airflow.serialization.serialized_objects import BaseSerialization

return json.dumps(BaseSerialization._serialize(value, fail_on_error=True)).encode('UTF-8')
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we should add a BaseSerialization.to_json() method for symmetry

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It already exists.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, excetp:

    @classmethod
    def to_json(cls, var: Union[DAG, BaseOperator, dict, list, set, tuple]) -> str:
        """Stringifies DAGs and operators contained by var and returns a JSON string of var."""
        return json.dumps(cls.to_dict(var), ensure_ascii=True)

    @classmethod
    def to_dict(cls, var: Union[DAG, BaseOperator, dict, list, set, tuple]) -> dict:
        """Stringifies DAGs and operators contained by var and returns a dict of var."""
        # Don't call on this class directly - only SerializedDAG or
        # SerializedBaseOperator should be used as the "entrypoint"
        raise NotImplementedError()

except (ValueError, TypeError):
log.error(
"Could not serialize the XCom value into JSON."
Expand All @@ -240,10 +242,14 @@ def deserialize_value(result: "XCom") -> Any:
try:
return pickle.loads(result.value)
except pickle.UnpicklingError:
return json.loads(result.value.decode('UTF-8'))
from airflow.serialization.serialized_objects import BaseSerialization

return BaseSerialization.from_json(result.value.decode('UTF-8'))
else:
try:
return json.loads(result.value.decode('UTF-8'))
from airflow.serialization.serialized_objects import BaseSerialization

return BaseSerialization.from_json(result.value.decode('UTF-8'))
except (json.JSONDecodeError, UnicodeDecodeError):
return pickle.loads(result.value)

Expand Down
4 changes: 3 additions & 1 deletion airflow/serialization/serialized_objects.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,7 @@ def serialize_to_json(
return serialized_object

@classmethod
def _serialize(cls, var: Any) -> Any: # Unfortunately there is no support for recursive types in mypy
def _serialize(cls, var: Any, fail_on_error: bool = False) -> Any:
"""Helper function of depth first search for serialization.

The serialization protocol is:
Expand Down Expand Up @@ -251,6 +251,8 @@ def _serialize(cls, var: Any) -> Any: # Unfortunately there is no support for r
elif isinstance(var, TaskGroup):
return SerializedTaskGroup.serialize_task_group(var)
else:
if fail_on_error:
raise TypeError(f"Can't serialize {var} - invalid type - {type(var)} ")
log.debug('Cast type %s to str in serialization.', type(var))
return str(var)

Expand Down
28 changes: 28 additions & 0 deletions tests/models/test_xcom.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,3 +259,31 @@ def test_xcom_init_on_load_uses_orm_deserialize_value(self, mock_orm_deserialize

instance.init_on_load()
mock_orm_deserialize.assert_called_once_with()

@conf_vars({("core", "xcom_enable_pickling"): "False"})
def test_xcom_roundtrip_with_set(self):
json_obj = {"key": {"value1", "value2"}}
execution_date = timezone.utcnow()
key = "xcom_test1"
dag_id = "test_dag1"
task_id = "test_task1"
XCom.set(key=key, value=json_obj, dag_id=dag_id, task_id=task_id, execution_date=execution_date)

ret_value = XCom.get_one(key=key, dag_id=dag_id, task_id=task_id, execution_date=execution_date)

assert ret_value == json_obj

session = settings.Session()
ret_value = (
session.query(XCom)
.filter(
XCom.key == key,
XCom.dag_id == dag_id,
XCom.task_id == task_id,
XCom.execution_date == execution_date,
)
.first()
.value
)

assert ret_value == json_obj