kaxil commented on code in PR #74858:
URL: https://github.com/apache/airflow/pull/74858#discussion_r4304559803
##########
providers/common/io/src/airflow/providers/common/io/xcom/backend.py:
##########
@@ -107,6 +112,19 @@ def _get_full_path(data: str) -> ObjectStoragePath:
raise ValueError(f"Not a valid url: {data}")
+ @classmethod
+ def store_value(cls, coordinate: XComCoordinate, key: str, value: Any) ->
str:
+ base = _get_base_path()
+ path = cls._write_large_value(
+ value,
+ lambda suffix: base.joinpath(
+ coordinate.build_path(
Review Comment:
The object name is now the quoted key plus the compression suffix, but
`deserialize_value` still opens it with `compression="infer"`. With offloading
on and compression off (the default), a key like `report.gz` or `data.zip` gets
written as plain JSON and read back as gzip/zip. fsspec raises `BadGzipFile` /
`BadZipFile` there, neither is in the `except (FileNotFoundError, TypeError,
ValueError)`, so the pull fails. The old `<uuid><suffix>` name couldn't hit
this.
Appending a fixed extension after the key (`<quoted key>.json` plus the
compression suffix) would keep the path deterministic and stop inference from
seeing the key's own extension. A round-trip test with a `.gz` key and
compression off would pin it.
##########
airflow-core/docs/core-concepts/xcoms.rst:
##########
@@ -106,11 +106,15 @@ The default XCom backend, BaseXCom, stores XComs in the
Airflow database, which
Custom XCom Backends
--------------------
-The XCom system has interchangeable backends, and you can set which backend is
being used via the ``xcom_backend`` configuration option.
-
-If you want to implement your own backend, you should subclass
:class:`~airflow.sdk.bases.xcom.BaseXCom`, and override the ``serialize_value``
and ``deserialize_value`` methods.
-
-You can override the ``purge`` method in the ``BaseXCom`` class to have
control over purging the XCom data from the custom backend. This will be called
as part of ``delete``.
+The XCom system has interchangeable backends, and you can set which backend is
being used via the ``xcom_backend`` configuration option. A backend is a
subclass of :class:`~airflow.sdk.bases.xcom.BaseXCom`.
+
+A backend that keeps values outside the metadata database overrides
``store_value``, which receives an
+:class:`~airflow.sdk.bases.xcom.XComCoordinate`, stores the value and returns
what the XCom row holds, and
+``load_value``, which turns that row value back into the value.
``XComCoordinate.build_path`` builds a storage
+location that is unique per task instance, or per write when the writer cannot
know the task instance. Override
+``purge`` to remove the stored data when an XCom is deleted. A backend that
only overrides ``serialize_value`` and
Review Comment:
Only the SDK's `BaseXCom.delete` calls `purge`. Deleting through the REST
API / UI (`delete_xcom_entry`) or `airflow db clean` drops the row and leaves
the object, so "when an XCom is deleted" promises more than happens. The old
"called as part of `delete`" was accurate.
##########
task-sdk/src/airflow/sdk/bases/xcom.py:
##########
@@ -31,12 +37,108 @@
XComSequenceSliceResult,
)
+if TYPE_CHECKING:
+ from pydantic import JsonValue
+
+ from airflow.sdk.types import TaskInstance
+
# Lightweight wrapper for XCom values
_XComValueWrapper = collections.namedtuple("_XComValueWrapper", "value")
log = structlog.get_logger(logger_name="task")
[email protected]
+class XComCoordinate:
Review Comment:
`XComCoordinate` is the parameter type of the new public `store_value`, and
the docs send backend authors to `build_path`, but it isn't exported from
`airflow.sdk` or listed in `task-sdk/docs/api.rst` (only `BaseXCom` is). Should
it go next to `BaseXCom` there? Otherwise backends will import it from
`airflow.sdk.bases.xcom` and that becomes the path we have to keep.
##########
devel-common/src/tests_common/pytest_plugin.py:
##########
@@ -3018,35 +3018,61 @@ def execute(self, context):
timezone = _import_timezone()
+ # kgb leaves a classmethod it unspied bound to the class it spied, which
breaks subclasses in later tests.
+ # A task's own pushes go through _push/_apush on newer Task SDK versions
and through set/aset on older ones.
+ push_methods = [name for name in ("set", "_push", "_apush") if
hasattr(XCom, name)]
+ spied_methods = [*push_methods, "get_one"]
+ classmethods = {
+ name: next((owner, vars(owner)[name]) for owner in XCom.__mro__ if
name in vars(owner))
+ for name in spied_methods
+ }
+
# Set up spies once at fixture level
- if hasattr(XCom.set, "spy"):
- spy_agency.unspy(XCom.set)
- if hasattr(XCom.get_one, "spy"):
- spy_agency.unspy(XCom.get_one)
- spy_agency.spy_on(XCom.set, call_original=True)
+ for name in spied_methods:
+ if hasattr(getattr(XCom, name), "spy"):
+ spy_agency.unspy(getattr(XCom, name))
+ for name in push_methods:
+ spy_agency.spy_on(getattr(XCom, name), call_original=True)
spy_agency.spy_on(
XCom.get_one, call_fake=lambda cls, *args, **kwargs:
_get_one_from_set_calls(*args, **kwargs)
)
+ def _recorded_pushes() -> list[tuple[str, Any, dict[str, Any]]]:
+ """List the pushed ``(key, value, coordinates and options)`` of every
spied push method in call order."""
Review Comment:
This says call order, but it collects every `set` call, then every `_push`,
then every `_apush`. A test that calls `XCom.set` for a key after
`ti.xcom_push` of the same key gets the `_push` value back from `get_one`.
Either record a sequence number per call or drop "in call order" from the
docstring.
--
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]
