You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "omkar-foss (via GitHub)" <gi...@apache.org> on 2024/11/20 19:28:01 UTC
[PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
omkar-foss opened a new pull request, #44223:
URL: https://github.com/apache/airflow/pull/44223
<!--
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied. See the License for the
specific language governing permissions and limitations
under the License.
-->
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
closes: https://github.com/apache/airflow/issues/43753
related: https://github.com/apache/airflow/issues/42370
This migrates the Patch Task Instance API from `api_connexion` to `api_fastapi`.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on code in PR #44223:
URL: https://github.com/apache/airflow/pull/44223#discussion_r1852577668
##########
airflow/api_fastapi/core_api/datamodels/task_instances.py:
##########
@@ -150,3 +152,41 @@ class TaskInstanceHistoryCollectionResponse(BaseModel):
task_instances: list[TaskInstanceHistoryResponse]
total_entries: int
+
+
+class PatchTaskInstanceBody(BaseModel):
+ """Request body for Clear Task Instances endpoint."""
+
+ dry_run: bool = True
+ new_state: str | None = None
+ note: str | None = None
Review Comment:
Yes that one.
Thanks for you input @Kludex :)
##########
airflow/api_fastapi/core_api/datamodels/task_instances.py:
##########
@@ -150,3 +152,41 @@ class TaskInstanceHistoryCollectionResponse(BaseModel):
task_instances: list[TaskInstanceHistoryResponse]
total_entries: int
+
+
+class PatchTaskInstanceBody(BaseModel):
+ """Request body for Clear Task Instances endpoint."""
+
+ dry_run: bool = True
+ new_state: str | None = None
+ note: str | None = None
Review Comment:
Yes that one.
Thanks for your input @Kludex :)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on code in PR #44223:
URL: https://github.com/apache/airflow/pull/44223#discussion_r1854023921
##########
airflow/api_fastapi/core_api/datamodels/task_instances.py:
##########
@@ -190,6 +192,33 @@ def validate_model(cls, data: Any) -> Any:
return data
+class PatchTaskInstanceBody(BaseModel):
+ """Request body for Clear Task Instances endpoint."""
+
+ dry_run: bool = True
+ new_state: str | None = None
+ note: Annotated[str, StringConstraints(max_length=1000)] | None = None
+ include_upstream: bool = False
Review Comment:
Defaulting these to `False` as per your suggestion @bbovenzi, jfyi.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on PR #44223:
URL: https://github.com/apache/airflow/pull/44223#issuecomment-2498531074
@pierrejeambrun all checks passing, and have rebased with `main` just now. Please review when you get a chance, thanks.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on code in PR #44223:
URL: https://github.com/apache/airflow/pull/44223#discussion_r1853924645
##########
airflow/api_fastapi/core_api/routes/public/task_instances.py:
##########
@@ -482,3 +484,84 @@ def get_mapped_task_instance_try_details(
map_index=map_index,
session=session,
)
+
+
+@task_instances_router.patch(
+ "/{task_id}",
+ responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST]),
+)
+@task_instances_router.patch(
+ "/{task_id}/{map_index}",
+ responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST]),
+)
+def patch_task_instance(
+ dag_id: str,
+ dag_run_id: str,
+ task_id: str,
+ request: Request,
+ body: PatchTaskInstanceBody,
+ session: Annotated[Session, Depends(get_session)],
+ map_index: int = -1,
+ update_mask: list[str] | None = Query(None),
+) -> TaskInstanceResponse:
+ """Update the state of a task instance."""
+ dag = request.app.state.dag_bag.get_dag(dag_id)
+ if not dag:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, f"DAG {dag_id} not found")
+
+ if not dag.has_task(task_id):
+ raise HTTPException(status.HTTP_404_NOT_FOUND, f"Task '{task_id}' not found in DAG '{dag_id}'")
+
+ query = (
+ select(TI)
+ .where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id)
+ .join(TI.dag_run)
+ .options(joinedload(TI.rendered_task_instance_fields))
+ )
+ if map_index == -1:
+ query = query.where(or_(TI.map_index == -1, TI.map_index is None))
+ else:
+ query = query.where(TI.map_index == map_index)
+
+ try:
+ ti = session.scalar(query)
+ except MultipleResultsFound:
+ raise HTTPException(
+ status.HTTP_400_BAD_REQUEST,
+ "Multiple task instances found. As the TI is mapped, add the map_index value to the URL",
+ )
+
+ err_msg_404 = f"Task Instance not found for dag_id={dag_id}, run_id={dag_run_id}, task_id={task_id}"
+ if ti is None:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, err_msg_404)
+
+ fields_to_update = body.model_fields_set
+ if update_mask:
+ fields_to_update = fields_to_update.intersection(update_mask)
+
+ for field in fields_to_update:
+ if field == "new_state":
+ if not body.dry_run:
+ tis: list[TI] = dag.set_task_instance_state(
Review Comment:
@pierrejeambrun I'm adding these anyway as the legacy Set TI State API has those. We can remove if not required.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "Kludex (via GitHub)" <gi...@apache.org>.
Kludex commented on code in PR #44223:
URL: https://github.com/apache/airflow/pull/44223#discussion_r1852470729
##########
airflow/api_fastapi/core_api/datamodels/task_instances.py:
##########
@@ -150,3 +152,41 @@ class TaskInstanceHistoryCollectionResponse(BaseModel):
task_instances: list[TaskInstanceHistoryResponse]
total_entries: int
+
+
+class PatchTaskInstanceBody(BaseModel):
+ """Request body for Clear Task Instances endpoint."""
+
+ dry_run: bool = True
+ new_state: str | None = None
+ note: str | None = None
+
+ @model_validator(mode="before")
+ @classmethod
+ def validate_model(cls, data: Any) -> Any:
+ if data.get("note") is None and data.get("new_state") is None:
+ raise ValueError("new_state is required.")
+ return data
+
+ @field_validator("note", mode="before")
+ @classmethod
+ def validate_note(cls, note: str | None) -> str | None:
+ """Validate note."""
+ if note is None:
+ return None
+ if len(note) > 1000:
+ raise ValueError("Note length should not exceed 1000 characters.")
Review Comment:
Shared the link above.
##########
airflow/api_fastapi/core_api/datamodels/task_instances.py:
##########
@@ -150,3 +152,41 @@ class TaskInstanceHistoryCollectionResponse(BaseModel):
task_instances: list[TaskInstanceHistoryResponse]
total_entries: int
+
+
+class PatchTaskInstanceBody(BaseModel):
+ """Request body for Clear Task Instances endpoint."""
+
+ dry_run: bool = True
+ new_state: str | None = None
+ note: str | None = None
Review Comment:
https://docs.pydantic.dev/latest/api/types/#pydantic.types.StringConstraints
```suggestion
note: Annotated[str, StringConstraints(max_length=1000)] | None = None
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on code in PR #44223:
URL: https://github.com/apache/airflow/pull/44223#discussion_r1854022880
##########
airflow/api_fastapi/core_api/routes/public/task_instances.py:
##########
@@ -482,3 +484,84 @@ def get_mapped_task_instance_try_details(
map_index=map_index,
session=session,
)
+
+
+@task_instances_router.patch(
+ "/{task_id}",
+ responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST]),
+)
+@task_instances_router.patch(
+ "/{task_id}/{map_index}",
+ responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST]),
+)
+def patch_task_instance(
+ dag_id: str,
+ dag_run_id: str,
+ task_id: str,
+ request: Request,
+ body: PatchTaskInstanceBody,
+ session: Annotated[Session, Depends(get_session)],
+ map_index: int = -1,
+ update_mask: list[str] | None = Query(None),
+) -> TaskInstanceResponse:
+ """Update the state of a task instance."""
+ dag = request.app.state.dag_bag.get_dag(dag_id)
+ if not dag:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, f"DAG {dag_id} not found")
+
+ if not dag.has_task(task_id):
+ raise HTTPException(status.HTTP_404_NOT_FOUND, f"Task '{task_id}' not found in DAG '{dag_id}'")
+
+ query = (
+ select(TI)
+ .where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id)
+ .join(TI.dag_run)
+ .options(joinedload(TI.rendered_task_instance_fields))
+ )
+ if map_index == -1:
+ query = query.where(or_(TI.map_index == -1, TI.map_index is None))
+ else:
+ query = query.where(TI.map_index == map_index)
+
+ try:
+ ti = session.scalar(query)
+ except MultipleResultsFound:
+ raise HTTPException(
+ status.HTTP_400_BAD_REQUEST,
+ "Multiple task instances found. As the TI is mapped, add the map_index value to the URL",
+ )
+
+ err_msg_404 = f"Task Instance not found for dag_id={dag_id}, run_id={dag_run_id}, task_id={task_id}"
+ if ti is None:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, err_msg_404)
+
+ fields_to_update = body.model_fields_set
+ if update_mask:
+ fields_to_update = fields_to_update.intersection(update_mask)
+
+ for field in fields_to_update:
+ if field == "new_state":
+ if not body.dry_run:
+ tis: list[TI] = dag.set_task_instance_state(
Review Comment:
This is done.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on code in PR #44223:
URL: https://github.com/apache/airflow/pull/44223#discussion_r1852965805
##########
airflow/api_fastapi/core_api/datamodels/task_instances.py:
##########
@@ -150,3 +152,41 @@ class TaskInstanceHistoryCollectionResponse(BaseModel):
task_instances: list[TaskInstanceHistoryResponse]
total_entries: int
+
+
+class PatchTaskInstanceBody(BaseModel):
+ """Request body for Clear Task Instances endpoint."""
+
+ dry_run: bool = True
+ new_state: str | None = None
+ note: str | None = None
+
+ @model_validator(mode="before")
+ @classmethod
+ def validate_model(cls, data: Any) -> Any:
+ if data.get("note") is None and data.get("new_state") is None:
+ raise ValueError("new_state is required.")
+ return data
+
+ @field_validator("note", mode="before")
+ @classmethod
+ def validate_note(cls, note: str | None) -> str | None:
+ """Validate note."""
+ if note is None:
+ return None
+ if len(note) > 1000:
+ raise ValueError("Note length should not exceed 1000 characters.")
Review Comment:
Done, updated PR. Thanks a lot for the suggestion @Kludex! :)
##########
airflow/api_fastapi/core_api/routes/public/task_instances.py:
##########
@@ -482,3 +484,78 @@ def get_mapped_task_instance_try_details(
map_index=map_index,
session=session,
)
+
+
+@task_instances_router.patch(
+ "/{task_id}",
+ responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST]),
+)
+@task_instances_router.patch(
+ "/{task_id}/{map_index}",
+ responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST]),
+)
+def patch_task_instance(
+ dag_id: str,
+ dag_run_id: str,
+ task_id: str,
+ request: Request,
+ body: PatchTaskInstanceBody,
+ session: Annotated[Session, Depends(get_session)],
+ map_index: int = -1,
Review Comment:
Done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss closed pull request #44223: Migrate public endpoint Patch Task Instance to FastAPI
URL: https://github.com/apache/airflow/pull/44223
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on code in PR #44223:
URL: https://github.com/apache/airflow/pull/44223#discussion_r1852966479
##########
airflow/api_fastapi/core_api/datamodels/task_instances.py:
##########
@@ -150,3 +152,41 @@ class TaskInstanceHistoryCollectionResponse(BaseModel):
task_instances: list[TaskInstanceHistoryResponse]
total_entries: int
+
+
+class PatchTaskInstanceBody(BaseModel):
+ """Request body for Clear Task Instances endpoint."""
+
+ dry_run: bool = True
+ new_state: str | None = None
+ note: str | None = None
Review Comment:
Done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on PR #44223:
URL: https://github.com/apache/airflow/pull/44223#issuecomment-2500061767
@pierrejeambrun PR rebased with `main` and ready to merge. Thank you.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun merged PR #44223:
URL: https://github.com/apache/airflow/pull/44223
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on PR #44223:
URL: https://github.com/apache/airflow/pull/44223#issuecomment-2490124978
@pierrejeambrun I'm adding the Set Task Instance Note functionality to this (Patch Task Instance) API as you mentioned [here](https://github.com/apache/airflow/issues/43755#issuecomment-2459950793). Will update this PR in a while.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on code in PR #44223:
URL: https://github.com/apache/airflow/pull/44223#discussion_r1852986783
##########
airflow/api_fastapi/core_api/routes/public/task_instances.py:
##########
@@ -482,3 +484,84 @@ def get_mapped_task_instance_try_details(
map_index=map_index,
session=session,
)
+
+
+@task_instances_router.patch(
+ "/{task_id}",
+ responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST]),
+)
+@task_instances_router.patch(
+ "/{task_id}/{map_index}",
+ responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST]),
+)
+def patch_task_instance(
+ dag_id: str,
+ dag_run_id: str,
+ task_id: str,
+ request: Request,
+ body: PatchTaskInstanceBody,
+ session: Annotated[Session, Depends(get_session)],
+ map_index: int = -1,
+ update_mask: list[str] | None = Query(None),
+) -> TaskInstanceResponse:
+ """Update the state of a task instance."""
+ dag = request.app.state.dag_bag.get_dag(dag_id)
+ if not dag:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, f"DAG {dag_id} not found")
+
+ if not dag.has_task(task_id):
+ raise HTTPException(status.HTTP_404_NOT_FOUND, f"Task '{task_id}' not found in DAG '{dag_id}'")
+
+ query = (
+ select(TI)
+ .where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id)
+ .join(TI.dag_run)
+ .options(joinedload(TI.rendered_task_instance_fields))
+ )
+ if map_index == -1:
+ query = query.where(or_(TI.map_index == -1, TI.map_index is None))
+ else:
+ query = query.where(TI.map_index == map_index)
+
+ try:
+ ti = session.scalar(query)
+ except MultipleResultsFound:
+ raise HTTPException(
+ status.HTTP_400_BAD_REQUEST,
+ "Multiple task instances found. As the TI is mapped, add the map_index value to the URL",
+ )
+
+ err_msg_404 = f"Task Instance not found for dag_id={dag_id}, run_id={dag_run_id}, task_id={task_id}"
+ if ti is None:
+ raise HTTPException(status.HTTP_404_NOT_FOUND, err_msg_404)
+
+ fields_to_update = body.model_fields_set
+ if update_mask:
+ fields_to_update = fields_to_update.intersection(update_mask)
+
+ for field in fields_to_update:
+ if field == "new_state":
+ if not body.dry_run:
+ tis: list[TI] = dag.set_task_instance_state(
Review Comment:
Are these to be included for updation as well? Let me know if required I'll add these too.
```python
upstream=body.include_upstream,
downstream=body.include_downstream,
future=body.include_future,
past=body.include_past,
...
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on code in PR #44223:
URL: https://github.com/apache/airflow/pull/44223#discussion_r1857267249
##########
tests/api_fastapi/core_api/routes/public/test_task_instances.py:
##########
@@ -2318,3 +2319,542 @@ def test_raises_404_for_non_existent_dag(self, test_client):
)
assert response.status_code == 404
assert "DAG non-existent-dag not found" in response.text
+
+
+class TestPatchTaskInstance(TestTaskInstanceEndpoint):
+ ENDPOINT_URL = (
+ "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context"
+ )
+
+ @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
+ def test_should_call_mocked_api(self, mock_set_task_instance_state, test_client, session):
+ self.create_task_instances(session)
+
+ NEW_STATE = "failed"
+ mock_set_task_instance_state.return_value = session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == "example_python_operator",
+ TaskInstance.task_id == "print_the_context",
+ TaskInstance.run_id == "TEST_DAG_RUN_ID",
+ TaskInstance.map_index == -1,
+ )
+ ).one_or_none()
+
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dry_run": False,
+ "new_state": NEW_STATE,
+ },
+ )
+ assert response.status_code == 200
+ assert response.json() == {
+ "dag_id": "example_python_operator",
+ "dag_run_id": "TEST_DAG_RUN_ID",
+ "logical_date": "2020-01-01T00:00:00Z",
+ "task_id": "print_the_context",
+ "duration": 10000.0,
+ "end_date": "2020-01-03T00:00:00Z",
+ "executor": None,
+ "executor_config": "{}",
+ "hostname": "",
+ "id": mock.ANY,
+ "map_index": -1,
+ "max_tries": 0,
+ "note": "placeholder-note",
+ "operator": "PythonOperator",
+ "pid": 100,
+ "pool": "default_pool",
+ "pool_slots": 1,
+ "priority_weight": 9,
+ "queue": "default_queue",
+ "queued_when": None,
+ "start_date": "2020-01-02T00:00:00Z",
+ "state": "running",
+ "task_display_name": "print_the_context",
+ "try_number": 0,
+ "unixname": getuser(),
+ "rendered_fields": {},
+ "rendered_map_index": None,
+ "trigger": None,
+ "triggerer_job": None,
+ }
+
+ mock_set_task_instance_state.assert_called_once()
Review Comment:
Done
##########
tests/api_fastapi/core_api/routes/public/test_task_instances.py:
##########
@@ -2318,3 +2319,542 @@ def test_raises_404_for_non_existent_dag(self, test_client):
)
assert response.status_code == 404
assert "DAG non-existent-dag not found" in response.text
+
+
+class TestPatchTaskInstance(TestTaskInstanceEndpoint):
+ ENDPOINT_URL = (
+ "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context"
+ )
+
+ @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
+ def test_should_call_mocked_api(self, mock_set_task_instance_state, test_client, session):
+ self.create_task_instances(session)
+
+ NEW_STATE = "failed"
+ mock_set_task_instance_state.return_value = session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == "example_python_operator",
+ TaskInstance.task_id == "print_the_context",
+ TaskInstance.run_id == "TEST_DAG_RUN_ID",
+ TaskInstance.map_index == -1,
+ )
+ ).one_or_none()
+
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dry_run": False,
+ "new_state": NEW_STATE,
+ },
+ )
+ assert response.status_code == 200
+ assert response.json() == {
+ "dag_id": "example_python_operator",
+ "dag_run_id": "TEST_DAG_RUN_ID",
+ "logical_date": "2020-01-01T00:00:00Z",
+ "task_id": "print_the_context",
+ "duration": 10000.0,
+ "end_date": "2020-01-03T00:00:00Z",
+ "executor": None,
+ "executor_config": "{}",
+ "hostname": "",
+ "id": mock.ANY,
+ "map_index": -1,
+ "max_tries": 0,
+ "note": "placeholder-note",
+ "operator": "PythonOperator",
+ "pid": 100,
+ "pool": "default_pool",
+ "pool_slots": 1,
+ "priority_weight": 9,
+ "queue": "default_queue",
+ "queued_when": None,
+ "start_date": "2020-01-02T00:00:00Z",
+ "state": "running",
+ "task_display_name": "print_the_context",
+ "try_number": 0,
+ "unixname": getuser(),
+ "rendered_fields": {},
+ "rendered_map_index": None,
+ "trigger": None,
+ "triggerer_job": None,
+ }
+
+ mock_set_task_instance_state.assert_called_once()
+
+ @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
+ def test_should_not_call_mocked_api_for_dry_run(self, mock_set_task_instance_state, test_client, session):
+ self.create_task_instances(session)
+
+ NEW_STATE = "failed"
+ mock_set_task_instance_state.return_value = session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == "example_python_operator",
+ TaskInstance.task_id == "print_the_context",
+ TaskInstance.run_id == "TEST_DAG_RUN_ID",
+ TaskInstance.map_index == -1,
+ )
+ ).one_or_none()
+
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dry_run": True,
+ "new_state": NEW_STATE,
+ },
+ )
+ assert response.status_code == 200
+ assert response.json() == {
+ "dag_id": "example_python_operator",
+ "dag_run_id": "TEST_DAG_RUN_ID",
+ "logical_date": "2020-01-01T00:00:00Z",
+ "task_id": "print_the_context",
+ "duration": 10000.0,
+ "end_date": "2020-01-03T00:00:00Z",
+ "executor": None,
+ "executor_config": "{}",
+ "hostname": "",
+ "id": mock.ANY,
+ "map_index": -1,
+ "max_tries": 0,
+ "note": "placeholder-note",
+ "operator": "PythonOperator",
+ "pid": 100,
+ "pool": "default_pool",
+ "pool_slots": 1,
+ "priority_weight": 9,
+ "queue": "default_queue",
+ "queued_when": None,
+ "start_date": "2020-01-02T00:00:00Z",
+ "state": "running",
+ "task_display_name": "print_the_context",
+ "try_number": 0,
+ "unixname": getuser(),
+ "rendered_fields": {},
+ "rendered_map_index": None,
+ "trigger": None,
+ "triggerer_job": None,
+ }
+
+ mock_set_task_instance_state.assert_not_called()
+
+ def test_should_update_task_instance_state(self, test_client, session):
+ self.create_task_instances(session)
+
+ NEW_STATE = "failed"
+
+ test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dry_run": False,
+ "new_state": NEW_STATE,
+ },
+ )
+
+ response2 = test_client.get(self.ENDPOINT_URL)
+ assert response2.status_code == 200
+ assert response2.json()["state"] == NEW_STATE
+
+ def test_should_update_task_instance_state_default_dry_run_to_true(self, test_client, session):
+ self.create_task_instances(session)
+
+ NEW_STATE = "running"
+
+ test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "new_state": NEW_STATE,
+ },
+ )
+
+ response2 = test_client.get(self.ENDPOINT_URL)
+ assert response2.status_code == 200
+ assert response2.json()["state"] == NEW_STATE
+
+ def test_should_update_mapped_task_instance_state(self, test_client, session):
+ NEW_STATE = "failed"
+ map_index = 1
+ tis = self.create_task_instances(session)
+ ti = TaskInstance(task=tis[0].task, run_id=tis[0].run_id, map_index=map_index)
+ ti.rendered_task_instance_fields = RTIF(ti, render_templates=False)
+ session.add(ti)
+ session.commit()
+
+ response = test_client.patch(
+ f"{self.ENDPOINT_URL}/{map_index}",
+ json={
+ "dry_run": False,
+ "new_state": NEW_STATE,
+ },
+ )
+ assert response.status_code == 200
+
+ response2 = test_client.get(f"{self.ENDPOINT_URL}/{map_index}")
+ assert response2.status_code == 200
+ assert response2.json()["state"] == NEW_STATE
+
+ @pytest.mark.parametrize(
+ "error, code, payload",
+ [
+ [
+ (
+ "Task Instance not found for dag_id=example_python_operator"
+ ", run_id=TEST_DAG_RUN_ID, task_id=print_the_context"
+ ),
+ 404,
+ {
+ "dry_run": True,
+ "new_state": "failed",
+ },
+ ]
+ ],
+ )
+ def test_should_handle_errors(self, error, code, payload, test_client, session):
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json=payload,
+ )
+ assert response.status_code == code
+ assert response.json()["detail"] == error
+
+ def test_should_200_for_unknown_fields(self, test_client, session):
+ self.create_task_instances(session)
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dryrun": True,
+ "new_state": "failed",
+ },
+ )
+ assert response.status_code == 200
+
+ def test_should_raise_404_for_non_existent_dag(self, test_client):
+ response = test_client.patch(
+ "/public/dags/non-existent-dag/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context",
+ json={
+ "dry_run": False,
+ "new_state": "failed",
+ },
+ )
+ assert response.status_code == 404
+ assert response.json() == {"detail": "DAG non-existent-dag not found"}
+
+ def test_should_raise_404_for_non_existent_task_in_dag(self, test_client):
+ response = test_client.patch(
+ "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/non_existent_task",
+ json={
+ "dry_run": False,
+ "new_state": "failed",
+ },
+ )
+ assert response.status_code == 404
+ assert response.json() == {
+ "detail": "Task 'non_existent_task' not found in DAG 'example_python_operator'"
+ }
+
+ def test_should_raise_404_not_found_dag(self, test_client):
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dry_run": True,
+ "new_state": "failed",
+ },
+ )
+ assert response.status_code == 404
+
+ def test_should_raise_404_not_found_task(self, test_client):
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dry_run": True,
+ "new_state": "failed",
+ },
+ )
+ assert response.status_code == 404
+
+ @pytest.mark.parametrize(
+ "payload, expected",
+ [
+ (
+ {
+ "dry_run": True,
+ "new_state": "failede",
+ },
+ f"'failede' is not one of ['{State.SUCCESS}', '{State.FAILED}', '{State.SKIPPED}']",
+ ),
+ (
+ {
+ "dry_run": True,
+ "new_state": "queued",
+ },
+ f"'queued' is not one of ['{State.SUCCESS}', '{State.FAILED}', '{State.SKIPPED}']",
+ ),
+ ],
+ )
+ def test_should_raise_422_for_invalid_task_instance_state(self, payload, expected, test_client, session):
+ self.create_task_instances(session)
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json=payload,
+ )
+ assert response.status_code == 422
+ assert response.json() == {
+ "detail": [
+ {
+ "type": "value_error",
+ "loc": ["body", "new_state"],
+ "msg": f"Value error, {expected}",
+ "input": payload["new_state"],
+ "ctx": {"error": {}},
+ }
+ ]
+ }
+
+ @pytest.mark.parametrize(
+ "new_state,expected_status_code,expected_json,set_ti_state_call_count",
+ [
+ (
+ "failed",
+ 200,
+ {
+ "dag_id": "example_python_operator",
+ "dag_run_id": "TEST_DAG_RUN_ID",
+ "logical_date": "2020-01-01T00:00:00Z",
+ "task_id": "print_the_context",
+ "duration": 10000.0,
+ "end_date": "2020-01-03T00:00:00Z",
+ "executor": None,
+ "executor_config": "{}",
+ "hostname": "",
+ "id": mock.ANY,
+ "map_index": -1,
+ "max_tries": 0,
+ "note": "placeholder-note",
+ "operator": "PythonOperator",
+ "pid": 100,
+ "pool": "default_pool",
+ "pool_slots": 1,
+ "priority_weight": 9,
+ "queue": "default_queue",
+ "queued_when": None,
+ "start_date": "2020-01-02T00:00:00Z",
+ "state": "running",
+ "task_display_name": "print_the_context",
+ "try_number": 0,
+ "unixname": getuser(),
+ "rendered_fields": {},
+ "rendered_map_index": None,
+ "trigger": None,
+ "triggerer_job": None,
+ },
+ 1,
+ ),
+ (
+ None,
+ 422,
+ {
+ "detail": [
+ {
+ "type": "value_error",
+ "loc": ["body", "new_state"],
+ "msg": "Value error, 'new_state' should not be empty",
+ "input": None,
+ "ctx": {"error": {}},
+ }
+ ]
+ },
+ 0,
+ ),
+ ],
+ )
+ @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
+ def test_update_mask_should_call_mocked_api(
+ self,
+ mock_set_ti_state,
+ test_client,
+ session,
+ new_state,
+ expected_status_code,
+ expected_json,
+ set_ti_state_call_count,
+ ):
+ self.create_task_instances(session)
+
+ mock_set_ti_state.return_value = session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == "example_python_operator",
+ TaskInstance.task_id == "print_the_context",
+ TaskInstance.run_id == "TEST_DAG_RUN_ID",
+ TaskInstance.map_index == -1,
+ )
+ ).one_or_none()
+
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ params={"update_mask": "new_state"},
+ json={
+ "dry_run": False,
+ "new_state": new_state,
+ },
+ )
+ assert response.status_code == expected_status_code
+ assert response.json() == expected_json
+ assert mock_set_ti_state.call_count == set_ti_state_call_count
+
+ @pytest.mark.parametrize(
+ "new_note_value",
+ [
+ "My super cool TaskInstance note.",
+ None,
+ ],
+ )
+ def test_update_mask_set_note_should_respond_200(self, test_client, session, new_note_value):
+ self.create_task_instances(session)
+ response = test_client.patch(
+ "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context",
+ params={"update_mask": "note"},
+ json={"note": new_note_value},
+ )
+ assert response.status_code == 200, response.text
+ assert response.json() == {
+ "dag_id": "example_python_operator",
+ "duration": 10000.0,
+ "end_date": "2020-01-03T00:00:00Z",
+ "logical_date": "2020-01-01T00:00:00Z",
+ "id": mock.ANY,
+ "executor": None,
+ "executor_config": "{}",
+ "hostname": "",
+ "map_index": -1,
+ "max_tries": 0,
+ "note": new_note_value,
+ "operator": "PythonOperator",
+ "pid": 100,
+ "pool": "default_pool",
+ "pool_slots": 1,
+ "priority_weight": 9,
+ "queue": "default_queue",
+ "queued_when": None,
+ "start_date": "2020-01-02T00:00:00Z",
+ "state": "running",
+ "task_id": "print_the_context",
+ "task_display_name": "print_the_context",
+ "try_number": 0,
+ "unixname": getuser(),
+ "dag_run_id": "TEST_DAG_RUN_ID",
+ "rendered_fields": {},
+ "rendered_map_index": None,
+ "trigger": None,
+ "triggerer_job": None,
+ }
+
+ def test_set_note_should_respond_200(self, test_client, session):
+ self.create_task_instances(session)
+ new_note_value = "My super cool TaskInstance note."
+ response = test_client.patch(
+ "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context",
+ json={"note": new_note_value},
+ )
+ assert response.status_code == 200, response.text
+ assert response.json() == {
+ "dag_id": "example_python_operator",
+ "duration": 10000.0,
+ "end_date": "2020-01-03T00:00:00Z",
+ "logical_date": "2020-01-01T00:00:00Z",
+ "id": mock.ANY,
+ "executor": None,
+ "executor_config": "{}",
+ "hostname": "",
+ "map_index": -1,
+ "max_tries": 0,
+ "note": new_note_value,
+ "operator": "PythonOperator",
+ "pid": 100,
+ "pool": "default_pool",
+ "pool_slots": 1,
+ "priority_weight": 9,
+ "queue": "default_queue",
+ "queued_when": None,
+ "start_date": "2020-01-02T00:00:00Z",
+ "state": "running",
+ "task_id": "print_the_context",
+ "task_display_name": "print_the_context",
+ "try_number": 0,
+ "unixname": getuser(),
+ "dag_run_id": "TEST_DAG_RUN_ID",
+ "rendered_fields": {},
+ "rendered_map_index": None,
+ "trigger": None,
+ "triggerer_job": None,
+ }
+ # @TODO: Uncomment the 2 lines below when permissions and auth is in place.
+ # ti = session.scalars(select(TaskInstance).where(TaskInstance.task_id == "print_the_context")).one()
+ # assert ti.task_instance_note.user_id is not None
Review Comment:
Done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on PR #44223:
URL: https://github.com/apache/airflow/pull/44223#issuecomment-2493928125
All conversations resolved, PR rebased with `main` and all conflicts resolved, please have a look when possible @pierrejeambrun @bbovenzi. Thanks.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on code in PR #44223:
URL: https://github.com/apache/airflow/pull/44223#discussion_r1856970678
##########
tests/api_fastapi/core_api/routes/public/test_task_instances.py:
##########
@@ -2318,3 +2319,542 @@ def test_raises_404_for_non_existent_dag(self, test_client):
)
assert response.status_code == 404
assert "DAG non-existent-dag not found" in response.text
+
+
+class TestPatchTaskInstance(TestTaskInstanceEndpoint):
+ ENDPOINT_URL = (
+ "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context"
+ )
+
+ @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
+ def test_should_call_mocked_api(self, mock_set_task_instance_state, test_client, session):
+ self.create_task_instances(session)
+
+ NEW_STATE = "failed"
+ mock_set_task_instance_state.return_value = session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == "example_python_operator",
+ TaskInstance.task_id == "print_the_context",
+ TaskInstance.run_id == "TEST_DAG_RUN_ID",
+ TaskInstance.map_index == -1,
+ )
+ ).one_or_none()
+
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dry_run": False,
+ "new_state": NEW_STATE,
+ },
+ )
+ assert response.status_code == 200
+ assert response.json() == {
+ "dag_id": "example_python_operator",
+ "dag_run_id": "TEST_DAG_RUN_ID",
+ "logical_date": "2020-01-01T00:00:00Z",
+ "task_id": "print_the_context",
+ "duration": 10000.0,
+ "end_date": "2020-01-03T00:00:00Z",
+ "executor": None,
+ "executor_config": "{}",
+ "hostname": "",
+ "id": mock.ANY,
+ "map_index": -1,
+ "max_tries": 0,
+ "note": "placeholder-note",
+ "operator": "PythonOperator",
+ "pid": 100,
+ "pool": "default_pool",
+ "pool_slots": 1,
+ "priority_weight": 9,
+ "queue": "default_queue",
+ "queued_when": None,
+ "start_date": "2020-01-02T00:00:00Z",
+ "state": "running",
+ "task_display_name": "print_the_context",
+ "try_number": 0,
+ "unixname": getuser(),
+ "rendered_fields": {},
+ "rendered_map_index": None,
+ "trigger": None,
+ "triggerer_job": None,
+ }
+
+ mock_set_task_instance_state.assert_called_once()
Review Comment:
Please assert the parameter the function was called with.
##########
tests/api_fastapi/core_api/routes/public/test_task_instances.py:
##########
@@ -2318,3 +2319,542 @@ def test_raises_404_for_non_existent_dag(self, test_client):
)
assert response.status_code == 404
assert "DAG non-existent-dag not found" in response.text
+
+
+class TestPatchTaskInstance(TestTaskInstanceEndpoint):
+ ENDPOINT_URL = (
+ "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context"
+ )
+
+ @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
+ def test_should_call_mocked_api(self, mock_set_task_instance_state, test_client, session):
+ self.create_task_instances(session)
+
+ NEW_STATE = "failed"
+ mock_set_task_instance_state.return_value = session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == "example_python_operator",
+ TaskInstance.task_id == "print_the_context",
+ TaskInstance.run_id == "TEST_DAG_RUN_ID",
+ TaskInstance.map_index == -1,
+ )
+ ).one_or_none()
+
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dry_run": False,
+ "new_state": NEW_STATE,
+ },
+ )
+ assert response.status_code == 200
+ assert response.json() == {
+ "dag_id": "example_python_operator",
+ "dag_run_id": "TEST_DAG_RUN_ID",
+ "logical_date": "2020-01-01T00:00:00Z",
+ "task_id": "print_the_context",
+ "duration": 10000.0,
+ "end_date": "2020-01-03T00:00:00Z",
+ "executor": None,
+ "executor_config": "{}",
+ "hostname": "",
+ "id": mock.ANY,
+ "map_index": -1,
+ "max_tries": 0,
+ "note": "placeholder-note",
+ "operator": "PythonOperator",
+ "pid": 100,
+ "pool": "default_pool",
+ "pool_slots": 1,
+ "priority_weight": 9,
+ "queue": "default_queue",
+ "queued_when": None,
+ "start_date": "2020-01-02T00:00:00Z",
+ "state": "running",
+ "task_display_name": "print_the_context",
+ "try_number": 0,
+ "unixname": getuser(),
+ "rendered_fields": {},
+ "rendered_map_index": None,
+ "trigger": None,
+ "triggerer_job": None,
+ }
+
+ mock_set_task_instance_state.assert_called_once()
+
+ @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
+ def test_should_not_call_mocked_api_for_dry_run(self, mock_set_task_instance_state, test_client, session):
+ self.create_task_instances(session)
+
+ NEW_STATE = "failed"
+ mock_set_task_instance_state.return_value = session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == "example_python_operator",
+ TaskInstance.task_id == "print_the_context",
+ TaskInstance.run_id == "TEST_DAG_RUN_ID",
+ TaskInstance.map_index == -1,
+ )
+ ).one_or_none()
+
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dry_run": True,
+ "new_state": NEW_STATE,
+ },
+ )
+ assert response.status_code == 200
+ assert response.json() == {
+ "dag_id": "example_python_operator",
+ "dag_run_id": "TEST_DAG_RUN_ID",
+ "logical_date": "2020-01-01T00:00:00Z",
+ "task_id": "print_the_context",
+ "duration": 10000.0,
+ "end_date": "2020-01-03T00:00:00Z",
+ "executor": None,
+ "executor_config": "{}",
+ "hostname": "",
+ "id": mock.ANY,
+ "map_index": -1,
+ "max_tries": 0,
+ "note": "placeholder-note",
+ "operator": "PythonOperator",
+ "pid": 100,
+ "pool": "default_pool",
+ "pool_slots": 1,
+ "priority_weight": 9,
+ "queue": "default_queue",
+ "queued_when": None,
+ "start_date": "2020-01-02T00:00:00Z",
+ "state": "running",
+ "task_display_name": "print_the_context",
+ "try_number": 0,
+ "unixname": getuser(),
+ "rendered_fields": {},
+ "rendered_map_index": None,
+ "trigger": None,
+ "triggerer_job": None,
+ }
+
+ mock_set_task_instance_state.assert_not_called()
+
+ def test_should_update_task_instance_state(self, test_client, session):
+ self.create_task_instances(session)
+
+ NEW_STATE = "failed"
+
+ test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dry_run": False,
+ "new_state": NEW_STATE,
+ },
+ )
+
+ response2 = test_client.get(self.ENDPOINT_URL)
+ assert response2.status_code == 200
+ assert response2.json()["state"] == NEW_STATE
+
+ def test_should_update_task_instance_state_default_dry_run_to_true(self, test_client, session):
+ self.create_task_instances(session)
+
+ NEW_STATE = "running"
+
+ test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "new_state": NEW_STATE,
+ },
+ )
+
+ response2 = test_client.get(self.ENDPOINT_URL)
+ assert response2.status_code == 200
+ assert response2.json()["state"] == NEW_STATE
+
+ def test_should_update_mapped_task_instance_state(self, test_client, session):
+ NEW_STATE = "failed"
+ map_index = 1
+ tis = self.create_task_instances(session)
+ ti = TaskInstance(task=tis[0].task, run_id=tis[0].run_id, map_index=map_index)
+ ti.rendered_task_instance_fields = RTIF(ti, render_templates=False)
+ session.add(ti)
+ session.commit()
+
+ response = test_client.patch(
+ f"{self.ENDPOINT_URL}/{map_index}",
+ json={
+ "dry_run": False,
+ "new_state": NEW_STATE,
+ },
+ )
+ assert response.status_code == 200
+
+ response2 = test_client.get(f"{self.ENDPOINT_URL}/{map_index}")
+ assert response2.status_code == 200
+ assert response2.json()["state"] == NEW_STATE
+
+ @pytest.mark.parametrize(
+ "error, code, payload",
+ [
+ [
+ (
+ "Task Instance not found for dag_id=example_python_operator"
+ ", run_id=TEST_DAG_RUN_ID, task_id=print_the_context"
+ ),
+ 404,
+ {
+ "dry_run": True,
+ "new_state": "failed",
+ },
+ ]
+ ],
+ )
+ def test_should_handle_errors(self, error, code, payload, test_client, session):
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json=payload,
+ )
+ assert response.status_code == code
+ assert response.json()["detail"] == error
+
+ def test_should_200_for_unknown_fields(self, test_client, session):
+ self.create_task_instances(session)
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dryrun": True,
+ "new_state": "failed",
+ },
+ )
+ assert response.status_code == 200
+
+ def test_should_raise_404_for_non_existent_dag(self, test_client):
+ response = test_client.patch(
+ "/public/dags/non-existent-dag/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context",
+ json={
+ "dry_run": False,
+ "new_state": "failed",
+ },
+ )
+ assert response.status_code == 404
+ assert response.json() == {"detail": "DAG non-existent-dag not found"}
+
+ def test_should_raise_404_for_non_existent_task_in_dag(self, test_client):
+ response = test_client.patch(
+ "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/non_existent_task",
+ json={
+ "dry_run": False,
+ "new_state": "failed",
+ },
+ )
+ assert response.status_code == 404
+ assert response.json() == {
+ "detail": "Task 'non_existent_task' not found in DAG 'example_python_operator'"
+ }
+
+ def test_should_raise_404_not_found_dag(self, test_client):
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dry_run": True,
+ "new_state": "failed",
+ },
+ )
+ assert response.status_code == 404
+
+ def test_should_raise_404_not_found_task(self, test_client):
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json={
+ "dry_run": True,
+ "new_state": "failed",
+ },
+ )
+ assert response.status_code == 404
+
+ @pytest.mark.parametrize(
+ "payload, expected",
+ [
+ (
+ {
+ "dry_run": True,
+ "new_state": "failede",
+ },
+ f"'failede' is not one of ['{State.SUCCESS}', '{State.FAILED}', '{State.SKIPPED}']",
+ ),
+ (
+ {
+ "dry_run": True,
+ "new_state": "queued",
+ },
+ f"'queued' is not one of ['{State.SUCCESS}', '{State.FAILED}', '{State.SKIPPED}']",
+ ),
+ ],
+ )
+ def test_should_raise_422_for_invalid_task_instance_state(self, payload, expected, test_client, session):
+ self.create_task_instances(session)
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ json=payload,
+ )
+ assert response.status_code == 422
+ assert response.json() == {
+ "detail": [
+ {
+ "type": "value_error",
+ "loc": ["body", "new_state"],
+ "msg": f"Value error, {expected}",
+ "input": payload["new_state"],
+ "ctx": {"error": {}},
+ }
+ ]
+ }
+
+ @pytest.mark.parametrize(
+ "new_state,expected_status_code,expected_json,set_ti_state_call_count",
+ [
+ (
+ "failed",
+ 200,
+ {
+ "dag_id": "example_python_operator",
+ "dag_run_id": "TEST_DAG_RUN_ID",
+ "logical_date": "2020-01-01T00:00:00Z",
+ "task_id": "print_the_context",
+ "duration": 10000.0,
+ "end_date": "2020-01-03T00:00:00Z",
+ "executor": None,
+ "executor_config": "{}",
+ "hostname": "",
+ "id": mock.ANY,
+ "map_index": -1,
+ "max_tries": 0,
+ "note": "placeholder-note",
+ "operator": "PythonOperator",
+ "pid": 100,
+ "pool": "default_pool",
+ "pool_slots": 1,
+ "priority_weight": 9,
+ "queue": "default_queue",
+ "queued_when": None,
+ "start_date": "2020-01-02T00:00:00Z",
+ "state": "running",
+ "task_display_name": "print_the_context",
+ "try_number": 0,
+ "unixname": getuser(),
+ "rendered_fields": {},
+ "rendered_map_index": None,
+ "trigger": None,
+ "triggerer_job": None,
+ },
+ 1,
+ ),
+ (
+ None,
+ 422,
+ {
+ "detail": [
+ {
+ "type": "value_error",
+ "loc": ["body", "new_state"],
+ "msg": "Value error, 'new_state' should not be empty",
+ "input": None,
+ "ctx": {"error": {}},
+ }
+ ]
+ },
+ 0,
+ ),
+ ],
+ )
+ @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
+ def test_update_mask_should_call_mocked_api(
+ self,
+ mock_set_ti_state,
+ test_client,
+ session,
+ new_state,
+ expected_status_code,
+ expected_json,
+ set_ti_state_call_count,
+ ):
+ self.create_task_instances(session)
+
+ mock_set_ti_state.return_value = session.scalars(
+ select(TaskInstance).where(
+ TaskInstance.dag_id == "example_python_operator",
+ TaskInstance.task_id == "print_the_context",
+ TaskInstance.run_id == "TEST_DAG_RUN_ID",
+ TaskInstance.map_index == -1,
+ )
+ ).one_or_none()
+
+ response = test_client.patch(
+ self.ENDPOINT_URL,
+ params={"update_mask": "new_state"},
+ json={
+ "dry_run": False,
+ "new_state": new_state,
+ },
+ )
+ assert response.status_code == expected_status_code
+ assert response.json() == expected_json
+ assert mock_set_ti_state.call_count == set_ti_state_call_count
+
+ @pytest.mark.parametrize(
+ "new_note_value",
+ [
+ "My super cool TaskInstance note.",
+ None,
+ ],
+ )
+ def test_update_mask_set_note_should_respond_200(self, test_client, session, new_note_value):
+ self.create_task_instances(session)
+ response = test_client.patch(
+ "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context",
+ params={"update_mask": "note"},
+ json={"note": new_note_value},
+ )
+ assert response.status_code == 200, response.text
+ assert response.json() == {
+ "dag_id": "example_python_operator",
+ "duration": 10000.0,
+ "end_date": "2020-01-03T00:00:00Z",
+ "logical_date": "2020-01-01T00:00:00Z",
+ "id": mock.ANY,
+ "executor": None,
+ "executor_config": "{}",
+ "hostname": "",
+ "map_index": -1,
+ "max_tries": 0,
+ "note": new_note_value,
+ "operator": "PythonOperator",
+ "pid": 100,
+ "pool": "default_pool",
+ "pool_slots": 1,
+ "priority_weight": 9,
+ "queue": "default_queue",
+ "queued_when": None,
+ "start_date": "2020-01-02T00:00:00Z",
+ "state": "running",
+ "task_id": "print_the_context",
+ "task_display_name": "print_the_context",
+ "try_number": 0,
+ "unixname": getuser(),
+ "dag_run_id": "TEST_DAG_RUN_ID",
+ "rendered_fields": {},
+ "rendered_map_index": None,
+ "trigger": None,
+ "triggerer_job": None,
+ }
+
+ def test_set_note_should_respond_200(self, test_client, session):
+ self.create_task_instances(session)
+ new_note_value = "My super cool TaskInstance note."
+ response = test_client.patch(
+ "/public/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context",
+ json={"note": new_note_value},
+ )
+ assert response.status_code == 200, response.text
+ assert response.json() == {
+ "dag_id": "example_python_operator",
+ "duration": 10000.0,
+ "end_date": "2020-01-03T00:00:00Z",
+ "logical_date": "2020-01-01T00:00:00Z",
+ "id": mock.ANY,
+ "executor": None,
+ "executor_config": "{}",
+ "hostname": "",
+ "map_index": -1,
+ "max_tries": 0,
+ "note": new_note_value,
+ "operator": "PythonOperator",
+ "pid": 100,
+ "pool": "default_pool",
+ "pool_slots": 1,
+ "priority_weight": 9,
+ "queue": "default_queue",
+ "queued_when": None,
+ "start_date": "2020-01-02T00:00:00Z",
+ "state": "running",
+ "task_id": "print_the_context",
+ "task_display_name": "print_the_context",
+ "try_number": 0,
+ "unixname": getuser(),
+ "dag_run_id": "TEST_DAG_RUN_ID",
+ "rendered_fields": {},
+ "rendered_map_index": None,
+ "trigger": None,
+ "triggerer_job": None,
+ }
+ # @TODO: Uncomment the 2 lines below when permissions and auth is in place.
+ # ti = session.scalars(select(TaskInstance).where(TaskInstance.task_id == "print_the_context")).one()
+ # assert ti.task_instance_note.user_id is not None
Review Comment:
You can remove that. All of the places where this should be updated are not tracked at the moment.
We'll have to just go over the code i guess. No need to leave partial comments here and there.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on PR #44223:
URL: https://github.com/apache/airflow/pull/44223#issuecomment-2489465091
All checks passed, PR is ready for review ✅
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on PR #44223:
URL: https://github.com/apache/airflow/pull/44223#issuecomment-2492438234
> I think we need to include the `update_mask` as we do for other patch endpoints. (The state part would actually be the `set_task_instance_state` of your other PR)
@pierrejeambrun I've added `update_mask` here in this PR to Patch Task Instance API.
Yes got it, sure, I've closed the other Set Task Instance State PR https://github.com/apache/airflow/pull/44246.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "pierrejeambrun (via GitHub)" <gi...@apache.org>.
pierrejeambrun commented on code in PR #44223:
URL: https://github.com/apache/airflow/pull/44223#discussion_r1852454251
##########
airflow/api_fastapi/core_api/datamodels/task_instances.py:
##########
@@ -150,3 +152,41 @@ class TaskInstanceHistoryCollectionResponse(BaseModel):
task_instances: list[TaskInstanceHistoryResponse]
total_entries: int
+
+
+class PatchTaskInstanceBody(BaseModel):
+ """Request body for Clear Task Instances endpoint."""
+
+ dry_run: bool = True
+ new_state: str | None = None
+ note: str | None = None
+
+ @model_validator(mode="before")
+ @classmethod
+ def validate_model(cls, data: Any) -> Any:
+ if data.get("note") is None and data.get("new_state") is None:
+ raise ValueError("new_state is required.")
+ return data
+
+ @field_validator("note", mode="before")
+ @classmethod
+ def validate_note(cls, note: str | None) -> str | None:
+ """Validate note."""
+ if note is None:
+ return None
+ if len(note) > 1000:
+ raise ValueError("Note length should not exceed 1000 characters.")
Review Comment:
This exists natively in pydantic Field. We can use that instead of coding a check manually.
##########
airflow/api_fastapi/core_api/routes/public/task_instances.py:
##########
@@ -482,3 +484,78 @@ def get_mapped_task_instance_try_details(
map_index=map_index,
session=session,
)
+
+
+@task_instances_router.patch(
+ "/{task_id}",
+ responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST]),
+)
+@task_instances_router.patch(
+ "/{task_id}/{map_index}",
+ responses=create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_400_BAD_REQUEST]),
+)
+def patch_task_instance(
+ dag_id: str,
+ dag_run_id: str,
+ task_id: str,
+ request: Request,
+ body: PatchTaskInstanceBody,
+ session: Annotated[Session, Depends(get_session)],
+ map_index: int = -1,
Review Comment:
We are missing the update mask.
You can take a look at others (not all of them are reworked yet), but `patch_dag_run` might be close to what you want.
Also, if `mask` is provided, `None` values should not be ignored but set to `None`.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Migrate public endpoint Patch Task Instance to FastAPI [airflow]
Posted by "omkar-foss (via GitHub)" <gi...@apache.org>.
omkar-foss commented on code in PR #44223:
URL: https://github.com/apache/airflow/pull/44223#discussion_r1854023921
##########
airflow/api_fastapi/core_api/datamodels/task_instances.py:
##########
@@ -190,6 +192,33 @@ def validate_model(cls, data: Any) -> Any:
return data
+class PatchTaskInstanceBody(BaseModel):
+ """Request body for Clear Task Instances endpoint."""
+
+ dry_run: bool = True
+ new_state: str | None = None
+ note: Annotated[str, StringConstraints(max_length=1000)] | None = None
+ include_upstream: bool = False
Review Comment:
Defaulting these to `False` as per your suggestion @bbovenzi.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org