You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ur...@apache.org on 2023/07/10 08:34:28 UTC
[airflow] branch main updated: Add xcom map_index as a filter to xcom endpoint (#32453)
This is an automated email from the ASF dual-hosted git repository.
uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new bc97646b26 Add xcom map_index as a filter to xcom endpoint (#32453)
bc97646b26 is described below
commit bc97646b262e7f338b4f3d4dce199e640e24e250
Author: Hussein Awala <hu...@awala.fr>
AuthorDate: Mon Jul 10 10:34:21 2023 +0200
Add xcom map_index as a filter to xcom endpoint (#32453)
---
airflow/api_connexion/endpoints/xcom_endpoint.py | 11 +-
airflow/api_connexion/openapi/v1.yaml | 13 ++
airflow/api_connexion/schemas/xcom_schema.py | 1 +
airflow/www/static/js/types/api-generated.ts | 9 +
.../api_connexion/endpoints/test_xcom_endpoint.py | 243 ++++++++++++++-------
tests/api_connexion/schemas/test_xcom_schema.py | 65 ++++--
tests/conftest.py | 2 +
7 files changed, 244 insertions(+), 100 deletions(-)
diff --git a/airflow/api_connexion/endpoints/xcom_endpoint.py b/airflow/api_connexion/endpoints/xcom_endpoint.py
index 81c0fe8f12..d95a3d15fc 100644
--- a/airflow/api_connexion/endpoints/xcom_endpoint.py
+++ b/airflow/api_connexion/endpoints/xcom_endpoint.py
@@ -49,6 +49,8 @@ def get_xcom_entries(
dag_id: str,
dag_run_id: str,
task_id: str,
+ map_index: int | None = None,
+ xcom_key: str | None = None,
limit: int | None,
offset: int | None = None,
session: Session = NEW_SESSION,
@@ -68,6 +70,10 @@ def get_xcom_entries(
query = query.where(XCom.task_id == task_id)
if dag_run_id != "~":
query = query.where(DR.run_id == dag_run_id)
+ if map_index is not None:
+ query = query.where(XCom.map_index == map_index)
+ if xcom_key is not None:
+ query = query.where(XCom.key == xcom_key)
query = query.order_by(DR.execution_date, XCom.task_id, XCom.dag_id, XCom.key)
total_entries = session.execute(select(func.count()).select_from(query)).scalar()
query = session.scalars(query.offset(offset).limit(limit))
@@ -89,6 +95,7 @@ def get_xcom_entry(
task_id: str,
dag_run_id: str,
xcom_key: str,
+ map_index: int = -1,
deserialize: bool = False,
session: Session = NEW_SESSION,
) -> APIResponse:
@@ -100,7 +107,9 @@ def get_xcom_entry(
else:
query = select(XCom)
- query = query.where(XCom.dag_id == dag_id, XCom.task_id == task_id, XCom.key == xcom_key)
+ query = query.where(
+ XCom.dag_id == dag_id, XCom.task_id == task_id, XCom.key == xcom_key, XCom.map_index == map_index
+ )
query = query.join(DR, and_(XCom.dag_id == DR.dag_id, XCom.run_id == DR.run_id))
query = query.where(DR.run_id == dag_run_id)
diff --git a/airflow/api_connexion/openapi/v1.yaml b/airflow/api_connexion/openapi/v1.yaml
index 3e1dcdedd7..6e3faa2f24 100644
--- a/airflow/api_connexion/openapi/v1.yaml
+++ b/airflow/api_connexion/openapi/v1.yaml
@@ -1573,6 +1573,8 @@ paths:
operationId: get_xcom_entries
tags: [XCom]
parameters:
+ - $ref: '#/components/parameters/FilterMapIndex'
+ - $ref: '#/components/parameters/FilterXcomKey'
- $ref: '#/components/parameters/PageLimit'
- $ref: '#/components/parameters/PageOffset'
responses:
@@ -1600,6 +1602,7 @@ paths:
operationId: get_xcom_entry
tags: [XCom]
parameters:
+ - $ref: '#/components/parameters/FilterMapIndex'
- in: query
name: deserialize
schema:
@@ -3489,6 +3492,8 @@ components:
execution_date:
type: string
format: datetime
+ map_index:
+ type: integer
task_id:
type: string
dag_id:
@@ -5095,6 +5100,14 @@ components:
*New in version 2.6.0*
+ FilterXcomKey:
+ in: query
+ name: xcom_key
+ schema:
+ type: string
+ required: false
+ description: Only filter the XCom records which have the provided key.
+
# Other parameters
FileToken:
in: path
diff --git a/airflow/api_connexion/schemas/xcom_schema.py b/airflow/api_connexion/schemas/xcom_schema.py
index 09d2505bf7..5894db8b1a 100644
--- a/airflow/api_connexion/schemas/xcom_schema.py
+++ b/airflow/api_connexion/schemas/xcom_schema.py
@@ -35,6 +35,7 @@ class XComCollectionItemSchema(SQLAlchemySchema):
key = auto_field()
timestamp = auto_field()
execution_date = auto_field()
+ map_index = auto_field()
task_id = auto_field()
dag_id = auto_field()
diff --git a/airflow/www/static/js/types/api-generated.ts b/airflow/www/static/js/types/api-generated.ts
index aa9bebf4d6..bcb18ad0e9 100644
--- a/airflow/www/static/js/types/api-generated.ts
+++ b/airflow/www/static/js/types/api-generated.ts
@@ -1425,6 +1425,7 @@ export interface components {
timestamp?: string;
/** Format: datetime */
execution_date?: string;
+ map_index?: number;
task_id?: string;
dag_id?: string;
};
@@ -2390,6 +2391,8 @@ export interface components {
* *New in version 2.6.0*
*/
Paused: boolean;
+ /** @description Only filter the XCom records which have the provided key. */
+ FilterXcomKey: string;
/**
* @description The key containing the encrypted path to the file. Encryption and decryption take place only on
* the server. This prevents the client from reading an non-DAG file. This also ensures API
@@ -3839,6 +3842,10 @@ export interface operations {
task_id: components["parameters"]["TaskID"];
};
query: {
+ /** Filter on map index for mapped task. */
+ map_index?: components["parameters"]["FilterMapIndex"];
+ /** Only filter the XCom records which have the provided key. */
+ xcom_key?: components["parameters"]["FilterXcomKey"];
/** The numbers of items to return. */
limit?: components["parameters"]["PageLimit"];
/** The number of items to skip before starting to collect the result set. */
@@ -3869,6 +3876,8 @@ export interface operations {
xcom_key: components["parameters"]["XComKey"];
};
query: {
+ /** Filter on map index for mapped task. */
+ map_index?: components["parameters"]["FilterMapIndex"];
/**
* Whether to deserialize an XCom value when using a custom XCom backend.
*
diff --git a/tests/api_connexion/endpoints/test_xcom_endpoint.py b/tests/api_connexion/endpoints/test_xcom_endpoint.py
index 7cd82fdaa3..0233197e24 100644
--- a/tests/api_connexion/endpoints/test_xcom_endpoint.py
+++ b/tests/api_connexion/endpoints/test_xcom_endpoint.py
@@ -82,6 +82,20 @@ def configured_app(minimal_app_for_api):
delete_user(app, username="test_no_permissions") # type: ignore
+def _compare_xcom_collections(collection1: dict, collection_2: dict):
+ assert collection1.get("total_entries") == collection_2.get("total_entries")
+ sort_key = lambda record: (
+ record.get("dag_id"),
+ record.get("task_id"),
+ record.get("execution_date"),
+ record.get("map_index"),
+ record.get("key"),
+ )
+ assert sorted(collection1.get("xcom_entries", []), key=sort_key) == sorted(
+ collection_2.get("xcom_entries", []), key=sort_key
+ )
+
+
class TestXComEndpoint:
@staticmethod
def clean_db():
@@ -128,6 +142,7 @@ class TestGetXComEntry(TestXComEndpoint):
"execution_date": execution_date,
"key": xcom_key,
"task_id": task_id,
+ "map_index": -1,
"timestamp": "TIMESTAMP",
"value": "TEST_VALUE",
}
@@ -258,25 +273,30 @@ class TestGetXComEntries(TestXComEndpoint):
response_data = response.json
for xcom_entry in response_data["xcom_entries"]:
xcom_entry["timestamp"] = "TIMESTAMP"
- assert response_data == {
- "xcom_entries": [
- {
- "dag_id": dag_id,
- "execution_date": execution_date,
- "key": "test-xcom-key-1",
- "task_id": task_id,
- "timestamp": "TIMESTAMP",
- },
- {
- "dag_id": dag_id,
- "execution_date": execution_date,
- "key": "test-xcom-key-2",
- "task_id": task_id,
- "timestamp": "TIMESTAMP",
- },
- ],
- "total_entries": 2,
- }
+ _compare_xcom_collections(
+ response_data,
+ {
+ "xcom_entries": [
+ {
+ "dag_id": dag_id,
+ "execution_date": execution_date,
+ "key": "test-xcom-key-1",
+ "task_id": task_id,
+ "timestamp": "TIMESTAMP",
+ "map_index": -1,
+ },
+ {
+ "dag_id": dag_id,
+ "execution_date": execution_date,
+ "key": "test-xcom-key-2",
+ "task_id": task_id,
+ "timestamp": "TIMESTAMP",
+ "map_index": -1,
+ },
+ ],
+ "total_entries": 2,
+ },
+ )
def test_should_respond_200_with_tilde_and_access_to_all_dags(self):
dag_id_1 = "test-dag-id-1"
@@ -300,39 +320,46 @@ class TestGetXComEntries(TestXComEndpoint):
response_data = response.json
for xcom_entry in response_data["xcom_entries"]:
xcom_entry["timestamp"] = "TIMESTAMP"
- assert response_data == {
- "xcom_entries": [
- {
- "dag_id": dag_id_1,
- "execution_date": execution_date,
- "key": "test-xcom-key-1",
- "task_id": task_id_1,
- "timestamp": "TIMESTAMP",
- },
- {
- "dag_id": dag_id_1,
- "execution_date": execution_date,
- "key": "test-xcom-key-2",
- "task_id": task_id_1,
- "timestamp": "TIMESTAMP",
- },
- {
- "dag_id": dag_id_2,
- "execution_date": execution_date,
- "key": "test-xcom-key-1",
- "task_id": task_id_2,
- "timestamp": "TIMESTAMP",
- },
- {
- "dag_id": dag_id_2,
- "execution_date": execution_date,
- "key": "test-xcom-key-2",
- "task_id": task_id_2,
- "timestamp": "TIMESTAMP",
- },
- ],
- "total_entries": 4,
- }
+ _compare_xcom_collections(
+ response_data,
+ {
+ "xcom_entries": [
+ {
+ "dag_id": dag_id_1,
+ "execution_date": execution_date,
+ "key": "test-xcom-key-1",
+ "task_id": task_id_1,
+ "timestamp": "TIMESTAMP",
+ "map_index": -1,
+ },
+ {
+ "dag_id": dag_id_1,
+ "execution_date": execution_date,
+ "key": "test-xcom-key-2",
+ "task_id": task_id_1,
+ "timestamp": "TIMESTAMP",
+ "map_index": -1,
+ },
+ {
+ "dag_id": dag_id_2,
+ "execution_date": execution_date,
+ "key": "test-xcom-key-1",
+ "task_id": task_id_2,
+ "timestamp": "TIMESTAMP",
+ "map_index": -1,
+ },
+ {
+ "dag_id": dag_id_2,
+ "execution_date": execution_date,
+ "key": "test-xcom-key-2",
+ "task_id": task_id_2,
+ "timestamp": "TIMESTAMP",
+ "map_index": -1,
+ },
+ ],
+ "total_entries": 4,
+ },
+ )
def test_should_respond_200_with_tilde_and_granular_dag_access(self):
dag_id_1 = "test-dag-id-1"
@@ -356,25 +383,74 @@ class TestGetXComEntries(TestXComEndpoint):
response_data = response.json
for xcom_entry in response_data["xcom_entries"]:
xcom_entry["timestamp"] = "TIMESTAMP"
- assert response_data == {
- "xcom_entries": [
- {
- "dag_id": dag_id_1,
- "execution_date": execution_date,
- "key": "test-xcom-key-1",
- "task_id": task_id_1,
- "timestamp": "TIMESTAMP",
- },
- {
- "dag_id": dag_id_1,
- "execution_date": execution_date,
- "key": "test-xcom-key-2",
- "task_id": task_id_1,
- "timestamp": "TIMESTAMP",
- },
- ],
- "total_entries": 2,
+ _compare_xcom_collections(
+ response_data,
+ {
+ "xcom_entries": [
+ {
+ "dag_id": dag_id_1,
+ "execution_date": execution_date,
+ "key": "test-xcom-key-1",
+ "task_id": task_id_1,
+ "timestamp": "TIMESTAMP",
+ "map_index": -1,
+ },
+ {
+ "dag_id": dag_id_1,
+ "execution_date": execution_date,
+ "key": "test-xcom-key-2",
+ "task_id": task_id_1,
+ "timestamp": "TIMESTAMP",
+ "map_index": -1,
+ },
+ ],
+ "total_entries": 2,
+ },
+ )
+
+ def test_should_respond_200_with_map_index(self):
+ dag_id = "test-dag-id"
+ task_id = "test-task-id"
+ execution_date = "2005-04-02T00:00:00+00:00"
+ execution_date_parsed = parse_execution_date(execution_date)
+ dag_run_id = DagRun.generate_run_id(DagRunType.MANUAL, execution_date_parsed)
+ self._create_xcom_entries(dag_id, dag_run_id, execution_date_parsed, task_id, mapped_ti=True)
+
+ def assert_expected_result(expected_entries, map_index=None):
+ response = self.client.get(
+ "/api/v1/dags/~/dagRuns/~/taskInstances/~/xcomEntries"
+ f"{('?map_index=' + str(map_index)) if map_index is not None else ''}",
+ environ_overrides={"REMOTE_USER": "test"},
+ )
+
+ assert 200 == response.status_code
+ response_data = response.json
+ for xcom_entry in response_data["xcom_entries"]:
+ xcom_entry["timestamp"] = "TIMESTAMP"
+ assert response_data == {
+ "xcom_entries": expected_entries,
+ "total_entries": len(expected_entries),
+ }
+
+ expected_entry1 = {
+ "dag_id": dag_id,
+ "execution_date": execution_date,
+ "key": "test-xcom-key",
+ "task_id": task_id,
+ "timestamp": "TIMESTAMP",
+ "map_index": 0,
+ }
+ expected_entry2 = {
+ "dag_id": dag_id,
+ "execution_date": execution_date,
+ "key": "test-xcom-key",
+ "task_id": task_id,
+ "timestamp": "TIMESTAMP",
+ "map_index": 1,
}
+ assert_expected_result([expected_entry1], map_index=0)
+ assert_expected_result([expected_entry2], map_index=1)
+ assert_expected_result([expected_entry1, expected_entry2], map_index=None)
def test_should_raises_401_unauthenticated(self):
dag_id = "test-dag-id"
@@ -390,7 +466,7 @@ class TestGetXComEntries(TestXComEndpoint):
assert_401(response)
- def _create_xcom_entries(self, dag_id, run_id, execution_date, task_id):
+ def _create_xcom_entries(self, dag_id, run_id, execution_date, task_id, mapped_ti=False):
with create_session() as session:
dag = DagModel(dag_id=dag_id)
session.add(dag)
@@ -402,17 +478,26 @@ class TestGetXComEntries(TestXComEndpoint):
run_type=DagRunType.MANUAL,
)
session.add(dagrun)
- ti = TaskInstance(EmptyOperator(task_id=task_id), run_id=run_id)
- ti.dag_id = dag_id
- session.add(ti)
+ if mapped_ti:
+ for i in [0, 1]:
+ ti = TaskInstance(EmptyOperator(task_id=task_id), run_id=run_id, map_index=i)
+ ti.dag_id = dag_id
+ session.add(ti)
+ else:
+ ti = TaskInstance(EmptyOperator(task_id=task_id), run_id=run_id)
+ ti.dag_id = dag_id
+ session.add(ti)
for i in [1, 2]:
+ if mapped_ti:
+ key = "test-xcom-key"
+ map_index = i - 1
+ else:
+ key = f"test-xcom-key-{i}"
+ map_index = -1
+
XCom.set(
- key=f"test-xcom-key-{i}",
- value="TEST",
- run_id=run_id,
- task_id=task_id,
- dag_id=dag_id,
+ key=key, value="TEST", run_id=run_id, task_id=task_id, dag_id=dag_id, map_index=map_index
)
def _create_invalid_xcom_entries(self, execution_date):
diff --git a/tests/api_connexion/schemas/test_xcom_schema.py b/tests/api_connexion/schemas/test_xcom_schema.py
index c8f96d1475..5762086086 100644
--- a/tests/api_connexion/schemas/test_xcom_schema.py
+++ b/tests/api_connexion/schemas/test_xcom_schema.py
@@ -39,19 +39,35 @@ def clean_xcom():
session.query(XCom).delete()
+def _compare_xcom_collections(collection1: dict, collection_2: dict):
+ assert collection1.get("total_entries") == collection_2.get("total_entries")
+ sort_key = lambda record: (
+ record.get("dag_id"),
+ record.get("task_id"),
+ record.get("execution_date"),
+ record.get("map_index"),
+ record.get("key"),
+ )
+ assert sorted(collection1.get("xcom_entries", []), key=sort_key) == sorted(
+ collection_2.get("xcom_entries", []), key=sort_key
+ )
+
+
@pytest.fixture()
def create_xcom(create_task_instance, session):
- def maker(dag_id, task_id, execution_date, key, value=None):
+ def maker(dag_id, task_id, execution_date, key, map_index=-1, value=None):
ti = create_task_instance(
dag_id=dag_id,
task_id=task_id,
execution_date=execution_date,
+ map_index=map_index,
session=session,
)
run: DagRun = ti.dag_run
xcom = XCom(
dag_run_id=run.id,
task_id=ti.task_id,
+ map_index=map_index,
key=key,
value=value,
timestamp=run.execution_date,
@@ -84,6 +100,7 @@ class TestXComCollectionItemSchema:
"execution_date": self.default_time,
"task_id": "test_task_id",
"dag_id": "test_dag",
+ "map_index": -1,
}
def test_deserialize(self):
@@ -93,6 +110,7 @@ class TestXComCollectionItemSchema:
"execution_date": self.default_time,
"task_id": "test_task_id",
"dag_id": "test_dag",
+ "map_index": 2,
}
result = xcom_collection_item_schema.load(xcom_dump)
assert result == {
@@ -101,6 +119,7 @@ class TestXComCollectionItemSchema:
"execution_date": self.default_time_parsed,
"task_id": "test_task_id",
"dag_id": "test_dag",
+ "map_index": 2,
}
@@ -133,25 +152,30 @@ class TestXComCollectionSchema:
total_entries=xcom_models_query.count(),
)
)
- assert deserialized_xcoms == {
- "xcom_entries": [
- {
- "key": "test_key_1",
- "timestamp": self.default_time_1,
- "execution_date": self.default_time_1,
- "task_id": "test_task_id_1",
- "dag_id": "test_dag_1",
- },
- {
- "key": "test_key_2",
- "timestamp": self.default_time_2,
- "execution_date": self.default_time_2,
- "task_id": "test_task_id_2",
- "dag_id": "test_dag_2",
- },
- ],
- "total_entries": 2,
- }
+ _compare_xcom_collections(
+ deserialized_xcoms,
+ {
+ "xcom_entries": [
+ {
+ "key": "test_key_1",
+ "timestamp": self.default_time_1,
+ "execution_date": self.default_time_1,
+ "task_id": "test_task_id_1",
+ "dag_id": "test_dag_1",
+ "map_index": -1,
+ },
+ {
+ "key": "test_key_2",
+ "timestamp": self.default_time_2,
+ "execution_date": self.default_time_2,
+ "task_id": "test_task_id_2",
+ "dag_id": "test_dag_2",
+ "map_index": -1,
+ },
+ ],
+ "total_entries": 2,
+ },
+ )
class TestXComSchema:
@@ -175,6 +199,7 @@ class TestXComSchema:
"task_id": "test_task_id",
"dag_id": "test_dag",
"value": "test_binary",
+ "map_index": -1,
}
def test_deserialize(self):
diff --git a/tests/conftest.py b/tests/conftest.py
index ca0751f1be..4c44fd9810 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -761,6 +761,7 @@ def create_task_instance(dag_maker, create_dummy_dag):
run_id=None,
run_type=None,
data_interval=None,
+ map_index=-1,
**kwargs,
) -> TaskInstance:
if execution_date is None:
@@ -780,6 +781,7 @@ def create_task_instance(dag_maker, create_dummy_dag):
(ti,) = dagrun.task_instances
ti.task = task
ti.state = state
+ ti.map_index = map_index
dag_maker.session.flush()
return ti