jason810496 commented on code in PR #52325:
URL: https://github.com/apache/airflow/pull/52325#discussion_r2311911351
##########
providers/common/io/tests/unit/common/io/xcom/test_backend.py:
##########
@@ -129,7 +135,25 @@ def test_value_storage(self, task_instance,
mock_supervisor_comms, session):
run_id=task_instance.run_id,
)
- if AIRFLOW_V_3_0_PLUS:
+ if AIRFLOW_V_3_1_PLUS:
+ XComModel.set(
+ key=XCOM_RETURN_KEY,
+ value=self.path,
+ dag_id=task_instance.dag_id,
+ task_id=task_instance.task_id,
+ run_id=task_instance.run_id,
+ )
+
+ res = session.execute(
+ XComModel.get_many(
+ key=XCOM_RETURN_KEY,
+ dag_ids=task_instance.dag_id,
+ task_ids=task_instance.task_id,
+ run_id=task_instance.run_id,
+ ).with_only_columns(XComModel.value)
+ ).first()
+ data = XComModel.deserialize_value(res)
Review Comment:
The `XComModel.set` is same for 3.0+ and 3.1+.
We only need to add if/ else with the `AIRFLOW_V_3_<0/1>_PLUS` for the
`XComModel.get_many` part.
##########
providers/common/io/tests/unit/common/io/xcom/test_backend.py:
##########
@@ -24,14 +24,20 @@
import airflow.models.xcom
from airflow.providers.common.io.xcom.backend import XComObjectStorageBackend
from airflow.providers.standard.operators.empty import EmptyOperator
-from airflow.utils import timezone
+
+try:
+ from airflow.utils import timezone # type: ignore[attr-defined]
+except ImportError:
+ from airflow.sdk import timezone
Review Comment:
```suggestion
try:
from airflow.sdk import timezone
except ImportError:
from airflow.utils import timezone # type: ignore[attr-defined,no-redef]
```
The fallback is on the opposite side.
##########
providers/common/io/tests/unit/common/io/xcom/test_backend.py:
##########
@@ -211,7 +245,31 @@ def test_clear(self, task_instance, session,
mock_supervisor_comms):
run_id=task_instance.run_id,
)
- if AIRFLOW_V_3_0_PLUS:
+ if AIRFLOW_V_3_1_PLUS:
Review Comment:
Same as above, only the `XComModel.get_many` need compatible if/else.
##########
providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py:
##########
@@ -112,7 +112,28 @@ def create_context(task, persist_to_db=False,
map_index=None):
dag = DAG(dag_id="dag", schedule=None, start_date=pendulum.now())
dag.add_task(task)
now = timezone.utcnow()
- if AIRFLOW_V_3_0_PLUS:
+ if AIRFLOW_V_3_1_PLUS:
+ with create_session() as session:
+ from airflow.models.dagbundle import DagBundleModel
+
+ bundle_name = "testing"
+ session.add(DagBundleModel(name=bundle_name))
+ session.flush()
+ session.add(DagModel(dag_id=dag.dag_id, bundle_name=bundle_name))
+ session.commit()
+ dag.sync_to_db()
+ SerializedDagModel.write_dag(dag, bundle_name="testing")
+ dag_run = DagRun(
+ run_id=DagRun.generate_run_id(
+ run_type=DagRunType.MANUAL, logical_date=DEFAULT_DATE,
run_after=DEFAULT_DATE
+ ),
+ run_type=DagRunType.MANUAL,
+ dag_id=dag.dag_id,
+ logical_date=now,
+ data_interval=(now, now),
+ run_after=now,
+ )
Review Comment:
It seems the `AIRFLOW_V_3_1_PLUS` block is exactly same as
`AIRFLOW_V_3_0_PLUS `.
Will there be CI failure if we remove the `AIRFLOW_V_3_1_PLUS` block ?
##########
providers/common/io/tests/unit/common/io/xcom/test_backend.py:
##########
@@ -24,14 +24,20 @@
import airflow.models.xcom
from airflow.providers.common.io.xcom.backend import XComObjectStorageBackend
from airflow.providers.standard.operators.empty import EmptyOperator
-from airflow.utils import timezone
+
+try:
+ from airflow.utils import timezone # type: ignore[attr-defined]
+except ImportError:
+ from airflow.sdk import timezone
from tests_common.test_utils import db
from tests_common.test_utils.config import conf_vars
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
XCOM_RETURN_KEY
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS, XCOM_RETURN_KEY
pytestmark = [pytest.mark.db_test]
+if AIRFLOW_V_3_1_PLUS:
+ from airflow.models.xcom import XComModel
Review Comment:
```suggestion
```
##########
providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py:
##########
@@ -55,9 +55,9 @@
from airflow.utils.types import DagRunType
from tests_common.test_utils import db
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS
-if AIRFLOW_V_3_0_PLUS:
+if AIRFLOW_V_3_0_PLUS or AIRFLOW_V_3_1_PLUS:
Review Comment:
```suggestion
if AIRFLOW_V_3_0_PLUS:
```
`AIRFLOW_V_3_0_PLUS` is defined as, in `AIRFLOW_V_3_1_PLUS` case, the
`AIRFLOW_V_3_0_PLUS` will also be `True`.
```python
AIRFLOW_V_3_0_PLUS: bool = get_base_airflow_version_tuple() >= (3, 0, 0)
```
##########
providers/common/io/tests/unit/common/io/xcom/test_backend.py:
##########
@@ -322,7 +403,25 @@ def test_compression(self, task_instance, session,
mock_supervisor_comms):
run_id=task_instance.run_id,
)
- if AIRFLOW_V_3_0_PLUS:
+ if AIRFLOW_V_3_1_PLUS:
+ XComModel.set(
+ key=XCOM_RETURN_KEY,
+ value=self.path + ".gz",
+ dag_id=task_instance.dag_id,
+ task_id=task_instance.task_id,
+ run_id=task_instance.run_id,
+ )
Review Comment:
Same as well.
##########
providers/common/io/tests/unit/common/io/xcom/test_backend.py:
##########
@@ -264,7 +322,30 @@ def test_clear(self, task_instance, session,
mock_supervisor_comms):
)
assert value
- if AIRFLOW_V_3_0_PLUS:
+ if AIRFLOW_V_3_1_PLUS:
+ mock_supervisor_comms.send.return_value =
XComResult(key=XCOM_RETURN_KEY, value=path)
+ XCom.delete(
+ dag_id=task_instance.dag_id,
+ task_id=task_instance.task_id,
+ run_id=task_instance.run_id,
+ key=XCOM_RETURN_KEY,
+ map_index=task_instance.map_index,
+ )
+ XComModel.clear(
+ dag_id=task_instance.dag_id,
+ task_id=task_instance.task_id,
+ run_id=task_instance.run_id,
+ map_index=task_instance.map_index,
+ )
Review Comment:
Same 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]
