justinpakzad commented on code in PR #59874:
URL: https://github.com/apache/airflow/pull/59874#discussion_r3170841263
##########
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:
Yea, good call. I made that update. The endpoint now returns an
`XComDeleteCountResponse` with the number of deleted rows. Updated everything
else accordingly.
##########
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:
Good find. Fixed.
##########
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:
Added.
--
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]
