amoghrajesh commented on code in PR #59874:
URL: https://github.com/apache/airflow/pull/59874#discussion_r3165988964
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -568,6 +568,30 @@ def delete(
# decouple from the server response string
return OKResponse(ok=True)
+ def delete_all(
+ self,
+ dag_id: str,
+ run_id: str,
+ task_id: str | None = None,
+ key: str | None = None,
+ map_index: int | None = None,
+ ):
Review Comment:
```suggestion
) -> OKResponse :
```
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -1732,6 +1733,27 @@ class RequestTestCase:
),
test_id="delete_xcom",
),
+ RequestTestCase(
+ message=BulkDeleteXCom(
+ dag_id="test_dag",
+ run_id="test_run",
+ ),
+ client_mock=ClientMock(
+ method_path="xcoms.delete_all",
+ args=("test_dag", "test_run", None, None, None),
+ response=OKResponse(ok=True),
+ ),
+ test_id="bulk_delete_xcoms",
+ ),
+ RequestTestCase(
+ message=BulkDeleteXCom(dag_id="test_dag", run_id="test_run",
task_id="t1"),
+ client_mock=ClientMock(
+ method_path="xcoms.delete_all",
+ args=("test_dag", "test_run", "t1", None, None),
+ response=OKResponse(ok=True),
+ ),
+ test_id="bulk_delete_xcoms",
+ ),
Review Comment:
Can you also add in a test for `map_index` variation?
##########
task-sdk/src/airflow/sdk/api/client.py:
##########
@@ -568,6 +568,30 @@ def delete(
# decouple from the server response string
return OKResponse(ok=True)
+ def delete_all(
+ self,
+ dag_id: str,
+ run_id: str,
+ task_id: str | None = None,
+ key: str | None = None,
+ map_index: int | None = None,
+ ):
+ """Bulk delete XCom values via the API server."""
+ params: dict[str, str | int] = {}
+
+ if map_index is not None and map_index >= 0:
+ params["map_index"] = map_index
+ if task_id is not None:
+ params["task_id"] = task_id
+ if key is not None:
+ params["key"] = key
+
+ self.client.delete(url=f"xcoms/{dag_id}/{run_id}", params=params)
+ # Any error from the server will anyway be propagated down to the
supervisor,
+ # so we choose to send a generic response to the supervisor over the
server response to
+ # decouple from the server response string
+ return OKResponse(ok=True)
Review Comment:
Wonder if it will be beneficial to return number of deleted rows instead?
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -1732,6 +1733,27 @@ class RequestTestCase:
),
test_id="delete_xcom",
),
+ RequestTestCase(
+ message=BulkDeleteXCom(
+ dag_id="test_dag",
+ run_id="test_run",
+ ),
+ client_mock=ClientMock(
+ method_path="xcoms.delete_all",
+ args=("test_dag", "test_run", None, None, None),
+ response=OKResponse(ok=True),
+ ),
+ test_id="bulk_delete_xcoms",
+ ),
+ RequestTestCase(
+ message=BulkDeleteXCom(dag_id="test_dag", run_id="test_run",
task_id="t1"),
+ client_mock=ClientMock(
+ method_path="xcoms.delete_all",
+ args=("test_dag", "test_run", "t1", None, None),
+ response=OKResponse(ok=True),
+ ),
+ test_id="bulk_delete_xcoms",
+ ),
Review Comment:
Both have same test_id
--
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]
