diff --git a/UPDATING.md b/UPDATING.md index 981de32f2736d..1a20eddb3235a 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -79,6 +79,11 @@ https://developers.google.com/style/inclusive-documentation --> +### Passing ``execution_date`` to ``XCom.set()``, ``XCom.clear()``, ``XCom.get_one()``, and ``XCom.get_many()`` is deprecated + +Continuing the effort to bind TaskInstance to a DagRun, XCom entries are now also tied to a DagRun. Use the ``run_id`` argument to specify the DagRun instead. + + ## Airflow 2.2.2 No breaking changes. diff --git a/airflow/api/common/experimental/get_lineage.py b/airflow/api/common/experimental/get_lineage.py index c6857a65fb279..6db225f485cf1 100644 --- a/airflow/api/common/experimental/get_lineage.py +++ b/airflow/api/common/experimental/get_lineage.py @@ -16,8 +16,9 @@ # specific language governing permissions and limitations # under the License. """Lineage apis""" +import collections import datetime -from typing import Any, Dict, List +from typing import Any, Dict from airflow.api.common.experimental import check_and_get_dag, check_and_get_dagrun from airflow.lineage import PIPELINE_INLETS, PIPELINE_OUTLETS @@ -26,23 +27,18 @@ @provide_session -def get_lineage(dag_id: str, execution_date: datetime.datetime, session=None) -> Dict[str, Dict[str, Any]]: - """Gets the lineage information for dag specified""" +def get_lineage(dag_id: str, execution_date: datetime.datetime, *, session) -> Dict[str, Dict[str, Any]]: + """Gets the lineage information for dag specified.""" dag = check_and_get_dag(dag_id) - check_and_get_dagrun(dag, execution_date) + dagrun = check_and_get_dagrun(dag, execution_date) - inlets: List[XCom] = XCom.get_many( - dag_ids=dag_id, execution_date=execution_date, key=PIPELINE_INLETS, session=session - ).all() - outlets: List[XCom] = XCom.get_many( - dag_ids=dag_id, execution_date=execution_date, key=PIPELINE_OUTLETS, session=session - ).all() + inlets = XCom.get_many(dag_ids=dag_id, run_id=dagrun.run_id, key=PIPELINE_INLETS, session=session) + outlets = XCom.get_many(dag_ids=dag_id, run_id=dagrun.run_id, key=PIPELINE_OUTLETS, session=session) - lineage: Dict[str, Dict[str, Any]] = {} + lineage: Dict[str, Dict[str, Any]] = collections.defaultdict(dict) for meta in inlets: - lineage[meta.task_id] = {'inlets': meta.value} - + lineage[meta.task_id]["inlets"] = meta.value for meta in outlets: - lineage[meta.task_id]['outlets'] = meta.value + lineage[meta.task_id]["outlets"] = meta.value - return {'task_ids': lineage} + return {"task_ids": {k: v for k, v in lineage.items()}} diff --git a/airflow/cli/commands/task_command.py b/airflow/cli/commands/task_command.py index 5b1f714e935d1..9ea4f4daf5bbe 100644 --- a/airflow/cli/commands/task_command.py +++ b/airflow/cli/commands/task_command.py @@ -36,6 +36,7 @@ from airflow.models import DagPickle, TaskInstance from airflow.models.dag import DAG from airflow.models.dagrun import DagRun +from airflow.models.xcom import IN_MEMORY_DAGRUN_ID from airflow.ti_deps.dep_context import DepContext from airflow.ti_deps.dependencies_deps import SCHEDULER_QUEUED_DEPS from airflow.utils import cli as cli_utils @@ -74,7 +75,7 @@ def _get_dag_run(dag, exec_date_or_run_id, create_if_necessary, session): ) except NoResultFound: if create_if_necessary: - return DagRun(dag.dag_id, execution_date=execution_date) + return DagRun(dag.dag_id, run_id=IN_MEMORY_DAGRUN_ID, execution_date=execution_date) raise DagRunNotFound( f"DagRun for {dag.dag_id} with run_id or execution_date of {exec_date_or_run_id!r} not found" ) from None @@ -87,7 +88,7 @@ def _get_ti(task, exec_date_or_run_id, create_if_necessary=False, session=None): ti = dag_run.get_task_instance(task.task_id) if not ti and create_if_necessary: - ti = TaskInstance(task, run_id=None) + ti = TaskInstance(task, run_id=dag_run.run_id) ti.dag_run = dag_run ti.refresh_from_task(task) return ti diff --git a/airflow/models/skipmixin.py b/airflow/models/skipmixin.py index 765a94712ca0e..de5f1fd597bbb 100644 --- a/airflow/models/skipmixin.py +++ b/airflow/models/skipmixin.py @@ -124,7 +124,7 @@ def skip( value={XCOM_SKIPMIXIN_SKIPPED: [d.task_id for d in tasks]}, task_id=task_id, dag_id=dag_run.dag_id, - execution_date=dag_run.execution_date, + run_id=dag_run.run_id, session=session, ) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index d8463b9aeeea7..660da432d9806 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -2123,31 +2123,31 @@ def xcom_push( """ Make an XCom available for tasks to pull. - :param key: A key for the XCom + :param key: Key to store the value under. :type key: str - :param value: A value for the XCom. The value is pickled and stored - in the database. - :type value: any picklable object - :param execution_date: if provided, the XCom will not be visible until - this date. This can be used, for example, to send a message to a - task on a future date without it being immediately visible. + :param value: Value to store. What types are possible depends on whether + ``enable_xcom_pickling`` is true or not. If so, this can be any + picklable object; only be JSON-serializable may be used otherwise. + :param execution_date: Deprecated parameter that has no effect. :type execution_date: datetime - :param session: Sqlalchemy ORM Session - :type session: Session """ - self_execution_date = self.get_dagrun(session).execution_date - if execution_date and execution_date < self_execution_date: - raise ValueError( - f'execution_date can not be in the past (current execution_date is ' - f'{self_execution_date}; received {execution_date})' - ) + if execution_date is not None: + self_execution_date = self.get_dagrun(session).execution_date + if execution_date < self_execution_date: + raise ValueError( + f'execution_date can not be in the past (current execution_date is ' + f'{self_execution_date}; received {execution_date})' + ) + elif execution_date is not None: + message = "Passing 'execution_date' to 'TaskInstance.xcom_push()' is deprecated." + warnings.warn(message, DeprecationWarning, stacklevel=3) XCom.set( key=key, value=value, task_id=self.task_id, dag_id=self.dag_id, - execution_date=execution_date or self_execution_date, + run_id=self.run_id, session=session, ) @@ -2194,11 +2194,9 @@ def xcom_pull( if dag_id is None: dag_id = self.dag_id - execution_date = self.get_dagrun(session).execution_date - query = XCom.get_many( - execution_date=execution_date, key=key, + run_id=self.run_id, dag_ids=dag_id, task_ids=task_ids, include_prior_dates=include_prior_dates, diff --git a/airflow/models/xcom.py b/airflow/models/xcom.py index a3134e2ca83f3..109540826b869 100644 --- a/airflow/models/xcom.py +++ b/airflow/models/xcom.py @@ -20,6 +20,7 @@ import json import logging import pickle +import warnings from typing import TYPE_CHECKING, Any, Iterable, Optional, Type, Union, cast, overload import pendulum @@ -42,6 +43,17 @@ MAX_XCOM_SIZE = 49344 XCOM_RETURN_KEY = 'return_value' +# Work around 'airflow task test' generating a temporary in-memory DAG run +# without storing it in the database. To avoid interfering with actual XCom +# entries but still behave _somewhat_ consistently, we store XCom to a distant +# time in the future. Eventually we want to migrate XCom's primary to use run_id +# instead, so execution_date can just be None for this case. +IN_MEMORY_DAGRUN_ID = "__airflow_in_memory_dagrun__" + +# This is the largest possible value we can store in MySQL. +# https://dev.mysql.com/doc/refman/5.7/en/datetime.html +_DISTANT_FUTURE = datetime.datetime(2038, 1, 19, 3, 14, 7, tzinfo=timezone.utc) + class BaseXCom(Base, LoggingMixin): """Base class for XCom objects.""" @@ -137,11 +149,19 @@ def set( if not exactly_one(execution_date is not None, run_id is not None): raise ValueError("Exactly one of execution_date or run_id must be passed") - if run_id: + if run_id == IN_MEMORY_DAGRUN_ID: + execution_date = _DISTANT_FUTURE + elif run_id is not None: from airflow.models.dagrun import DagRun - dag_run = session.query(DagRun).filter_by(dag_id=dag_id, run_id=run_id).one() - execution_date = dag_run.execution_date + execution_date = ( + session.query(DagRun.execution_date) + .filter(DagRun.dag_id == dag_id, DagRun.run_id == run_id) + .scalar() + ) + else: # Guarantees execution_date is not None. + message = "Passing 'execution_date' to 'XCom.set()' is deprecated. Use 'run_id' instead." + warnings.warn(message, DeprecationWarning, stacklevel=3) # Remove duplicate XComs and insert a new one. session.query(cls).filter( @@ -238,6 +258,9 @@ def get_one( session=session, ) elif execution_date is not None: + message = "Passing 'execution_date' to 'XCom.get_one()' is deprecated. Use 'run_id' instead." + warnings.warn(message, PendingDeprecationWarning, stacklevel=3) + query = cls.get_many( execution_date=execution_date, key=key, @@ -319,48 +342,48 @@ def get_many( run_id: Optional[str] = None, ) -> Query: """:sphinx-autoapi-skip:""" + from airflow.models.dagrun import DagRun + if not exactly_one(execution_date is not None, run_id is not None): raise ValueError("Exactly one of execution_date or run_id must be passed") + if execution_date is not None: + message = "Passing 'execution_date' to 'XCom.get_many()' is deprecated. Use 'run_id' instead." + warnings.warn(message, PendingDeprecationWarning, stacklevel=3) - filters = [] + query = session.query(cls) if key: - filters.append(cls.key == key) + query = query.filter(cls.key == key) - if task_ids: - if is_container(task_ids): - filters.append(cls.task_id.in_(task_ids)) - else: - filters.append(cls.task_id == task_ids) + if is_container(task_ids): + query = query.filter(cls.task_id.in_(task_ids)) + elif task_ids is not None: + query = query.filter(cls.task_id == task_ids) - if dag_ids: - if is_container(dag_ids): - filters.append(cls.dag_id.in_(dag_ids)) - else: - filters.append(cls.dag_id == dag_ids) + if is_container(dag_ids): + query = query.filter(cls.dag_id.in_(dag_ids)) + elif dag_ids is not None: + query = query.filter(cls.dag_id == dag_ids) if include_prior_dates: - if execution_date is None: - # In theory it would be possible to build a subquery that joins to DagRun and then gets the - # execution dates. Lets do that for 2.3 - raise ValueError("Using include_prior_dates needs an execution_date to be passed") - filters.append(cls.execution_date <= execution_date) + if execution_date is not None: + query = query.filter(cls.execution_date <= execution_date) + else: + # This returns an empty query result for IN_MEMORY_DAGRUN_ID, + # but that is impossible to implement. Sorry? + dr = session.query(DagRun.execution_date).filter(DagRun.run_id == run_id).subquery() + query = query.filter(cls.execution_date <= dr.c.execution_date) elif execution_date is not None: - filters.append(cls.execution_date == execution_date) - - query = session.query(cls).filter(*filters) - - if run_id: - from airflow.models.dagrun import DagRun - + query = query.filter(cls.execution_date == execution_date) + elif run_id == IN_MEMORY_DAGRUN_ID: + query = query.filter(cls.execution_date == _DISTANT_FUTURE) + else: query = query.join(cls.dag_run).filter(DagRun.run_id == run_id) query = query.order_by(cls.execution_date.desc(), cls.timestamp.desc()) - if limit: return query.limit(limit) - else: - return query + return query @classmethod @provide_session @@ -423,17 +446,18 @@ def clear( if not exactly_one(execution_date is not None, run_id is not None): raise ValueError("Exactly one of execution_date or run_id must be passed") - query = session.query(cls).filter( - cls.dag_id == dag_id, - cls.task_id == task_id, - ) - + query = session.query(cls).filter(cls.dag_id == dag_id, cls.task_id == task_id) if execution_date is not None: + message = "Passing 'execution_date' to 'XCom.clear()' is deprecated. Use 'run_id' instead." + warnings.warn(message, DeprecationWarning, stacklevel=3) query = query.filter(cls.execution_date == execution_date) + elif run_id == IN_MEMORY_DAGRUN_ID: + query = query.filter(cls.execution_date == _DISTANT_FUTURE) else: from airflow.models.dagrun import DagRun - query = query.join(cls.dag_run).filter(DagRun.run_id == run_id) + execution_date = session.query(DagRun.execution_date).filter(DagRun.run_id == run_id).scalar() + query = query.filter(cls.execution_date == execution_date) return query.delete() diff --git a/airflow/providers/amazon/aws/operators/ecs.py b/airflow/providers/amazon/aws/operators/ecs.py index f560dff4e7471..f651cc7dd758d 100644 --- a/airflow/providers/amazon/aws/operators/ecs.py +++ b/airflow/providers/amazon/aws/operators/ecs.py @@ -393,7 +393,7 @@ def _xcom_set(self, context, key, value, task_id): value=value, task_id=task_id, dag_id=self.dag_id, - execution_date=context["ti"].execution_date, + run_id=context["run_id"], ) def _try_reattach_task(self, context): @@ -431,8 +431,9 @@ def _aws_logs_enabled(self): return self.awslogs_group and self.awslogs_stream_prefix def _get_task_log_fetcher(self) -> ECSTaskLogFetcher: + if not self.awslogs_group: + raise ValueError("must specify awslogs_group to fetch task logs") log_stream_name = f"{self.awslogs_stream_prefix}/{self.ecs_task_id}" - return ECSTaskLogFetcher( aws_conn_id=self.aws_conn_id, region_name=self.awslogs_region, diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 7e84d677e60f1..4cbaebbd73026 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -1086,6 +1086,7 @@ pem performant personalizations pformat +picklable pid pigcmd pinodb diff --git a/tests/models/test_xcom.py b/tests/models/test_xcom.py index 35c7e60d3ca75..77d3fbcb099ce 100644 --- a/tests/models/test_xcom.py +++ b/tests/models/test_xcom.py @@ -14,20 +14,52 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import datetime +import operator import os from unittest import mock import pytest from airflow.configuration import conf +from airflow.models.dagrun import DagRun, DagRunType from airflow.models.xcom import XCOM_RETURN_KEY, BaseXCom, XCom, resolve_xcom_backend from airflow.utils import timezone +from airflow.utils.session import create_session from tests.test_utils.config import conf_vars class CustomXCom(BaseXCom): - def orm_deserialize_value(self): - return 'Short value...' + orm_deserialize_value = mock.Mock() + + +@pytest.fixture(scope="module", autouse=True) +def reset_db(): + """Delete XCom entries left over by other test modules before we start.""" + with create_session() as session: + session.query(XCom).delete() + + +@pytest.fixture() +def dag_run_factory(session): + def func(dag_id, execution_date): + run = DagRun( + dag_id=dag_id, + run_type=DagRunType.SCHEDULED, + run_id=DagRun.generate_run_id(DagRunType.SCHEDULED, execution_date), + execution_date=execution_date, + ) + session.add(run) + session.flush() + return run + + yield func + session.flush() + + +@pytest.fixture() +def dag_run(dag_run_factory): + return dag_run_factory(dag_id="dag", execution_date=timezone.datetime(2021, 12, 3, 4, 56)) class TestXCom: @@ -50,110 +82,48 @@ def test_resolve_xcom_class_fallback_to_basexcom_no_config(self): assert issubclass(cls, BaseXCom) assert cls().serialize_value([1]) == b"[1]" - @pytest.mark.parametrize( - ("enable_xcom_pickling",), - [ - pytest.param(True, id='enable_xcom_pickling=True'), - pytest.param(False, id='enable_xcom_pickling=False'), - ], - ) - def test_xcom_get_one_get_many(self, enable_xcom_pickling, session): - json_obj = {"key": "value"} - execution_date = timezone.utcnow() - key = "xcom_test1" - dag_id = "test_dag1" - task_id = "test_task1" - - with conf_vars({("core", "enable_xcom_pickling"): str(enable_xcom_pickling)}): - XCom.set( - key=key, - value=json_obj, - dag_id=dag_id, - task_id=task_id, - execution_date=execution_date, - session=session, - ) - - ret_value = ( - XCom.get_many( - key=key, dag_ids=dag_id, task_ids=task_id, execution_date=execution_date, session=session - ) - .first() - .value - ) - - assert ret_value == json_obj - - ret_value = XCom.get_one( - key=key, dag_id=dag_id, task_id=task_id, execution_date=execution_date, session=session - ) - - assert ret_value == json_obj - - 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 - - def test_xcom_deserialize_with_json_to_pickle_switch(self, session): - json_obj = {"key": "value"} - execution_date = timezone.utcnow() - key = "xcom_test3" - dag_id = "test_dag" - task_id = "test_task3" - + def test_xcom_deserialize_with_json_to_pickle_switch(self, dag_run, session): with conf_vars({("core", "enable_xcom_pickling"): "False"}): XCom.set( - key=key, - value=json_obj, - dag_id=dag_id, - task_id=task_id, - execution_date=execution_date, + key="xcom_test3", + value={"key": "value"}, + dag_id=dag_run.dag_id, + task_id="test_task3", + run_id=dag_run.run_id, session=session, ) - with conf_vars({("core", "enable_xcom_pickling"): "True"}): ret_value = XCom.get_one( - key=key, dag_id=dag_id, task_id=task_id, execution_date=execution_date, session=session + key="xcom_test3", + dag_id=dag_run.dag_id, + task_id="test_task3", + run_id=dag_run.run_id, + session=session, ) + assert ret_value == {"key": "value"} - assert ret_value == json_obj - - def test_xcom_deserialize_with_pickle_to_json_switch(self, session): - json_obj = {"key": "value"} - execution_date = timezone.utcnow() - key = "xcom_test3" - dag_id = "test_dag" - task_id = "test_task3" - + def test_xcom_deserialize_with_pickle_to_json_switch(self, dag_run, session): with conf_vars({("core", "enable_xcom_pickling"): "True"}): XCom.set( - key=key, - value=json_obj, - dag_id=dag_id, - task_id=task_id, - execution_date=execution_date, + key="xcom_test3", + value={"key": "value"}, + dag_id=dag_run.dag_id, + task_id="test_task3", + run_id=dag_run.run_id, session=session, ) - with conf_vars({("core", "enable_xcom_pickling"): "False"}): ret_value = XCom.get_one( - key=key, dag_id=dag_id, task_id=task_id, execution_date=execution_date, session=session + key="xcom_test3", + dag_id=dag_run.dag_id, + task_id="test_task3", + run_id=dag_run.run_id, + session=session, ) - - assert ret_value == json_obj + assert ret_value == {"key": "value"} @conf_vars({("core", "xcom_enable_pickling"): "False"}) - def test_xcom_disable_pickle_type_fail_on_non_json(self, session): + def test_xcom_disable_pickle_type_fail_on_non_json(self, dag_run, session): class PickleRce: def __reduce__(self): return os.system, ("ls -alt",) @@ -162,48 +132,14 @@ def __reduce__(self): XCom.set( key="xcom_test3", value=PickleRce(), - dag_id="test_dag3", + dag_id=dag_run.dag_id, task_id="test_task3", - execution_date=timezone.utcnow(), + run_id=dag_run.run_id, session=session, ) - @conf_vars({("core", "xcom_enable_pickling"): "True"}) - def test_xcom_get_many(self, session): - json_obj = {"key": "value"} - execution_date = timezone.utcnow() - key = "xcom_test4" - dag_id1 = "test_dag4" - task_id1 = "test_task4" - dag_id2 = "test_dag5" - task_id2 = "test_task5" - - XCom.set( - key=key, - value=json_obj, - dag_id=dag_id1, - task_id=task_id1, - execution_date=execution_date, - session=session, - ) - - XCom.set( - key=key, - value=json_obj, - dag_id=dag_id2, - task_id=task_id2, - execution_date=execution_date, - session=session, - ) - - results = XCom.get_many(key=key, execution_date=execution_date, session=session) - - for result in results: - assert result.value == json_obj - @mock.patch("airflow.models.xcom.XCom.orm_deserialize_value") def test_xcom_init_on_load_uses_orm_deserialize_value(self, mock_orm_deserialize): - instance = BaseXCom( key="key", value="value", @@ -212,29 +148,347 @@ def test_xcom_init_on_load_uses_orm_deserialize_value(self, mock_orm_deserialize task_id="task_id", dag_id="dag_id", ) - instance.init_on_load() mock_orm_deserialize.assert_called_once_with() @conf_vars({("core", "xcom_backend"): "tests.models.test_xcom.CustomXCom"}) - def test_get_one_doesnt_use_orm_deserialize_value(self, session): + def test_get_one_custom_backend_no_use_orm_deserialize_value(self, dag_run, session): """Test that XCom.get_one does not call orm_deserialize_value""" - json_obj = {"key": "value"} - execution_date = timezone.utcnow() - key = XCOM_RETURN_KEY - dag_id = "test_dag" - task_id = "test_task" - XCom = resolve_xcom_backend() XCom.set( + key=XCOM_RETURN_KEY, + value={"key": "value"}, + dag_id=dag_run.dag_id, + task_id="test_task", + run_id=dag_run.run_id, + session=session, + ) + + value = XCom.get_one( + dag_id=dag_run.dag_id, + task_id="test_task", + run_id=dag_run.run_id, + session=session, + ) + assert value == {"key": "value"} + XCom.orm_deserialize_value.assert_not_called() + + +@pytest.fixture( + params=[ + pytest.param("true", id="enable_xcom_pickling=true"), + pytest.param("false", id="enable_xcom_pickling=false"), + ], +) +def setup_xcom_pickling(request): + with conf_vars({("core", "enable_xcom_pickling"): str(request.param)}): + yield + + +@pytest.fixture() +def push_simple_json_xcom(session): + def func(*, dag_run: DagRun, task_id: str, key: str, value): + return XCom.set( key=key, - value=json_obj, - dag_id=dag_id, + value=value, + dag_id=dag_run.dag_id, task_id=task_id, - execution_date=execution_date, + run_id=dag_run.run_id, session=session, ) - value = XCom.get_one(dag_id=dag_id, task_id=task_id, execution_date=execution_date, session=session) + return func + + +@pytest.mark.usefixtures("setup_xcom_pickling") +class TestXComGet: + @pytest.fixture() + def setup_for_xcom_get_one(self, dag_run, push_simple_json_xcom): + push_simple_json_xcom(dag_run=dag_run, task_id="task_id_1", key="xcom_1", value={"key": "value"}) + + @pytest.mark.usefixtures("setup_for_xcom_get_one") + def test_xcom_get_one(self, session, dag_run): + stored_value = XCom.get_one( + key="xcom_1", + dag_id=dag_run.dag_id, + task_id="task_id_1", + run_id=dag_run.run_id, + session=session, + ) + assert stored_value == {"key": "value"} + + @pytest.mark.usefixtures("setup_for_xcom_get_one") + def test_xcom_get_one_with_execution_date(self, session, dag_run): + with pytest.deprecated_call(): + stored_value = XCom.get_one( + key="xcom_1", + dag_id=dag_run.dag_id, + task_id="task_id_1", + execution_date=dag_run.logical_date, + session=session, + ) + assert stored_value == {"key": "value"} + + @pytest.fixture() + def dag_runs_for_xcom_get_one_from_prior_date(self, dag_run_factory, push_simple_json_xcom): + date1 = timezone.datetime(2021, 12, 3, 4, 56) + dr1 = dag_run_factory(dag_id="dag", execution_date=date1) + dr2 = dag_run_factory(dag_id="dag", execution_date=date1 + datetime.timedelta(days=1)) + + # The earlier run pushes an XCom, but not the later run, but the later + # run can get this earlier XCom with ``include_prior_dates``. + push_simple_json_xcom(dag_run=dr1, task_id="task_1", key="xcom_1", value={"key": "value"}) + + return dr1, dr2 + + def test_xcom_get_one_from_prior_date(self, session, dag_runs_for_xcom_get_one_from_prior_date): + _, dr2 = dag_runs_for_xcom_get_one_from_prior_date + retrieved_value = XCom.get_one( + run_id=dr2.run_id, + key="xcom_1", + task_id="task_1", + dag_id="dag", + include_prior_dates=True, + session=session, + ) + assert retrieved_value == {"key": "value"} + + def test_xcom_get_one_from_prior_with_execution_date( + self, + session, + dag_runs_for_xcom_get_one_from_prior_date, + ): + _, dr2 = dag_runs_for_xcom_get_one_from_prior_date + with pytest.deprecated_call(): + retrieved_value = XCom.get_one( + execution_date=dr2.execution_date, + key="xcom_1", + task_id="task_1", + dag_id="dag", + include_prior_dates=True, + session=session, + ) + assert retrieved_value == {"key": "value"} + + @pytest.fixture() + def setup_for_xcom_get_many_single_argument_value(self, dag_run, push_simple_json_xcom): + push_simple_json_xcom(dag_run=dag_run, task_id="task_id_1", key="xcom_1", value={"key": "value"}) + + @pytest.mark.usefixtures("setup_for_xcom_get_many_single_argument_value") + def test_xcom_get_many_single_argument_value(self, session, dag_run): + stored_xcoms = XCom.get_many( + key="xcom_1", + dag_ids=dag_run.dag_id, + task_ids="task_id_1", + run_id=dag_run.run_id, + session=session, + ).all() + assert len(stored_xcoms) == 1 + assert stored_xcoms[0].key == "xcom_1" + assert stored_xcoms[0].value == {"key": "value"} + + @pytest.mark.usefixtures("setup_for_xcom_get_many_single_argument_value") + def test_xcom_get_many_single_argument_value_with_execution_date(self, session, dag_run): + with pytest.deprecated_call(): + stored_xcoms = XCom.get_many( + execution_date=dag_run.logical_date, + key="xcom_1", + dag_ids=dag_run.dag_id, + task_ids="task_id_1", + session=session, + ).all() + assert len(stored_xcoms) == 1 + assert stored_xcoms[0].key == "xcom_1" + assert stored_xcoms[0].value == {"key": "value"} + + @pytest.fixture() + def setup_for_xcom_get_many_multiple_tasks(self, dag_run, push_simple_json_xcom): + push_simple_json_xcom(dag_run=dag_run, key="xcom_1", value={"key1": "value1"}, task_id="task_id_1") + push_simple_json_xcom(dag_run=dag_run, key="xcom_1", value={"key2": "value2"}, task_id="task_id_2") + + @pytest.mark.usefixtures("setup_for_xcom_get_many_multiple_tasks") + def test_xcom_get_many_multiple_tasks(self, session, dag_run): + stored_xcoms = XCom.get_many( + key="xcom_1", + dag_ids=dag_run.dag_id, + task_ids=["task_id_1", "task_id_2"], + run_id=dag_run.run_id, + session=session, + ) + sorted_values = [x.value for x in sorted(stored_xcoms, key=operator.attrgetter("task_id"))] + assert sorted_values == [{"key1": "value1"}, {"key2": "value2"}] + + @pytest.mark.usefixtures("setup_for_xcom_get_many_multiple_tasks") + def test_xcom_get_many_multiple_tasks_with_execution_date(self, session, dag_run): + with pytest.deprecated_call(): + stored_xcoms = XCom.get_many( + execution_date=dag_run.logical_date, + key="xcom_1", + dag_ids=dag_run.dag_id, + task_ids=["task_id_1", "task_id_2"], + session=session, + ) + sorted_values = [x.value for x in sorted(stored_xcoms, key=operator.attrgetter("task_id"))] + assert sorted_values == [{"key1": "value1"}, {"key2": "value2"}] + + @pytest.fixture() + def dag_runs_for_xcom_get_many_from_prior_dates(self, dag_run_factory, push_simple_json_xcom): + date1 = timezone.datetime(2021, 12, 3, 4, 56) + date2 = date1 + datetime.timedelta(days=1) + dr1 = dag_run_factory(dag_id="dag", execution_date=date1) + dr2 = dag_run_factory(dag_id="dag", execution_date=date2) + push_simple_json_xcom(dag_run=dr1, task_id="task_1", key="xcom_1", value={"key1": "value1"}) + push_simple_json_xcom(dag_run=dr2, task_id="task_1", key="xcom_1", value={"key2": "value2"}) + return dr1, dr2 + + def test_xcom_get_many_from_prior_dates(self, session, dag_runs_for_xcom_get_many_from_prior_dates): + dr1, dr2 = dag_runs_for_xcom_get_many_from_prior_dates + stored_xcoms = XCom.get_many( + run_id=dr2.run_id, + key="xcom_1", + dag_ids="dag", + task_ids="task_1", + include_prior_dates=True, + session=session, + ) - assert value == json_obj + # The retrieved XComs should be ordered by logical date, latest first. + assert [x.value for x in stored_xcoms] == [{"key2": "value2"}, {"key1": "value1"}] + assert [x.execution_date for x in stored_xcoms] == [dr2.logical_date, dr1.logical_date] + + def test_xcom_get_many_from_prior_dates_with_execution_date( + self, + session, + dag_runs_for_xcom_get_many_from_prior_dates, + ): + dr1, dr2 = dag_runs_for_xcom_get_many_from_prior_dates + with pytest.deprecated_call(): + stored_xcoms = XCom.get_many( + execution_date=dr2.execution_date, + key="xcom_1", + dag_ids="dag", + task_ids="task_1", + include_prior_dates=True, + session=session, + ) + + # The retrieved XComs should be ordered by logical date, latest first. + assert [x.value for x in stored_xcoms] == [{"key2": "value2"}, {"key1": "value1"}] + assert [x.execution_date for x in stored_xcoms] == [dr2.logical_date, dr1.logical_date] + + +@pytest.mark.usefixtures("setup_xcom_pickling") +class TestXComSet: + def test_xcom_set(self, session, dag_run): + XCom.set( + key="xcom_1", + value={"key": "value"}, + dag_id=dag_run.dag_id, + task_id="task_1", + run_id=dag_run.run_id, + session=session, + ) + stored_xcoms = session.query(XCom).all() + assert stored_xcoms[0].key == "xcom_1" + assert stored_xcoms[0].value == {"key": "value"} + assert stored_xcoms[0].dag_id == "dag" + assert stored_xcoms[0].task_id == "task_1" + assert stored_xcoms[0].execution_date == dag_run.logical_date + + def test_xcom_set_with_execution_date(self, session, dag_run): + with pytest.deprecated_call(): + XCom.set( + key="xcom_1", + value={"key": "value"}, + dag_id=dag_run.dag_id, + task_id="task_1", + execution_date=dag_run.execution_date, + session=session, + ) + stored_xcoms = session.query(XCom).all() + assert stored_xcoms[0].key == "xcom_1" + assert stored_xcoms[0].value == {"key": "value"} + assert stored_xcoms[0].dag_id == "dag" + assert stored_xcoms[0].task_id == "task_1" + assert stored_xcoms[0].execution_date == dag_run.logical_date + + @pytest.fixture() + def setup_for_xcom_set_again_replace(self, dag_run, push_simple_json_xcom): + push_simple_json_xcom(dag_run=dag_run, task_id="task_1", key="xcom_1", value={"key1": "value1"}) + + @pytest.mark.usefixtures("setup_for_xcom_set_again_replace") + def test_xcom_set_again_replace(self, session, dag_run): + assert session.query(XCom).one().value == {"key1": "value1"} + XCom.set( + key="xcom_1", + value={"key2": "value2"}, + dag_id=dag_run.dag_id, + task_id="task_1", + run_id=dag_run.run_id, + session=session, + ) + assert session.query(XCom).one().value == {"key2": "value2"} + + @pytest.mark.usefixtures("setup_for_xcom_set_again_replace") + def test_xcom_set_again_replace_with_execution_date(self, session, dag_run): + assert session.query(XCom).one().value == {"key1": "value1"} + with pytest.deprecated_call(): + XCom.set( + key="xcom_1", + value={"key2": "value2"}, + dag_id=dag_run.dag_id, + task_id="task_1", + execution_date=dag_run.logical_date, + session=session, + ) + assert session.query(XCom).one().value == {"key2": "value2"} + + +@pytest.mark.usefixtures("setup_xcom_pickling") +class TestXComClear: + @pytest.fixture() + def setup_for_xcom_clear(self, dag_run, push_simple_json_xcom): + push_simple_json_xcom(dag_run=dag_run, task_id="task_1", key="xcom_1", value={"key": "value"}) + + @pytest.mark.usefixtures("setup_for_xcom_clear") + def test_xcom_clear(self, session, dag_run): + assert session.query(XCom).count() == 1 + XCom.clear( + dag_id=dag_run.dag_id, + task_id="task_1", + run_id=dag_run.run_id, + session=session, + ) + assert session.query(XCom).count() == 0 + + @pytest.mark.usefixtures("setup_for_xcom_clear") + def test_xcom_clear_with_execution_date(self, session, dag_run): + assert session.query(XCom).count() == 1 + with pytest.deprecated_call(): + XCom.clear( + dag_id=dag_run.dag_id, + task_id="task_1", + execution_date=dag_run.execution_date, + session=session, + ) + assert session.query(XCom).count() == 0 + + @pytest.mark.usefixtures("setup_for_xcom_clear") + def test_xcom_clear_different_run(self, session, dag_run): + XCom.clear( + dag_id=dag_run.dag_id, + task_id="task_1", + run_id="different_run", + session=session, + ) + assert session.query(XCom).count() == 1 + + @pytest.mark.usefixtures("setup_for_xcom_clear") + def test_xcom_clear_different_execution_date(self, session, dag_run): + XCom.clear( + dag_id=dag_run.dag_id, + task_id="task_1", + execution_date=timezone.utcnow(), + session=session, + ) + assert session.query(XCom).count() == 1 diff --git a/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py b/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py index dfb99ed3e0850..5caa50062c5b3 100644 --- a/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py +++ b/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py @@ -24,6 +24,7 @@ from airflow.exceptions import AirflowException from airflow.models import DAG, DagRun, TaskInstance +from airflow.models.xcom import IN_MEMORY_DAGRUN_ID from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator from airflow.utils import timezone from airflow.utils.state import State @@ -50,8 +51,8 @@ def setUp(self): @staticmethod def create_context(task): dag = DAG(dag_id="dag") - task_instance = TaskInstance(task=task, run_id="kub_pod_test") - task_instance.dag_run = DagRun(run_id="kub_pod_test", execution_date=DEFAULT_DATE) + task_instance = TaskInstance(task=task, run_id=IN_MEMORY_DAGRUN_ID) + task_instance.dag_run = DagRun(run_id=IN_MEMORY_DAGRUN_ID, execution_date=DEFAULT_DATE) return { "dag": dag, "ts": DEFAULT_DATE.isoformat(), @@ -681,8 +682,8 @@ def test_push_xcom_pod_info(self): do_xcom_push=False, ) pod = self.run_pod(k) - ti = TaskInstance(task=k, run_id="test_push_xcom_pod_info") - ti.dag_run = DagRun(run_id="test_push_xcom_pod_info", execution_date=DEFAULT_DATE) + ti = TaskInstance(task=k, run_id=IN_MEMORY_DAGRUN_ID) + ti.dag_run = DagRun(run_id=IN_MEMORY_DAGRUN_ID, execution_date=DEFAULT_DATE) pod_name = ti.xcom_pull(task_ids=k.task_id, key='pod_name') pod_namespace = ti.xcom_pull(task_ids=k.task_id, key='pod_namespace') assert pod_name and pod_name == pod.metadata.name