pierrejeambrun commented on code in PR #26457:
URL: https://github.com/apache/airflow/pull/26457#discussion_r982757590
##########
airflow/api_connexion/endpoints/task_instance_endpoint.py:
##########
@@ -545,3 +546,54 @@ def post_set_task_instances_state(*, dag_id: str, session:
Session = NEW_SESSION
session=session,
)
return
task_instance_reference_collection_schema.dump(TaskInstanceReferenceCollection(task_instances=tis))
+
+
[email protected]_access(
+ [
+ (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG),
+ (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_DAG_RUN),
+ ],
+)
+@provide_session
+def set_task_instance_note(
+ *, dag_id: str, dag_run_id: str, task_id: str, session: Session =
NEW_SESSION
+) -> APIResponse:
+ """Set the note for a dag run."""
+ try:
+ post_body =
set_task_instance_note_form_schema.load(get_json_request_dict())
+ new_value_for_notes = post_body["notes"]
+ # Note: We can't add map_index to the url as subpaths can't start with
dashes.
+ map_index = post_body["map_index"]
+ except ValidationError as err:
+ raise BadRequest(detail=str(err))
+
+ query = (
+ session.query(TI)
+ .filter(TI.dag_id == dag_id)
+ .filter(TI.run_id == dag_run_id)
+ .filter(TI.task_id == task_id)
+ )
+ if map_index == -1:
+ query = query.filter(or_(TI.map_index == -1, TI.map_index is None))
+ else:
+ query = query.filter(TI.map_index == map_index)
+
+ try:
+ ti: TI | None = query.one_or_none()
+ except MultipleResultsFound:
+ raise NotFound(
+ "Task instance not found", detail="Task instance is mapped, add
the map_index value to the URL"
+ )
+ if ti is None:
+ error_message = f'Task Instance not found for dag_id={dag_id},
run_id={dag_run_id}, task_id={task_id}'
+ raise NotFound(error_message)
+
+ ti.notes = new_value_for_notes
+ session.commit()
+
+ if map_index == -1:
+ return get_task_instance(dag_id=dag_id, dag_run_id=dag_run_id,
task_id=task_id)
Review Comment:
We already have the TaskInstance available in the scope, is it possible to
directly serialize this instance and return it, instead of calling different
functions that will make an additional db query. Also Avoid the `if else`
logic:
```
return task_instance_schema.dump((ti, None)) # The task_instance_schema
seems a little special :)
```
If we need the `SlaMiss` and `rendered_task_instance_fields` you can add
them to your default query and update the dump call.
--
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]
